Activation + GroupedLinear Fusion for MOE and other MOE optimizations

qa/L0_pytorch_unittest/test.sh

@@ -70,6 +70,7 @@ NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_
 PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py"
 PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_ops_grouped_linear_distributed_weight.xml $TE_PATH/tests/pytorch/test_ops_grouped_linear_distributed_weight.py || test_fail "test_ops_grouped_linear_distributed_weight.py"
 NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py"
+NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_mlp_without_cutedsl_fusion.xml $TE_PATH/tests/pytorch/test_grouped_mlp.py || test_fail "test_grouped_mlp.py without CuTe DSL fusion"
 
 if [ "$RET" -ne 0 ]; then
     echo "Error in the following test cases:$FAILED_CASES"

tests/pytorch/test_grouped_mlp.py

@@ -24,6 +24,12 @@
     OUTPUT_BUFFER_KEY,
     GRAD_INPUT_BUFFER_KEY,
 )
+from transformer_engine.pytorch.ops.fused.backward_activation_grouped_linear import (
+    BackwardScaledActivationGroupedLinear,
+)
+from transformer_engine.pytorch.ops.fused.forward_activation_grouped_linear import (
+    ForwardScaledActivationGroupedLinear,
+)
 from transformer_engine.pytorch import (
     QuantizedTensor,
     Float8CurrentScalingQuantizer,
@@ -73,6 +79,10 @@
 if nvfp4_available:
     _grouped_mlp_quantization_list.append("nvfp4_rht")
 
+# Quantization recipes supported by ScaledActivation + GroupedLinear fusion
+if fp8_available:
+    _grouped_mlp_quantization_list.append("fp8_current_scaling")
+
 
 @pytest.fixture(autouse=True, scope="function")
 def _reset_rng_states_per_test():
@@ -755,6 +765,13 @@ def test_grouped_mlp(
         maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype)
         if dtype == torch.bfloat16 and not is_bf16_available():
             pytest.skip("BF16 requires SM 8.0+")
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and (
+            single_grouped_weight or single_grouped_bias
+        ):
+            pytest.skip(
+                "single_grouped_weight/single_grouped_bias requires"
+                " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1"
+            )
         if single_grouped_weight and quantization != "mxfp8":
             pytest.skip("single_grouped_weight is only supported for MXFP8 quantization")
         if single_grouped_bias and not bias:
@@ -1026,25 +1043,46 @@ def _make_module():
             or (
                 quantization == "nvfp4_rht"
                 and dtype == torch.bfloat16
-                and activation == "scaled_srelu"
-                and glu_interleave_size is None
+                and (
+                    (not activation_is_glu and glu_interleave_size is None)
+                    or (activation_is_glu and glu_interleave_size == 32)
+                )
             )
         )
+        forward_ops = module._module_groups[0]._forward_ops
+        backward_ops = module._module_groups[0]._backward_ops
+        full_grouped_mlp_fusion = False
         if expected_grouped_mlp_fusion:
             if activation_is_glu:
                 fused_cls = te.ops.fused.GroupedMLP_CuTeGEMMGLU
             else:
                 fused_cls = te.ops.fused.GroupedMLP_CuTeGEMMUnary
             if fused_cls.is_supported():
-                forward_ops = module._module_groups[0]._forward_ops
-                backward_ops = module._module_groups[0]._backward_ops
                 assert len(forward_ops) == 1
                 assert len(backward_ops) == 1
                 assert isinstance(
                     forward_ops[0][0],
                     fused_cls,
                 )
                 assert backward_ops[0][0] is forward_ops[0][0]
+                full_grouped_mlp_fusion = True
+
+        # When the full FC1 + activation + FC2 fusion is unavailable, verify
+        # that ScaledActivation + GroupedLinear fusions cover both boundaries
+        # whenever grouped quantized compute is supported.
+        act_grouped_linear_fusion_expected = (
+            not full_grouped_mlp_fusion
+            and te.ops.fused.act_grouped_linear_fusion_supported(fc2, module[1], recipe)
+            and te.ops.fused.act_grouped_linear_fusion_supported(fc1, module[1], recipe)
+        )
+        assert (
+            any(isinstance(op, ForwardScaledActivationGroupedLinear) for op, _ in forward_ops)
+            == act_grouped_linear_fusion_expected
+        )
+        assert (
+            any(isinstance(op, BackwardScaledActivationGroupedLinear) for op, _ in backward_ops)
+            == act_grouped_linear_fusion_expected
+        )
 
         # Loose tols for sanity checking
         tols = {"rtol": 0.125, "atol": 0.25}
@@ -1266,6 +1304,8 @@ def test_grouped_mlp_single_weight_numerics(
     ) -> None:
         """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP."""
 
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0":
+            pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1")
         if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
             pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
 
@@ -1584,6 +1624,8 @@ def test_grouped_mlp_overwrite_main_grad(
         that read ``.grad`` don't see stale bytes from the cached dummy).
         """
 
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight:
+            pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1")
         if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
             pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
 
@@ -1715,6 +1757,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8(
     ) -> None:
         """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8)."""
 
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight:
+            pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1")
         if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
             pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
         if dtype not in (torch.bfloat16, torch.float16):

transformer_engine/pytorch/csrc/extensions.h

@@ -290,6 +290,68 @@ py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float l
 
 py::object clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer,
                            float limit, float alpha, float glu_linear_offset);
+
+/* Scaled activation */
+py::object scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
+                         py::handle quantizer, int64_t glu_interleave_size);
+
+py::object scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
+                                 py::handle quantizer, float limit, float alpha,
+                                 float glu_linear_offset, int64_t glu_interleave_size);
+
+py::object scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales,
+                        py::handle quantizer);
+
+py::tuple scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input,
+                         const at::Tensor &act_scales, py::handle quantizer,
+                         int64_t glu_interleave_size, bool compute_scale_grad);
+
+py::tuple scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input,
+                                 const at::Tensor &act_scales, py::handle quantizer, float limit,
+                                 float alpha, float glu_linear_offset, int64_t glu_interleave_size,
+                                 bool compute_scale_grad);
+
+py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input,
+                        const at::Tensor &act_scales, py::handle quantizer,
+                        bool compute_scale_grad);
+
+/* Scaled activation + grouped quantize */
+py::object grouped_scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
+                                 py::handle quantizer, const size_t num_tensors,
+                                 std::optional<at::Tensor> first_dims,
+                                 std::optional<at::Tensor> tensor_offsets,
+                                 int64_t glu_interleave_size);
+
+py::object grouped_scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
+                                         py::handle quantizer, const size_t num_tensors,
+                                         std::optional<at::Tensor> first_dims,
+                                         std::optional<at::Tensor> tensor_offsets, float limit,
+                                         float alpha, float glu_linear_offset,
+                                         int64_t glu_interleave_size);
+
+py::object grouped_scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales,
+                                py::handle quantizer, const size_t num_tensors,
+                                std::optional<at::Tensor> first_dims,
+                                std::optional<at::Tensor> tensor_offsets);
+
+py::tuple grouped_scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input,
+                                 const at::Tensor &act_scales, py::handle quantizer,
+                                 const size_t num_tensors, std::optional<at::Tensor> first_dims,
+                                 std::optional<at::Tensor> tensor_offsets,
+                                 int64_t glu_interleave_size, bool compute_scale_grad);
+
+py::tuple grouped_scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input,
+                                         const at::Tensor &act_scales, py::handle quantizer,
+                                         const size_t num_tensors,
+                                         std::optional<at::Tensor> first_dims,
+                                         std::optional<at::Tensor> tensor_offsets, float limit,
+                                         float alpha, float glu_linear_offset,
+                                         int64_t glu_interleave_size, bool compute_scale_grad);
+
+py::tuple grouped_scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input,
+                                const at::Tensor &act_scales, py::handle quantizer,
+                                const size_t num_tensors, std::optional<at::Tensor> first_dims,
+                                std::optional<at::Tensor> tensor_offsets, bool compute_scale_grad);
 /***************************************************************************************************
  * LayerNorm
  **************************************************************************************************/