Activation + GroupedLinear Fusion for MOE

tests/pytorch/test_grouped_mlp.py

@@ -23,6 +23,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,
@@ -72,6 +78,15 @@
 if nvfp4_available:
     _grouped_mlp_quantization_list.append("nvfp4_rht")
 
+# Quantization recipes for ScaledActivation + GroupedLinear fusion
+_grouped_mlp_act_grouped_linear_quantization_list: list[str] = []
+if fp8_available:
+    _grouped_mlp_act_grouped_linear_quantization_list.append("fp8_current_scaling")
+if mxfp8_available:
+    _grouped_mlp_act_grouped_linear_quantization_list.append("mxfp8")
+if nvfp4_available:
+    _grouped_mlp_act_grouped_linear_quantization_list.append("nvfp4_rht")
+
 
 @pytest.fixture(autouse=True, scope="function")
 def _reset_rng_states_per_test():
@@ -724,6 +739,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:
@@ -999,21 +1021,40 @@ def _make_module():
                 and glu_interleave_size is None
             )
         )
+        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}
@@ -1076,6 +1117,25 @@ def _make_module():
             assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols)
             assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols)
 
+    @pytest.mark.parametrize("quantization", _grouped_mlp_act_grouped_linear_quantization_list)
+    def test_grouped_mlp_act_grouped_linear_fusion(
+        self,
+        *,
+        quantization: str,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        """Exercise ScaledActivation + GroupedLinear fusion without the full CuTe DSL MLP fusion."""
+        monkeypatch.setenv("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "0")
+        self.test_grouped_mlp(
+            group_size=2,
+            bias=False,
+            hidden_size=128,
+            quantization=quantization,
+            single_grouped_weight=False,
+            split_alignment=128,
+            activation="scaled_swiglu",
+        )
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list)
     @pytest.mark.parametrize(
@@ -1148,6 +1208,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")
 
@@ -1466,6 +1528,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")
 
@@ -1597,6 +1661,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,44 @@ 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 + 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
  **************************************************************************************************/

transformer_engine/pytorch/csrc/extensions/activation.cpp

@@ -342,5 +342,151 @@ py::object clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, py::
                                                               glu_linear_offset);
 }
 
+/* Scaled activation + grouped quantize helpers (mirrors activation_helper / dactivation_helper). */
+
+template <auto act_func, typename... Args>
+py::object grouped_scaled_activation_helper(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,
+                                            int shape_divisor, Args&&... args) {
+  init_extension();
+  NVTE_CHECK(input.dim() == 2, "grouped scaled activation input must be 2D");
+  NVTE_CHECK(act_scales.numel() == input.size(0),
+             "grouped scaled activation expects one scale per input row");
+  NVTE_CHECK(shape_divisor > 0 && input.size(1) % shape_divisor == 0,
+             "grouped scaled activation input width is not compatible with activation");
+
+  auto input_tensor = input.contiguous();
+  auto scales_tensor = act_scales.contiguous();
+  const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor);
+  const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_tensor);
+
+  // Keep the dense activation in the input dtype. It is only a transient buffer when a
+  // quantizer is provided; the quantized return path never exposes it to the caller.
+  auto output = at::empty({input.size(0), input.size(1) / shape_divisor}, input_tensor.options());
+  const TensorWrapper& output_nvte = makeTransformerEngineTensor(output);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  NVTE_SCOPED_GIL_RELEASE({
+    act_func(input_nvte.data(), scales_nvte.data(), output_nvte.data(), std::forward<Args>(args)...,
+             stream);
+  });
+
+  if (quantizer.is_none()) {
+    return py::cast(output);
+  }
+  return group_quantize(output, quantizer, num_tensors, first_dims, std::nullopt, tensor_offsets,
+                        std::nullopt);
+}
+
+template <auto dact_func, typename... Args>
+py::tuple grouped_scaled_dactivation_helper(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, Args&&... args) {
+  init_extension();
+  NVTE_CHECK(input.dim() == 2 && grad.dim() == 2,
+             "grouped scaled dactivation input and grad must be 2D");
+  NVTE_CHECK(act_scales.numel() == input.size(0),
+             "grouped scaled dactivation expects one scale per input row");
+
+  auto grad_tensor = grad.contiguous();
+  auto input_tensor = input.contiguous();
+  auto scales_tensor = act_scales.contiguous();
+  auto grad_input = at::empty_like(input_tensor);
+  auto grad_scales = compute_scale_grad ? at::empty_like(scales_tensor) : at::Tensor();
+
+  const TensorWrapper& grad_nvte = makeTransformerEngineTensor(grad_tensor);
+  const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor);
+  const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_tensor);
+  const TensorWrapper& grad_input_nvte = makeTransformerEngineTensor(grad_input);
+  std::optional<TensorWrapper> grad_scales_nvte;
+  if (compute_scale_grad) {
+    grad_scales_nvte.emplace(makeTransformerEngineTensor(grad_scales));
+  }
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  NVTE_SCOPED_GIL_RELEASE({
+    dact_func(grad_nvte.data(), input_nvte.data(), scales_nvte.data(), grad_input_nvte.data(),
+              compute_scale_grad ? grad_scales_nvte->data() : nullptr, std::forward<Args>(args)...,
+              stream);
+  });
+
+  // Return both the (optionally) grouped-quantized grad input for the next
+  // grouped GEMM and the dense high-precision grad input so callers can reuse
+  // it (e.g. bias gradient) without a lossy dequantize.
+  py::object grad_input_out = py::cast(grad_input);
+  if (!quantizer.is_none()) {
+    grad_input_out = group_quantize(grad_input, quantizer, num_tensors, first_dims, std::nullopt,
+                                    tensor_offsets, std::nullopt);
+  }
+  return py::make_tuple(grad_input_out, py::cast(grad_input),
+                        compute_scale_grad ? py::cast(grad_scales) : py::none());
+}
+
+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) {
+  return grouped_scaled_activation_helper<nvte_scaled_swiglu>(
+      input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets, /*shape_divisor=*/2,
+      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) {
+  return grouped_scaled_activation_helper<nvte_scaled_clamped_swiglu>(
+      input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets, /*shape_divisor=*/2,
+      limit, alpha, glu_linear_offset, 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) {
+  return grouped_scaled_activation_helper<nvte_scaled_srelu>(
+      input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets,
+      /*shape_divisor=*/1);
+}
+
+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) {
+  return grouped_scaled_dactivation_helper<nvte_scaled_dswiglu>(
+      grad, input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets,
+      compute_scale_grad, glu_interleave_size);
+}
+
+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) {
+  return grouped_scaled_dactivation_helper<nvte_scaled_clamped_dswiglu>(
+      grad, input, act_scales, quantizer, num_tensors, first_dims, tensor_offsets,
+      compute_scale_grad, limit, alpha, glu_linear_offset, glu_interleave_size);
+}
+
+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) {
+  return grouped_scaled_dactivation_helper<nvte_scaled_dsrelu>(grad, input, act_scales, quantizer,
+                                                               num_tensors, first_dims,
+                                                               tensor_offsets, compute_scale_grad);
+}
+
 }  // namespace pytorch
 }  // namespace transformer_engine