[PyTorch]Fused ScaledActivation Kernel in TE Ops and reducing CPU overheads for MOE

tests/pytorch/test_grouped_mlp.py

@@ -755,6 +755,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,8 +1033,10 @@ 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)
+                )
             )
         )
         if expected_grouped_mlp_fusion:
@@ -1266,6 +1275,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 +1595,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 +1728,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,31 @@ 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);
+
 /***************************************************************************************************
  * LayerNorm
  **************************************************************************************************/

transformer_engine/pytorch/csrc/extensions/activation.cpp

@@ -342,5 +342,150 @@ py::object clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, py::
                                                               glu_linear_offset);
 }
 
+/* Scaled activation helpers (activation + per-row scale via nvte_scaled_*). */
+
+template <auto act_func, typename... Args>
+at::Tensor scaled_activation_compute(const at::Tensor& input, const at::Tensor& act_scales,
+                                     int shape_divisor, Args&&... args) {
+  init_extension();
+  NVTE_CHECK(input.dim() >= 1, "scaled activation input must have at least 1 dimension");
+  NVTE_CHECK(shape_divisor > 0 && input.size(-1) % shape_divisor == 0,
+             "scaled activation input width is not compatible with activation");
+
+  auto input_tensor = input.contiguous();
+  auto scales_tensor = act_scales.contiguous().reshape({-1});
+  const int64_t rows = input_tensor.numel() / input_tensor.size(-1);
+  NVTE_CHECK(scales_tensor.numel() == rows, "scaled activation expects one scale per input row");
+
+  std::vector<int64_t> output_sizes(input_tensor.sizes().begin(), input_tensor.sizes().end());
+  output_sizes.back() /= shape_divisor;
+  auto output = at::empty(output_sizes, input_tensor.options());
+
+  const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor);
+  const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_tensor);
+  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);
+  });
+  return output;
+}
+
+template <auto dact_func, typename... Args>
+std::tuple<at::Tensor, at::Tensor> scaled_dactivation_compute(const at::Tensor& grad,
+                                                              const at::Tensor& input,
+                                                              const at::Tensor& act_scales,
+                                                              bool compute_scale_grad,
+                                                              Args&&... args) {
+  init_extension();
+  NVTE_CHECK(input.dim() >= 1 && grad.dim() >= 1,
+             "scaled dactivation input and grad must have at least 1 dimension");
+
+  auto grad_tensor = grad.contiguous();
+  auto input_tensor = input.contiguous();
+  auto scales_tensor = act_scales.contiguous();
+  const int64_t rows = input_tensor.numel() / input_tensor.size(-1);
+  NVTE_CHECK(scales_tensor.numel() == rows, "scaled dactivation expects one scale per input row");
+
+  auto scales_flat = scales_tensor.reshape({-1});
+  auto grad_input = at::empty_like(input_tensor);
+  auto grad_scales = compute_scale_grad ? at::empty_like(scales_tensor) : at::Tensor();
+  auto grad_scales_flat = compute_scale_grad ? grad_scales.reshape({-1}) : at::Tensor();
+
+  const TensorWrapper& grad_nvte = makeTransformerEngineTensor(grad_tensor);
+  const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor);
+  const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_flat);
+  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_flat));
+  }
+
+  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 {grad_input, grad_scales};
+}
+
+py::object maybe_quantize(const at::Tensor& tensor, py::handle quantizer) {
+  if (quantizer.is_none()) {
+    return py::cast(tensor);
+  }
+  auto quantizer_cpp = convert_quantizer(quantizer);
+  const TensorWrapper& tensor_nvte = makeTransformerEngineTensor(tensor);
+  const auto shape_te = tensor_nvte.shape();
+  const std::vector<size_t> shape(shape_te.data, shape_te.data + shape_te.ndim);
+  auto fake_dtype = GetTransformerEngineDType(tensor.scalar_type());
+  auto [out_nvte, out_py] = quantizer_cpp->create_tensor(shape, fake_dtype);
+  quantizer_cpp->quantize(tensor_nvte, out_nvte);
+  return out_py;
+}
+
+template <auto act_func, typename... Args>
+py::object scaled_activation_helper(const at::Tensor& input, const at::Tensor& act_scales,
+                                    py::handle quantizer, int shape_divisor, Args&&... args) {
+  auto output = scaled_activation_compute<act_func>(input, act_scales, shape_divisor,
+                                                    std::forward<Args>(args)...);
+  return maybe_quantize(output, quantizer);
+}
+
+template <auto dact_func, typename... Args>
+py::tuple scaled_dactivation_helper(const at::Tensor& grad, const at::Tensor& input,
+                                    const at::Tensor& act_scales, py::handle quantizer,
+                                    bool compute_scale_grad, Args&&... args) {
+  auto [grad_input, grad_scales] = scaled_dactivation_compute<dact_func>(
+      grad, input, act_scales, compute_scale_grad, std::forward<Args>(args)...);
+  return py::make_tuple(maybe_quantize(grad_input, quantizer),
+                        compute_scale_grad ? py::cast(grad_scales) : py::none());
+}
+
+py::object scaled_swiglu(const at::Tensor& input, const at::Tensor& act_scales,
+                         py::handle quantizer, int64_t glu_interleave_size) {
+  return scaled_activation_helper<nvte_scaled_swiglu>(input, act_scales, quantizer,
+                                                      /*shape_divisor=*/2, 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) {
+  return scaled_activation_helper<nvte_scaled_clamped_swiglu>(
+      input, act_scales, quantizer, /*shape_divisor=*/2, limit, alpha, glu_linear_offset,
+      glu_interleave_size);
+}
+
+py::object scaled_srelu(const at::Tensor& input, const at::Tensor& act_scales,
+                        py::handle quantizer) {
+  return scaled_activation_helper<nvte_scaled_srelu>(input, act_scales, quantizer,
+                                                     /*shape_divisor=*/1);
+}
+
+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) {
+  return scaled_dactivation_helper<nvte_scaled_dswiglu>(grad, input, act_scales, quantizer,
+                                                        compute_scale_grad, glu_interleave_size);
+}
+
+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) {
+  return scaled_dactivation_helper<nvte_scaled_clamped_dswiglu>(
+      grad, input, act_scales, quantizer, compute_scale_grad, limit, alpha, glu_linear_offset,
+      glu_interleave_size);
+}
+
+py::tuple scaled_dsrelu(const at::Tensor& grad, const at::Tensor& input,
+                        const at::Tensor& act_scales, py::handle quantizer,
+                        bool compute_scale_grad) {
+  return scaled_dactivation_helper<nvte_scaled_dsrelu>(grad, input, act_scales, quantizer,
+                                                       compute_scale_grad);
+}
+
 }  // namespace pytorch
 }  // namespace transformer_engine

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -285,6 +285,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "Backward of SwiGLU used in GPT OSS", py::arg("grad"), py::arg("fwd_input"),
         py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f,
         py::arg("glu_linear_offset") = 1.0f);
+  /* Scaled activation */
+  m.def("scaled_swiglu", transformer_engine::pytorch::scaled_swiglu, "Scaled SwiGLU activation",
+        py::arg("input"), py::arg("act_scales"), py::arg("quantizer"),
+        py::arg("glu_interleave_size") = 0);
+  m.def("scaled_clamped_swiglu", transformer_engine::pytorch::scaled_clamped_swiglu,
+        "Scaled clamped SwiGLU activation", py::arg("input"), py::arg("act_scales"),
+        py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f,
+        py::arg("glu_linear_offset") = 1.0f, py::arg("glu_interleave_size") = 0);
+  m.def("scaled_srelu", transformer_engine::pytorch::scaled_srelu, "Scaled SReLU activation",
+        py::arg("input"), py::arg("act_scales"), py::arg("quantizer"));
+  m.def("scaled_dswiglu", transformer_engine::pytorch::scaled_dswiglu, "Scaled SwiGLU backward",
+        py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"),
+        py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true);
+  m.def("scaled_clamped_dswiglu", transformer_engine::pytorch::scaled_clamped_dswiglu,
+        "Scaled clamped SwiGLU backward", py::arg("grad"), py::arg("fwd_input"),
+        py::arg("act_scales"), py::arg("quantizer"), py::arg("limit") = 7.0f,
+        py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f,
+        py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true);
+  m.def("scaled_dsrelu", transformer_engine::pytorch::scaled_dsrelu, "Scaled SReLU backward",
+        py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"),
+        py::arg("compute_scale_grad") = true);
   /* DBias + DAct fusions*/
   m.def("dbias_dgelu", transformer_engine::pytorch::dbias_dgelu, "DGeLU + DBias + Quantize",
         py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"));

transformer_engine/pytorch/module/grouped_linear.py

@@ -561,7 +561,7 @@ def _make_grouped_tensor(
         *,
         num_gemms: int,
         split_sizes: torch.Tensor,
-        base_split_offsets: torch.Tensor,
+        tensor_offsets: torch.Tensor,
         last_dim: int,
         dtype: torch.dtype,
     ) -> GroupedTensorStorage:
@@ -573,7 +573,7 @@ def _make_grouped_tensor(
             quantizer=None,
             data=data.reshape(-1),
             first_dims=split_sizes,
-            tensor_offsets=base_split_offsets * last_dim,
+            tensor_offsets=tensor_offsets,
         )
 
     @staticmethod
@@ -704,8 +704,18 @@ def _forward_grouped_tensor(
         out_features = weights[0].size(0)
         weight_requires_grad = weights[0].requires_grad
 
-        split_sizes = m_splits.to(device=device)
-        base_split_offsets = tex.splits_to_offsets(split_sizes, 1)
+        split_sizes, (
+            base_split_offsets,
+            input_tensor_offsets,
+            output_tensor_offsets,
+        ) = tex.splits_to_offsets_multi(
+            m_splits,
+            device,
+            strides=[1, in_features, out_features],
+            include_leading_zero=[True, True, True],
+            dtypes=[torch.int64, torch.int64, torch.int64],
+            bulk_allocate=True,
+        )
 
         inp_view = inp.reshape(-1, in_features)
         x = cast_if_needed(inp_view, activation_dtype)
@@ -716,13 +726,19 @@ def _forward_grouped_tensor(
                 columnwise=is_grad_enabled and weight_requires_grad,
             )
             input_quantizer.optimize_for_gemm = True
-            grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes)
+            grouped_x = tex.group_quantize(
+                x,
+                input_quantizer,
+                num_gemms,
+                split_sizes,
+                tensor_offsets=input_tensor_offsets,
+            )
         else:
             grouped_x = _GroupedLinear._make_grouped_tensor(
                 x,
                 num_gemms=num_gemms,
                 split_sizes=split_sizes,
-                base_split_offsets=base_split_offsets,
+                tensor_offsets=input_tensor_offsets,
                 last_dim=in_features,
                 dtype=activation_dtype,
             )
@@ -751,7 +767,7 @@ def _forward_grouped_tensor(
             out,
             num_gemms=num_gemms,
             split_sizes=split_sizes,
-            base_split_offsets=base_split_offsets,
+            tensor_offsets=output_tensor_offsets,
             last_dim=out_features,
             dtype=activation_dtype,
         )
@@ -796,6 +812,8 @@ def _forward_grouped_tensor(
                 *weights_to_save,
                 split_sizes,
                 base_split_offsets,
+                input_tensor_offsets,
+                output_tensor_offsets,
             )
             ctx.save_for_backward(*tensors_to_save)
             ctx.tensor_objects = tensor_objects
@@ -1217,6 +1235,8 @@ def _backward_grouped_tensor(
         weights = saved_tensors[1 : 1 + N]
         split_sizes = saved_tensors[1 + N]
         base_split_offsets = saved_tensors[2 + N]
+        input_tensor_offsets = saved_tensors[3 + N]
+        output_tensor_offsets = saved_tensors[4 + N]
 
         origin_weights = [None] * N
         main_grads = [None] * N
@@ -1253,20 +1273,22 @@ def _backward_grouped_tensor(
                     grad_output_quantizer,
                     N,
                     split_sizes,
+                    tensor_offsets=output_tensor_offsets,
                 )
             else:
                 grouped_dy = tex.group_quantize(
                     dy_2d,
                     grad_output_quantizer,
                     N,
                     split_sizes,
+                    tensor_offsets=output_tensor_offsets,
                 )
         else:
             grouped_dy = _GroupedLinear._make_grouped_tensor(
                 dy_2d,
                 num_gemms=N,
                 split_sizes=split_sizes,
-                base_split_offsets=base_split_offsets,
+                tensor_offsets=output_tensor_offsets,
                 last_dim=ctx.weights_shape_0,
                 dtype=ctx.activation_dtype,
             )
@@ -1298,7 +1320,7 @@ def _backward_grouped_tensor(
                 dgrad,
                 num_gemms=N,
                 split_sizes=split_sizes,
-                base_split_offsets=base_split_offsets,
+                tensor_offsets=input_tensor_offsets,
                 last_dim=ctx.weights_shape_1,
                 dtype=ctx.activation_dtype,
             )