[Pytorch] Add support for row-wise quanted input for grouped gemm

tests/pytorch/test_grouped_mlp.py

@@ -15,6 +15,7 @@
 import torch
 
 import transformer_engine.pytorch as te
+from transformer_engine.pytorch.constants import TE_DType
 from transformer_engine.pytorch.ops.fused.grouped_mlp import (
     _cudnn_frontend_supports_grouped_gemm_srelu,
     _cudnn_frontend_version_supported,
@@ -436,6 +437,93 @@ def test_grouped_linear(
                     else:
                         assert b_test.grad is None
 
+    @staticmethod
+    def _make_rowwise_mxfp8_wire_input(
+        x_hp: torch.Tensor,
+        group_size: int,
+        split_sizes: torch.Tensor,
+    ) -> "GroupedTensor":
+        """Rowwise-only MXFP8 GroupedTensor with compact scales (FP8 dispatch wire format)."""
+        wire_quantizer = MXFP8Quantizer(
+            fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False
+        )
+        x_wire = tex.group_quantize(
+            x_hp, wire_quantizer, group_size, split_sizes.to(dtype=torch.int64)
+        )
+        # Sanity: the wire tensor is rowwise-only with unswizzled scales.
+        assert x_wire.columnwise_data is None
+        assert not x_wire._with_gemm_swizzled_scales
+        return x_wire
+
+    @pytest.mark.parametrize("weight_requires_grad", (False, True))
+    def test_grouped_linear_prequantized_mxfp8_input(
+        self,
+        *,
+        group_size: int = 4,
+        weight_shape: tuple[int, int] = (256, 256),
+        split_alignment: int = 128,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        weight_requires_grad: bool,
+    ) -> None:
+        """Rowwise-only MXFP8 GroupedTensor input (FP8 token dispatch wire format).
+
+        The input arrives already rowwise-quantized with compact scales. The op
+        must feed the rowwise data to the forward GEMM as-is and manufacture the
+        columnwise copy for the wgrad GEMM. The reference run consumes the
+        *dequantized* wire tensor (the only data a layer can see after FP8
+        dispatch) on the normal quantize-from-BF16 path; because MXFP8 requant
+        is idempotent along the rowwise axis and both paths derive the
+        columnwise copy from the same dequantized data, the two runs must match
+        bit-for-bit.
+        """
+        if not mxfp8_available:
+            pytest.skip(reason_for_no_mxfp8)
+        maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype)
+
+        # Split sizes (including an empty group)
+        split_sizes = [split_alignment * i for i in range(group_size)]
+        random.shuffle(split_sizes)
+        split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device)
+
+        out_features, in_features = weight_shape
+        total_tokens = int(split_sizes.sum().item())
+        in_shape = (total_tokens, in_features)
+
+        # Wire-format input and its exact dequantization (the reference input).
+        x_hp = torch.rand(in_shape, dtype=dtype, device=device) - 0.5
+        x_wire = self._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes)
+        x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view(in_shape)
+
+        dy = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5
+        recipe = make_recipe("mxfp8")
+        op = te.ops.GroupedLinear(
+            group_size, in_features, out_features, bias=False, device=device, dtype=dtype
+        )
+        with torch.no_grad():
+            for param in op.parameters():
+                param.requires_grad_(requires_grad=weight_requires_grad)
+
+        def _run(x):
+            with te.autocast(enabled=True, recipe=recipe):
+                y = op(x, split_sizes)
+            wgrads = []
+            if weight_requires_grad:
+                y.backward(dy)
+                for group_idx in range(group_size):
+                    weight = getattr(op, f"weight{group_idx}")
+                    wgrads.append(weight.grad.detach().clone())
+                    weight.grad = None
+            return y.detach(), wgrads
+
+        y_ref, wgrads_ref = _run(x_ref)
+        y_test, wgrads_test = _run(x_wire)
+
+        # Bit-exact match expected (identical quantized inputs and kernels).
+        torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0)
+        for wgrad_test, wgrad_ref in zip(wgrads_test, wgrads_ref):
+            torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0)
+
     @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16))
     @pytest.mark.parametrize(
         "quantization",
@@ -1076,6 +1164,95 @@ 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("weight_requires_grad", (False, True))
+    def test_grouped_mlp_prequantized_mxfp8_input(
+        self,
+        *,
+        group_size: int = 4,
+        hidden_size: int = 256,
+        split_alignment: int = 256,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        weight_requires_grad: bool,
+    ) -> None:
+        """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor input.
+
+        Production (FP8 token dispatch) path: FC1 receives an already
+        rowwise-quantized input with compact scales. The fused op must feed the
+        rowwise data to the forward GEMM and manufacture FC1's columnwise copy
+        for wgrad. Compared bit-for-bit against a run on the dequantized wire
+        input (see ``test_grouped_linear_prequantized_mxfp8_input``).
+        """
+        if not mxfp8_available:
+            pytest.skip(reason_for_no_mxfp8)
+        if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
+            pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system")
+        maybe_skip_quantization(
+            "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype
+        )
+
+        # Split sizes (including an empty group); sum is a multiple of 128.
+        split_sizes = [split_alignment * i for i in range(group_size)]
+        random.shuffle(split_sizes)
+        split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device)
+        total_tokens = int(split_sizes.sum().item())
+        glu_interleave_size = 32
+
+        # Wire-format FC1 input and its exact dequantization (reference input).
+        x_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+        x_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes)
+        x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view(
+            total_tokens, hidden_size
+        )
+
+        probs = torch.rand((total_tokens,), dtype=dtype, device=device)
+        dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+
+        recipe = make_recipe("mxfp8")
+        with te.quantized_model_init(enabled=True, recipe=recipe):
+            fc1 = te.ops.GroupedLinear(
+                group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype
+            )
+            fc2 = te.ops.GroupedLinear(
+                group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype
+            )
+            module = te.ops.Sequential(
+                fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2
+            )
+        with torch.no_grad():
+            for param in module.parameters():
+                param.requires_grad_(requires_grad=weight_requires_grad)
+
+        def _run(x):
+            with te.autocast(enabled=True, recipe=recipe):
+                y = module(x, split_sizes, probs, split_sizes)
+            fc1_wgrads, fc2_wgrads = [], []
+            if weight_requires_grad:
+                y.backward(dy)
+                for group_idx in range(group_size):
+                    fc1_w = getattr(fc1, f"weight{group_idx}")
+                    fc2_w = getattr(fc2, f"weight{group_idx}")
+                    fc1_wgrads.append(fc1_w.grad.detach().clone())
+                    fc2_wgrads.append(fc2_w.grad.detach().clone())
+                    fc1_w.grad = None
+                    fc2_w.grad = None
+            return y.detach(), fc1_wgrads, fc2_wgrads
+
+        y_ref, fc1_wgrads_ref, fc2_wgrads_ref = _run(x_ref)
+        y_test, fc1_wgrads_test, fc2_wgrads_test = _run(x_wire)
+
+        # Confirm the CuTeDSL fused op was actually formed (not the fallback).
+        forward_ops = module._module_groups[0]._forward_ops
+        assert len(forward_ops) == 1
+        assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU)
+
+        # Bit-exact match expected (identical quantized inputs and kernels).
+        torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0)
+        for wgrad_test, wgrad_ref in zip(fc1_wgrads_test, fc1_wgrads_ref):
+            torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0)
+        for wgrad_test, wgrad_ref in zip(fc2_wgrads_test, fc2_wgrads_ref):
+            torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0)
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list)
     @pytest.mark.parametrize(

transformer_engine/pytorch/ops/_common.py

@@ -10,12 +10,17 @@
 
 import torch
 
+import transformer_engine_torch as tex
 from transformer_engine_torch import FP8TensorMeta
+from ..constants import TE_DType
 from ..torch_version import torch_version
 from ..quantization import FP8GlobalStateManager
 from ..tensor.float8_tensor import Float8Tensor
+from ..tensor.grouped_tensor import GroupedTensor
+from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
+from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage
 from ..quantized_tensor import QuantizedTensorStorage
-from ..utils import canonicalize_dtype
+from ..utils import canonicalize_dtype, round_up_to_nearest_multiple
 
 
 def validate_or_alloc_output(
@@ -66,6 +71,128 @@ def maybe_dequantize(
     return tensor
 
 
+def grouped_storage_from_grouped_tensor(tensor: GroupedTensor) -> GroupedTensorStorage:
+    """Repack a ``GroupedTensor`` into a ``GroupedTensorStorage``.
+
+    ``GroupedTensor`` is a ``torch.Tensor`` subclass, so the CPU offload
+    infrastructure's ``prepare_for_saving`` treats it as a plain tensor and
+    does not decompose it into its component data tensors. By repacking into
+    a ``GroupedTensorStorage`` (not a ``torch.Tensor``), the fuser's
+    ``prepare_for_saving`` call correctly decomposes the activation before
+    ``save_for_backward``.
+    """
+    return GroupedTensorStorage(
+        shape=tensor.logical_shape,
+        dtype=tensor.fake_dtype,
+        num_tensors=tensor.num_tensors,
+        shapes=tensor.tensor_shapes,
+        quantizer=tensor.quantizer,
+        data=tensor.rowwise_data,
+        columnwise_data=tensor.columnwise_data,
+        scale_inv=tensor.scale_inv,
+        columnwise_scale_inv=tensor.columnwise_scale_inv,
+        amax=tensor.amax,
+        columnwise_amax=tensor.columnwise_amax,
+        scale=tensor.scale,
+        first_dims=tensor.first_dims,
+        last_dims=tensor.last_dims,
+        tensor_offsets=tensor.tensor_offsets,
+        offsets=tensor.offsets,
+        scale_inv_offsets=tensor.scale_inv_offsets,
+        columnwise_scale_inv_offsets=tensor.columnwise_scale_inv_offsets,
+        with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales,
+        row_scaled_nvfp4=tensor.row_scaled_nvfp4,
+        nvfp4_use_4over6=tensor.nvfp4_use_4over6,
+        nvfp4_e4m3_max=tensor.nvfp4_e4m3_max,
+    )
+
+
+def prepare_prequantized_mxfp8_grouped_input(
+    grouped_x: GroupedTensorStorage,
+    quantizer: MXFP8Quantizer,
+    num_groups: int,
+    split_sizes: torch.Tensor,
+    dtype: torch.dtype,
+    *,
+    with_columnwise: bool,
+    tensor_offsets: Optional[torch.Tensor] = None,
+) -> None:
+    """Prepare a rowwise-only MXFP8 grouped input for grouped GEMM (in place).
+
+    Supports inputs that arrive already rowwise-quantized (e.g. FP8 token
+    dispatch): the rowwise data is fed to the forward GEMM as-is, while the
+    columnwise copy needed by the wgrad GEMM cannot be derived from the
+    rowwise data (per-block scales differ per direction), so it is
+    manufactured by dequantizing the rowwise data and requantizing
+    columnwise-only. Rowwise scales must arrive in compact (unswizzled)
+    format; they are converted to the GEMM-swizzled layout afterwards.
+    TODO: optimize and fuse the round-trips requant
+    """
+    if grouped_x.rowwise_data is None:
+        raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise data.")
+    if grouped_x._with_gemm_swizzled_scales:
+        raise NotImplementedError(
+            "Pre-quantized MXFP8 grouped input must have scales in compact format."
+        )
+    if grouped_x.columnwise_data is not None:
+        # Columnwise grouped scales have a per-group layout, so the global
+        # single-tensor swizzle below cannot convert them.
+        raise NotImplementedError(
+            "Pre-quantized MXFP8 grouped input with compact scales must be rowwise-only."
+        )
+    if grouped_x.quantizer is not None and grouped_x.quantizer.dtype != quantizer.dtype:
+        # The forward GEMM consumes the input's rowwise data verbatim while the
+        # wgrad GEMM consumes the columnwise copy we manufacture with ``quantizer``.
+        # A dtype mismatch would make the two directions disagree numerically.
+        raise ValueError(
+            f"Pre-quantized MXFP8 grouped input has FP8 dtype {grouped_x.quantizer.dtype}, "
+            f"but the op's input quantizer expects {quantizer.dtype}."
+        )
+
+    # Manufacture columnwise data for the wgrad GEMM: dequantize the rowwise
+    # wire data and requantize columnwise-only.
+    if with_columnwise:
+        hp_x = tex.group_dequantize(grouped_x, TE_DType[dtype])
+        colwise_quantizer = quantizer.copy()
+        colwise_quantizer.set_usage(rowwise=False, columnwise=True)
+        colwise_quantizer.optimize_for_gemm = True
+        colwise_quantizer.internal = True
+        colwise_x = tex.group_quantize(
+            hp_x.rowwise_data.view(grouped_x.logical_shape),
+            colwise_quantizer,
+            num_groups,
+            split_sizes,
+            tensor_offsets=tensor_offsets,
+        )
+        grouped_x.columnwise_data = colwise_x.columnwise_data
+        grouped_x.columnwise_scale_inv = colwise_x.columnwise_scale_inv
+
+    # Convert rowwise scales to the GEMM-swizzled layout. The grouped GEMM
+    # reads activation scales as one (total_tokens, cols) matrix, so the
+    # single-tensor swizzle applies. Swizzling allocates a new scale buffer;
+    # the original compact scales are left untouched.
+    total_tokens, cols = grouped_x.logical_shape
+    scale_shape = (
+        round_up_to_nearest_multiple(total_tokens, 128),
+        round_up_to_nearest_multiple(cols // 32, 4),
+    )
+    tmp = MXFP8Tensor(
+        shape=(total_tokens, cols),
+        dtype=dtype,
+        fp8_dtype=quantizer.dtype,
+        rowwise_data=grouped_x.rowwise_data.view(total_tokens, cols),
+        rowwise_scale_inv=grouped_x.scale_inv.view(scale_shape),
+        columnwise_data=None,
+        columnwise_scale_inv=None,
+        quantizer=quantizer,
+        requires_grad=False,
+        with_gemm_swizzled_scales=False,
+    )
+    tex.swizzle_scales_for_gemm_(tmp)
+    grouped_x.scale_inv = tmp._rowwise_scale_inv.view(-1)
+    grouped_x._with_gemm_swizzled_scales = True
+
+
 def maybe_autocast_dtype(
     *,
     device_type: str = "cuda",