[Common/PyTorch] Grouped weighted-SwiGLU MXFP8 kernel

tests/cpp/operator/CMakeLists.txt

@@ -13,6 +13,7 @@ add_executable(test_operator
                test_qdq.cu
                test_cast_mxfp8.cu
                test_cast_mxfp8_grouped.cu
+               test_cast_mxfp8_grouped_scaled_swiglu.cu
                test_cast_nvfp4_transpose.cu
                test_cast_float8blockwise.cu
                test_cast_float8blockwise_grouped.cu

tests/pytorch/test_grouped_tensor.py

@@ -157,6 +157,105 @@ def align4(x: int) -> int:
     return blocks_x * align4(blocks_y) if columnwise else blocks_y * align4(blocks_x)
 
 
+def _scaled_swiglu_reference(
+    input_2h: torch.Tensor,
+    prob: torch.Tensor,
+    *,
+    limit: Optional[float] = None,
+    alpha: float = 1.0,
+    glu_linear_offset: float = 0.0,
+) -> torch.Tensor:
+    """PyTorch reference for the (clamped) scaled SwiGLU that the grouped kernel recomputes.
+
+    ``limit=None`` selects plain SwiGLU, where ``alpha`` and ``glu_linear_offset`` sit at
+    their identity values. The clamps use ``torch.where`` rather than ``Tensor.clamp`` so
+    this stays character-for-character the formulation that #3424 settled on in
+    tests/pytorch/test_fusible_ops.py.
+
+    Evaluated in fp32 because the kernel also widens its bf16 input to fp32 and quantizes
+    from there; a bf16 reference would carry a rounding step the kernel never takes.
+    """
+    x_glu, x_linear = input_2h.float().chunk(2, dim=-1)
+    if limit is not None:
+        x_glu = torch.where(x_glu > limit, limit, x_glu)
+        x_linear = torch.where(x_linear < -limit, -limit, x_linear)
+        x_linear = torch.where(x_linear > limit, limit, x_linear)
+    out_glu = x_glu * torch.sigmoid(alpha * x_glu)
+    return out_glu * (x_linear + glu_linear_offset) * prob.float().unsqueeze(-1)
+
+
+# MXFP8 scales one e8m0 exponent per 32-row block of a column, and e4m3 carries 3 mantissa
+# bits, so within a block the representable values are spaced 2^-3 of the leading binade
+# apart. Using the block amax as a stand-in for that binade (it is within 2x of it) puts one
+# code step at ``amax * 2^-3``; the bound below allows two, which covers a block whose two
+# e8m0 exponents land one step apart and so quantizes at twice the spacing.
+#
+# Note the scope: this bounds the resolution the kernel is entitled to lose, so it catches a
+# defect that moves some element far. It is measured against each block's largest element, so
+# it is much weaker against a defect that perturbs many elements a little without moving that
+# one -- a wrong ``alpha`` is the case in point, since sigmoid has saturated where the largest
+# element sits. Measured on GB200, a doubled alpha lands at 1.14 steps, inside this bound.
+_MXFP8_STEPS_ALLOWED = 2.0
+_MXFP8_E4M3_MANTISSA_BITS = 3
+
+_MXFP8_SCALE_BLOCK_ROWS = 32
+
+
+def _assert_mxfp8_matches_reference(actual, expected, *, context: str) -> None:
+    """Compare two columnwise-MXFP8 tensors produced by the same quantizer.
+
+    Both sides must come from the same quantizer instance so the scale layout (compact or
+    GEMM-swizzled) and the usage flags agree; only the values being quantized differ, which
+    is what makes a resolution-level comparison meaningful.
+
+    Two independent assertions. The e8m0 scales are compared as raw exponents, where any
+    disagreement beyond one step means the two sides indexed different data rather than
+    rounded differently. The values are then compared after dequantization, elementwise,
+    against the quantization step of their own 32-row block.
+    """
+    actual_scale = actual._columnwise_scale_inv.reshape(-1).view(torch.uint8)
+    expected_scale = expected._columnwise_scale_inv.reshape(-1).view(torch.uint8)
+    assert actual_scale.shape == expected_scale.shape, (
+        f"{context}: scale size {tuple(actual_scale.shape)} != reference"
+        f" {tuple(expected_scale.shape)}"
+    )
+
+    scale_delta = (actual_scale.int() - expected_scale.int()).abs()
+    scale_bad = int((scale_delta > 1).sum())
+    assert scale_bad == 0, (
+        f"{context}: {scale_bad} of {expected_scale.numel()} scales differ by more than one"
+        f" e8m0 step (max delta {int(scale_delta.max())}); that is an indexing or layout"
+        " error, not rounding"
+    )
+
+    actual_deq = actual.dequantize().float()
+    expected_deq = expected.dequantize().float()
+    assert actual_deq.shape == expected_deq.shape, (
+        f"{context}: dequantized shape {tuple(actual_deq.shape)} != reference"
+        f" {tuple(expected_deq.shape)}"
+    )
+    rows, cols = actual_deq.shape
+    assert rows % _MXFP8_SCALE_BLOCK_ROWS == 0, (
+        f"{context}: {rows} rows is not a whole number of {_MXFP8_SCALE_BLOCK_ROWS}-row scale"
+        " blocks, so the per-block bound below would not line up"
+    )
+
+    blocked = (rows // _MXFP8_SCALE_BLOCK_ROWS, _MXFP8_SCALE_BLOCK_ROWS, cols)
+    block_amax = expected_deq.reshape(blocked).abs().amax(dim=1, keepdim=True)
+    tolerance = block_amax * _MXFP8_STEPS_ALLOWED * 2.0**-_MXFP8_E4M3_MANTISSA_BITS
+    deviation = (actual_deq - expected_deq).reshape(blocked).abs()
+
+    n_over = int((deviation > tolerance).sum())
+    floor = torch.finfo(torch.float32).tiny
+    worst = float((deviation / tolerance.clamp(min=floor)).max()) * _MXFP8_STEPS_ALLOWED
+    assert n_over == 0, (
+        f"{context}: {n_over} of {deviation.numel()} elements deviate by more than"
+        f" {_MXFP8_STEPS_ALLOWED:g} MXFP8 code steps of their block"
+        f" (worst {worst:.3f} steps); scales differing by one step:"
+        f" {int((scale_delta != 0).sum())}"
+    )
+
+
 class TestGroupedTensor:
     @staticmethod
     def setup_class(cls) -> None:
@@ -637,6 +736,182 @@ def test_group_quantize_precomputed_offsets(self, output_dbias: bool) -> None:
         assert torch.equal(grouped_output.rowwise_data, expected_output.rowwise_data)
         assert torch.equal(grouped_output.scale_inv, expected_output.scale_inv)
 
+    @pytest.mark.parametrize("optimize_for_gemm", [False, True])
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_group_scaled_swiglu_shapes(self, optimize_for_gemm: bool) -> None:
+        """Test the grouped scaled SwiGLU MXFP8 recompute binding plumbs shapes/dtypes.
+
+        Covers the pybind layer and its rejection paths: a [N, 2H] input plus a [N] prob must
+        come back as a columnwise-MXFP8 [N, H] grouped output, and non-contiguous or
+        wrong-dtype operands must raise. Values are checked in
+        ``test_group_scaled_swiglu_matches_reference``.
+        """
+        num_tensors = 3
+        last_dim = 256
+        split_sizes_list = [128, 384, 512]
+        total_tokens = sum(split_sizes_list)
+
+        input_2h = torch.randn(total_tokens, 2 * last_dim, dtype=torch.bfloat16, device="cuda")
+        # prob rides in the model dtype, matching TE's cuDNN fc1_prob_tensor convention.
+        prob = torch.rand(total_tokens, dtype=torch.bfloat16, device="cuda")
+        first_dims = torch.tensor(split_sizes_list, dtype=torch.int64, device="cuda")
+
+        quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)
+        quantizer.set_usage(rowwise=False, columnwise=True)
+        quantizer.optimize_for_gemm = optimize_for_gemm
+
+        grouped_output = tex.group_scaled_swiglu(input_2h, prob, quantizer, num_tensors, first_dims)
+
+        outputs = grouped_output.split_into_quantized_tensors()
+        assert len(outputs) == num_tensors
+        for rows, output in zip(split_sizes_list, outputs):
+            assert output.shape == (rows, last_dim)
+            assert output._columnwise_data.numel() == rows * last_dim
+            # One e8m0 exponent per 32-row block of every column. Both the compact and the
+            # GEMM-swizzled layout need the same number of scales.
+            assert output._columnwise_scale_inv.numel() == (rows // 32) * last_dim
+
+        with pytest.raises(RuntimeError):
+            tex.group_scaled_swiglu(input_2h, prob.float(), quantizer, num_tensors, first_dims)
+
+        # Both operands reach the kernel as raw pointers over a densely packed range, so a
+        # strided view must be rejected instead of being read as if it were contiguous.
+        wide = torch.randn(total_tokens, 4 * last_dim, dtype=torch.bfloat16, device="cuda")
+        with pytest.raises(RuntimeError):
+            tex.group_scaled_swiglu(
+                wide[:, : 2 * last_dim], prob, quantizer, num_tensors, first_dims
+            )
+
+        strided_prob = torch.rand(2 * total_tokens, dtype=torch.bfloat16, device="cuda")[::2]
+        with pytest.raises(RuntimeError):
+            tex.group_scaled_swiglu(input_2h, strided_prob, quantizer, num_tensors, first_dims)
+
+    @pytest.mark.parametrize("optimize_for_gemm", [False, True])
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_group_scaled_clamped_swiglu_shapes(self, optimize_for_gemm: bool) -> None:
+        """Test the clamped variant's binding plumbs shapes, dtypes and clamp parameters.
+
+        Covers the pybind layer, the identity case that must collapse onto the unclamped
+        kernel, and rejection of a non-positive limit. Values are checked in
+        ``test_group_scaled_swiglu_matches_reference``.
+        """
+        num_tensors = 3
+        last_dim = 256
+        split_sizes_list = [128, 384, 512]
+        total_tokens = sum(split_sizes_list)
+
+        input_2h = torch.randn(total_tokens, 2 * last_dim, dtype=torch.bfloat16, device="cuda")
+        prob = torch.rand(total_tokens, dtype=torch.bfloat16, device="cuda")
+        first_dims = torch.tensor(split_sizes_list, dtype=torch.int64, device="cuda")
+
+        quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)
+        quantizer.set_usage(rowwise=False, columnwise=True)
+        quantizer.optimize_for_gemm = optimize_for_gemm
+
+        grouped_output = tex.group_scaled_clamped_swiglu(
+            input_2h, prob, quantizer, num_tensors, 7.0, 1.702, 1.0, first_dims
+        )
+
+        outputs = grouped_output.split_into_quantized_tensors()
+        assert len(outputs) == num_tensors
+        for rows, output in zip(split_sizes_list, outputs):
+            assert output.shape == (rows, last_dim)
+            assert output._columnwise_data.numel() == rows * last_dim
+            assert output._columnwise_scale_inv.numel() == (rows // 32) * last_dim
+
+        # A clamp wider than any input, with alpha and offset at their identity values,
+        # must reproduce the unclamped kernel exactly.
+        wide_limit = float(input_2h.abs().max()) + 1.0
+        clamped_identity = tex.group_scaled_clamped_swiglu(
+            input_2h, prob, quantizer, num_tensors, wide_limit, 1.0, 0.0, first_dims
+        )
+        plain = tex.group_scaled_swiglu(input_2h, prob, quantizer, num_tensors, first_dims)
+        for lhs, rhs in zip(
+            clamped_identity.split_into_quantized_tensors(),
+            plain.split_into_quantized_tensors(),
+        ):
+            assert torch.equal(lhs._columnwise_data, rhs._columnwise_data)
+            assert torch.equal(lhs._columnwise_scale_inv, rhs._columnwise_scale_inv)
+
+        for bad_limit in (0.0, -1.0):
+            with pytest.raises(RuntimeError):
+                tex.group_scaled_clamped_swiglu(
+                    input_2h, prob, quantizer, num_tensors, bad_limit, 1.702, 1.0, first_dims
+                )
+
+    @pytest.mark.parametrize("optimize_for_gemm", [False, True])
+    @pytest.mark.parametrize("clamped", [False, True])
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_group_scaled_swiglu_matches_reference(
+        self, clamped: bool, optimize_for_gemm: bool
+    ) -> None:
+        """Check both grouped kernels against a quantized PyTorch reference, expert by expert.
+
+        Each expert's reference is computed from that expert's own rows and quantized with
+        the very quantizer the kernel used, so a wrong expert offset, a dropped prob multiply
+        or a mis-applied clamp all surface as codes differing far beyond the rounding budget
+        that ``_assert_mxfp8_matches_reference`` allows.
+        """
+        limit, alpha, glu_linear_offset = (7.0, 1.702, 1.0) if clamped else (None, 1.0, 0.0)
+
+        # Row counts are multiples of the kernel's 128-row chunk, and deliberately unequal so
+        # that a per-expert offset computed from the wrong group is caught.
+        split_sizes_list = [128, 256, 384, 512]
+        num_tensors = len(split_sizes_list)
+        last_dim = 256
+        total_tokens = sum(split_sizes_list)
+
+        # Scaled up so the tails reach past the clamp limit: standard normals essentially
+        # never exceed 7.0, which would leave the clamped path untested on its own numerics.
+        input_2h = 4.0 * torch.randn(
+            total_tokens, 2 * last_dim, dtype=torch.bfloat16, device="cuda"
+        )
+        prob = torch.rand(total_tokens, dtype=torch.bfloat16, device="cuda")
+        first_dims = torch.tensor(split_sizes_list, dtype=torch.int64, device="cuda")
+
+        quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)
+        quantizer.set_usage(rowwise=False, columnwise=True)
+        quantizer.optimize_for_gemm = optimize_for_gemm
+
+        if clamped:
+            grouped_output = tex.group_scaled_clamped_swiglu(
+                input_2h,
+                prob,
+                quantizer,
+                num_tensors,
+                limit,
+                alpha,
+                glu_linear_offset,
+                first_dims,
+            )
+        else:
+            grouped_output = tex.group_scaled_swiglu(
+                input_2h, prob, quantizer, num_tensors, first_dims
+            )
+
+        outputs = grouped_output.split_into_quantized_tensors()
+        assert len(outputs) == num_tensors
+
+        row = 0
+        for expert, (rows, output) in enumerate(zip(split_sizes_list, outputs)):
+            reference = _scaled_swiglu_reference(
+                input_2h[row : row + rows],
+                prob[row : row + rows],
+                limit=limit,
+                alpha=alpha,
+                glu_linear_offset=glu_linear_offset,
+            )
+            assert reference.shape == (rows, last_dim)
+            _assert_mxfp8_matches_reference(
+                output,
+                quantizer(reference),
+                context=(
+                    f"expert {expert} ({rows}x{last_dim}, clamped={clamped},"
+                    f" optimize_for_gemm={optimize_for_gemm})"
+                ),
+            )
+            row += rows
+
     @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
     def test_bgrad_group_quantize_zero_size_tensor(self) -> None:
         """Test bgrad_group_quantize handles zero-row input without error."""

transformer_engine/common/activation/swiglu_grouped.cu

@@ -15,6 +15,29 @@ void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cu
                                                                        stream);
 }
 
+void nvte_group_scaled_swiglu(const NVTEGroupedTensor input, const NVTETensor prob,
+                              NVTEGroupedTensor output, cudaStream_t stream) {
+  NVTE_API_CALL(nvte_group_scaled_swiglu);
+  using namespace transformer_engine;
+  // Scaled SwiGLU recompute: (silu(act) * gate) * prob -> columnwise MXFP8.
+  // The kernel halves prob, so the activation is the doubled approximate SiLU.
+  Empty e = {};
+  dispatch::group_scaled_swiglu_fwd_helper<Empty, silu_approx_x2<fp32, fp32>>(input, prob, output,
+                                                                              e, nullptr, stream);
+}
+
+void nvte_group_scaled_clamped_swiglu(const NVTEGroupedTensor input, const NVTETensor prob,
+                                      NVTEGroupedTensor output, float limit, float alpha,
+                                      float glu_linear_offset, cudaStream_t stream) {
+  NVTE_API_CALL(nvte_group_scaled_clamped_swiglu);
+  using namespace transformer_engine;
+  // glu_linear_offset is a direct argument rather than hard-coded to 1.0f with a _v2
+  // alongside: this entry point is new, so there is no prior behavior to preserve.
+  ClampedSwiGLUParam param = {limit, alpha, glu_linear_offset};
+  dispatch::group_scaled_swiglu_fwd_helper<ClampedSwiGLUParam, clamped_silu_approx_x2<fp32, fp32>>(
+      input, prob, output, param, nullptr, stream);
+}
+
 void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input,
                       NVTEGroupedTensor output, cudaStream_t stream) {
   NVTE_API_CALL(nvte_group_dsilu);

transformer_engine/common/cast/dispatch/quantize.cuh

@@ -21,6 +21,7 @@
 #include "../fp8/quantize_fp8.cuh"
 #include "../fp8_blockwise/group_quantize_fp8_blockwise.cuh"
 #include "../mxfp8/group_quantize_mxfp8.cuh"
+#include "../mxfp8/group_scaled_swiglu_mxfp8.cuh"
 #include "../mxfp8/quantize_mxfp8.cuh"
 #include "../nvfp4/group_quantize_transpose_nvfp4.cuh"
 #include "../nvfp4/quantize_4over6_nvfp4.cuh"
@@ -520,6 +521,46 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor
   }
 }
 
+// Grouped scaled SwiGLU recompute: input [N, 2H] ([act|gate]) + prob [N]
+// -> columnwise MXFP8 of (silu(act) * gate) * prob.
+template <typename ParamOP, float (*OP)(float, const ParamOP &)>
+void group_scaled_swiglu_fwd_helper(const NVTEGroupedTensor input, const NVTETensor prob,
+                                    NVTEGroupedTensor output, const ParamOP &p,
+                                    const NVTEQuantizationConfig quant_config,
+                                    cudaStream_t stream) {
+  using namespace detail;
+
+  NVTEScalingMode scaling_mode = nvte_grouped_tensor_scaling_mode(output);
+
+  const GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input);
+  GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output);
+  const Tensor *prob_tensor = convertNVTETensorCheck(prob);
+
+  // Quantization config
+  QuantizationConfig quant_config_cpp;
+  if (quant_config != nullptr) {
+    quant_config_cpp = *reinterpret_cast<QuantizationConfig *>(quant_config);
+  }
+
+  // Noop flag (graph-safe skip)
+  Tensor dummy_tensor;
+  Tensor *noop_tensor = &dummy_tensor;
+  if (quant_config_cpp.noop_tensor != nullptr) {
+    noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor);
+  }
+
+  switch (scaling_mode) {
+    case NVTE_MXFP8_1D_SCALING: {
+      mxfp8::group_scaled_swiglu<ParamOP, OP>(input_tensor, prob_tensor, noop_tensor, output_tensor,
+                                              p, &quant_config_cpp, stream);
+      break;
+    }
+    default:
+      NVTE_ERROR("group_scaled_swiglu only supports NVTE_MXFP8_1D_SCALING, got: " +
+                 to_string(scaling_mode) + ".");
+  }
+}
+
 template <bool IS_DBIAS, bool IS_DACT, typename ParamOP, float (*OP)(float, const ParamOP &)>
 void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTensor input,
                                NVTEGroupedTensor output, NVTEGroupedTensor dbias,