From 11e90074d350262a42c65e86b584cf7f4c8f111c Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Thu, 23 Jul 2026 11:20:47 -0700
Subject: [PATCH 01/18] add support for row-wise quanted input for grouped gemm

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py             | 177 ++++++++++++++++++
 transformer_engine/pytorch/ops/_common.py     | 129 ++++++++++++-
 .../pytorch/ops/basic/grouped_linear.py       |  44 ++++-
 .../pytorch/ops/fused/grouped_mlp.py          |  60 +++---
 4 files changed, 375 insertions(+), 35 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index 76550db5f8..07dd1587a2 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/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(
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index 607346ce30..d45301df10 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/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",
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index af3f8b5930..7232a8fff9 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -46,8 +46,10 @@
     get_accumulate_flag_in_param,
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
+    grouped_storage_from_grouped_tensor,
     is_quantized_tensor,
     maybe_dequantize,
+    prepare_prequantized_mxfp8_grouped_input,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
 )
@@ -1172,6 +1174,11 @@ def _fuser_forward_split_quantize(
         out_buffer: Optional[torch.Tensor] = None,
     ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]:
         """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow."""
+        if isinstance(input_, GroupedTensor):
+            raise NotImplementedError(
+                "Pre-quantized GroupedTensor input is only supported on the "
+                "graph-safe grouped-tensor path."
+            )
         num_groups = self.num_groups
         has_bias = self.has_bias
 
@@ -1296,11 +1303,42 @@ def _fuser_forward_grouped_tensor(
 
         # Flatten to 2D so the first dim is the total token count.
         original_shape = list(input_.size())
-        x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features)
-        total_tokens = x.size(0)
+        prequantized_mxfp8_input = (
+            with_quantized_compute
+            and isinstance(input_, GroupedTensor)
+            and isinstance(input_quantizers[0], MXFP8Quantizer)
+            and isinstance(input_.quantizer, MXFP8Quantizer)
+        )
+        if prequantized_mxfp8_input:
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, in_features) layout; just validate the shape.
+            if input_.dim() != 2 or input_.size(-1) != self.in_features:
+                raise ValueError(
+                    "GroupedTensor input must have shape (total_tokens, "
+                    f"{self.in_features}), but got {tuple(input_.size())}."
+                )
+            total_tokens = input_.size(0)
+        else:
+            x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features)
+            total_tokens = x.size(0)
 
         # Build the input GroupedTensor.
-        if with_quantized_compute:
+        if prequantized_mxfp8_input:
+            # Rowwise-only MXFP8 input (e.g. FP8 token dispatch): feed the
+            # rowwise data to the forward GEMM as-is, manufacture the
+            # columnwise copy needed by the wgrad GEMM, and swizzle the
+            # rowwise scales for the GEMM.
+            grouped_x = grouped_storage_from_grouped_tensor(input_)
+            prepare_prequantized_mxfp8_grouped_input(
+                grouped_x,
+                input_quantizers[0],
+                num_groups,
+                split_sizes,
+                dtype,
+                with_columnwise=weight_requires_grad,
+                tensor_offsets=base_split_offsets * self.in_features,
+            )
+        elif with_quantized_compute:
             input_quantizer = input_quantizers[0]
             input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
             input_quantizer.optimize_for_gemm = True
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 83954a9b3d..00c2dca213 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -52,8 +52,10 @@
     get_accumulate_flag_in_param,
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
+    grouped_storage_from_grouped_tensor,
     is_quantized_tensor,
     maybe_dequantize,
+    prepare_prequantized_mxfp8_grouped_input,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
 )
@@ -919,7 +921,16 @@ def fuser_forward(
         # Tensor properties
         fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features)
         fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features)
-        input_ = input_.reshape(-1, fc1_weight_shape[1])
+        if isinstance(input_, GroupedTensor):
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, in_features) layout; just validate the shape.
+            if input_.dim() != 2 or input_.size(-1) != fc1_weight_shape[1]:
+                raise ValueError(
+                    "GroupedTensor input must have shape (total_tokens, "
+                    f"{fc1_weight_shape[1]}), but got {tuple(input_.size())}."
+                )
+        else:
+            input_ = input_.reshape(-1, fc1_weight_shape[1])
         in_shape = list(input_.size())
         if in_shape[0] % 128 != 0:
             raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).")
@@ -1080,36 +1091,23 @@ def fuser_forward(
             or isinstance(fc1_input_quantizer, NVFP4Quantizer)
             and isinstance(input_quantizer, NVFP4Quantizer)
         ):
-            # 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), we
-            # ensure the fuser's prepare_for_saving call correctly decomposes
-            # the activation before save_for_backward.
-            grouped_fc1_x = GroupedTensorStorage(
-                shape=input_.logical_shape,
-                dtype=input_.fake_dtype,
-                num_tensors=input_.num_tensors,
-                shapes=input_.tensor_shapes,
-                quantizer=input_.quantizer,
-                data=input_.rowwise_data,
-                columnwise_data=input_.columnwise_data,
-                scale_inv=input_.scale_inv,
-                columnwise_scale_inv=input_.columnwise_scale_inv,
-                amax=input_.amax,
-                columnwise_amax=input_.columnwise_amax,
-                scale=input_.scale,
-                first_dims=input_.first_dims,
-                last_dims=input_.last_dims,
-                tensor_offsets=input_.tensor_offsets,
-                offsets=input_.offsets,
-                scale_inv_offsets=input_.scale_inv_offsets,
-                columnwise_scale_inv_offsets=input_.columnwise_scale_inv_offsets,
-                with_gemm_swizzled_scales=input_._with_gemm_swizzled_scales,
-                row_scaled_nvfp4=input_.row_scaled_nvfp4,
-                nvfp4_use_4over6=input_.nvfp4_use_4over6,
-                nvfp4_e4m3_max=input_.nvfp4_e4m3_max,
-            )
+            grouped_fc1_x = grouped_storage_from_grouped_tensor(input_)
+            if (
+                isinstance(fc1_input_quantizer, MXFP8Quantizer)
+                and not grouped_fc1_x._with_gemm_swizzled_scales
+            ):
+                # Rowwise-only MXFP8 input (e.g. FP8 token dispatch):
+                # manufacture the columnwise copy needed by the wgrad GEMM
+                # and swizzle the rowwise scales for the forward GEMM.
+                prepare_prequantized_mxfp8_grouped_input(
+                    grouped_fc1_x,
+                    fc1_input_quantizer,
+                    num_groups,
+                    split_sizes,
+                    dtype,
+                    with_columnwise=weight_requires_grad,
+                    tensor_offsets=fc1_x_tensor_offsets,
+                )
         else:
             fc1_x = maybe_dequantize(input_, dtype)
             grouped_fc1_x = _group_quantize_for_grouped_mlp(

From d234cd81ca749f6765c7584d7be2af2025e93840 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Sat, 25 Jul 2026 12:05:07 -0700
Subject: [PATCH 02/18] doc change

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 transformer_engine/pytorch/ops/_common.py     | 51 +++++++++++--------
 .../pytorch/ops/basic/grouped_linear.py       |  4 +-
 .../pytorch/ops/fused/grouped_mlp.py          |  4 +-
 3 files changed, 34 insertions(+), 25 deletions(-)

diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index d45301df10..7cb1ff2a2a 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -12,7 +12,7 @@
 
 import transformer_engine_torch as tex
 from transformer_engine_torch import FP8TensorMeta
-from ..constants import TE_DType
+from ..constants import MXFP8_BLOCK_SCALING_SIZE, TE_DType
 from ..torch_version import torch_version
 from ..quantization import FP8GlobalStateManager
 from ..tensor.float8_tensor import Float8Tensor
@@ -20,7 +20,7 @@
 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, round_up_to_nearest_multiple
+from ..utils import canonicalize_dtype
 
 
 def validate_or_alloc_output(
@@ -107,7 +107,7 @@ def grouped_storage_from_grouped_tensor(tensor: GroupedTensor) -> GroupedTensorS
     )
 
 
-def prepare_prequantized_mxfp8_grouped_input(
+def prepare_prequantized_mxfp8_input_for_gemm(
     grouped_x: GroupedTensorStorage,
     quantizer: MXFP8Quantizer,
     num_groups: int,
@@ -117,28 +117,29 @@ def prepare_prequantized_mxfp8_grouped_input(
     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.
+    """Make an already-quantized MXFP8 grouped input GEMM-ready (in place).
+
+    For inputs that arrive rowwise-quantized (e.g. FP8 token dispatch). On return
+    the rowwise scales are GEMM-swizzled; columnwise data/scales are populated iff
+    ``with_columnwise``, manufactured by dequantize + columnwise-only requantize
+    since the two directions scale along perpendicular axes.
+
+    Requires rowwise-only input with unswizzled scales and per-group token counts
+    that are multiples of 128 (caller contract: ``split_sizes`` is on device).
+
     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.scale_inv is None:
+        raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise scales.")
     if grouped_x._with_gemm_swizzled_scales:
-        raise NotImplementedError(
-            "Pre-quantized MXFP8 grouped input must have scales in compact format."
-        )
+        raise NotImplementedError("Pre-quantized MXFP8 grouped input must have unswizzled scales.")
     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."
+            "Pre-quantized MXFP8 grouped input with unswizzled 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
@@ -170,12 +171,20 @@ def prepare_prequantized_mxfp8_grouped_input(
     # 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.
+    # the original unswizzled scales are left untouched.
+    # 128-alignment (see docstring) means the scale array needs no padding.
     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),
-    )
+    if total_tokens % 128 != 0 or cols % 128 != 0:
+        raise ValueError(
+            "Pre-quantized MXFP8 grouped input requires dims that are multiples of 128, "
+            f"but got ({total_tokens}, {cols})."
+        )
+    scale_shape = (total_tokens, cols // MXFP8_BLOCK_SCALING_SIZE)
+    if grouped_x.scale_inv.numel() != math.prod(scale_shape):
+        raise ValueError(
+            f"Pre-quantized MXFP8 grouped input has {grouped_x.scale_inv.numel()} rowwise "
+            f"scales, but expected {math.prod(scale_shape)} for shape {scale_shape}."
+        )
     tmp = MXFP8Tensor(
         shape=(total_tokens, cols),
         dtype=dtype,
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index 7232a8fff9..1c946ef020 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -49,7 +49,7 @@
     grouped_storage_from_grouped_tensor,
     is_quantized_tensor,
     maybe_dequantize,
-    prepare_prequantized_mxfp8_grouped_input,
+    prepare_prequantized_mxfp8_input_for_gemm,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
 )
@@ -1329,7 +1329,7 @@ def _fuser_forward_grouped_tensor(
             # columnwise copy needed by the wgrad GEMM, and swizzle the
             # rowwise scales for the GEMM.
             grouped_x = grouped_storage_from_grouped_tensor(input_)
-            prepare_prequantized_mxfp8_grouped_input(
+            prepare_prequantized_mxfp8_input_for_gemm(
                 grouped_x,
                 input_quantizers[0],
                 num_groups,
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 00c2dca213..4d9d3972f6 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -55,7 +55,7 @@
     grouped_storage_from_grouped_tensor,
     is_quantized_tensor,
     maybe_dequantize,
-    prepare_prequantized_mxfp8_grouped_input,
+    prepare_prequantized_mxfp8_input_for_gemm,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
 )
@@ -1099,7 +1099,7 @@ def fuser_forward(
                 # Rowwise-only MXFP8 input (e.g. FP8 token dispatch):
                 # manufacture the columnwise copy needed by the wgrad GEMM
                 # and swizzle the rowwise scales for the forward GEMM.
-                prepare_prequantized_mxfp8_grouped_input(
+                prepare_prequantized_mxfp8_input_for_gemm(
                     grouped_fc1_x,
                     fc1_input_quantizer,
                     num_groups,

From 376b94c43c22e17e771d71271c5ae591df145c1d Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Sat, 25 Jul 2026 21:36:28 -0700
Subject: [PATCH 03/18] implement for backward

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py             | 186 ++++++++++++++++++
 transformer_engine/pytorch/ops/_common.py     |  48 ++++-
 .../pytorch/ops/basic/grouped_linear.py       |  51 ++++-
 .../pytorch/ops/fused/grouped_mlp.py          |  46 ++++-
 4 files changed, 316 insertions(+), 15 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index 07dd1587a2..487d080566 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -230,6 +230,23 @@ def make_reference_and_test_tensors(
     return ref, test
 
 
+class _InjectGrad(torch.autograd.Function):
+    """Replace the gradient flowing into ``x`` with ``grad``.
+
+    Mirrors how an FP8 token dispatch delivers a pre-quantized ``GroupedTensor``
+    grad output: the downstream op simply returns one as its grad input.
+    """
+
+    @staticmethod
+    def forward(ctx, x, grad):  # pylint: disable=arguments-differ
+        ctx.injected_grad = grad
+        return x
+
+    @staticmethod
+    def backward(ctx, grad_output):  # pylint: disable=arguments-differ
+        return ctx.injected_grad, None
+
+
 class TestGroupedLinearOp:
     """Tests for advanced features with grouped linear basic op"""
 
@@ -524,6 +541,77 @@ def _run(x):
         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("bias", (False, True))
+    def test_grouped_linear_prequantized_mxfp8_grad(
+        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",
+        bias: bool,
+    ) -> None:
+        """Rowwise-only MXFP8 GroupedTensor grad output (FP8 token dispatch, backward).
+
+        Mirrors ``test_grouped_linear_prequantized_mxfp8_input`` on the backward
+        side: the rowwise data feeds the dgrad GEMM and the columnwise copy is
+        manufactured for wgrad. With ``bias`` the per-group bias gradient rides
+        the columnwise stage of the quantize kernel.
+        """
+        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())
+
+        x = torch.rand((total_tokens, in_features), dtype=dtype, device=device) - 0.5
+        dy_hp = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5
+
+        # Wire-format grad output and its exact dequantization (the reference grad).
+        dy_wire = self._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes)
+        dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view(
+            total_tokens, out_features
+        )
+
+        recipe = make_recipe("mxfp8")
+        op = te.ops.GroupedLinear(
+            group_size, in_features, out_features, bias=bias, device=device, dtype=dtype
+        )
+
+        def _run(grad):
+            x_in = x.detach().clone().requires_grad_()
+            with te.autocast(enabled=True, recipe=recipe):
+                y = op(x_in, split_sizes)
+            # Deliver ``grad`` as the op's grad output, as FP8 dispatch would.
+            _InjectGrad.apply(y, grad).backward(torch.ones_like(y))
+            wgrads, bgrads = [], []
+            for group_idx in range(group_size):
+                weight = getattr(op, f"weight{group_idx}")
+                wgrads.append(weight.grad.detach().clone())
+                weight.grad = None
+                if bias:
+                    bias_param = getattr(op, f"bias{group_idx}")
+                    bgrads.append(bias_param.grad.detach().clone())
+                    bias_param.grad = None
+            return x_in.grad, wgrads, bgrads
+
+        dx_ref, wgrads_ref, bgrads_ref = _run(dy_ref)
+        dx_test, wgrads_test, bgrads_test = _run(dy_wire)
+
+        # Bit-exact match expected (identical quantized grads and kernels).
+        torch.testing.assert_close(dx_test, dx_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)
+        for bgrad_test, bgrad_ref in zip(bgrads_test, bgrads_ref):
+            torch.testing.assert_close(bgrad_test, bgrad_ref, rtol=0, atol=0)
+
     @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16))
     @pytest.mark.parametrize(
         "quantization",
@@ -1253,6 +1341,104 @@ def _run(x):
         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))
+    def test_grouped_mlp_prequantized_mxfp8_grad(
+        self,
+        *,
+        group_size: int = 4,
+        hidden_size: int = 256,
+        split_alignment: int = 256,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        bias: bool,
+    ) -> None:
+        """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor grad output.
+
+        FC2 receives the pre-quantized grad, as an FP8 token dispatch delivers it
+        on the backward pass. With ``bias`` FC2 uses ``scale_bias``, whose
+        dbias/dscales need the dequantized grad rather than the fused dbias.
+        """
+        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
+
+        x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+        probs = torch.rand((total_tokens,), dtype=dtype, device=device)
+        dy_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+
+        # Wire-format grad output and its exact dequantization (the reference grad).
+        dy_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes)
+        dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view(
+            total_tokens, hidden_size
+        )
+
+        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=bias, device=device, dtype=dtype
+            )
+            fc2 = te.ops.GroupedLinear(
+                group_size,
+                hidden_size,
+                hidden_size,
+                bias=bias,
+                device=device,
+                dtype=dtype,
+                scale_bias=bias,
+            )
+            module = te.ops.Sequential(
+                fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2
+            )
+
+        def _run(grad):
+            x_in = x.detach().clone().requires_grad_()
+            fc2_extra = (split_sizes, probs) if bias else (split_sizes,)
+            with te.autocast(enabled=True, recipe=recipe):
+                y = module(x_in, split_sizes, probs, *fc2_extra)
+            _InjectGrad.apply(y, grad).backward(torch.ones_like(y))
+            grads = [("dx", x_in.grad)]
+            for name, fc in (("fc1", fc1), ("fc2", fc2)):
+                for group_idx in range(group_size):
+                    weight = getattr(fc, f"weight{group_idx}")
+                    grads.append((f"{name}_w{group_idx}", weight.grad.detach().clone()))
+                    weight.grad = None
+                    if bias:
+                        bias_param = getattr(fc, f"bias{group_idx}")
+                        grads.append((f"{name}_b{group_idx}", bias_param.grad.detach().clone()))
+                        bias_param.grad = None
+            return grads
+
+        grads_ref = _run(dy_ref)
+        grads_test = _run(dy_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 grads and kernels), except
+        # bias gradients: the fused kernels generate them with an accumulation
+        # that is not reproducible run to run (two runs on identical inputs differ
+        # by one BF16 ulp). Same tolerances as
+        # ``test_grouped_mlp_single_weight_numerics``.
+        bias_tols = {"rtol": 0.05, "atol": 0.015625}
+        for (name, grad_test), (_, grad_ref) in zip(grads_test, grads_ref):
+            if "_b" in name:
+                torch.testing.assert_close(grad_test, grad_ref, **bias_tols)
+            else:
+                torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0)
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list)
     @pytest.mark.parametrize(
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index 7cb1ff2a2a..53e01365cf 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -115,8 +115,9 @@ def prepare_prequantized_mxfp8_input_for_gemm(
     dtype: torch.dtype,
     *,
     with_columnwise: bool,
+    with_dbias: bool = False,
     tensor_offsets: Optional[torch.Tensor] = None,
-) -> None:
+) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
     """Make an already-quantized MXFP8 grouped input GEMM-ready (in place).
 
     For inputs that arrive rowwise-quantized (e.g. FP8 token dispatch). On return
@@ -124,6 +125,13 @@ def prepare_prequantized_mxfp8_input_for_gemm(
     ``with_columnwise``, manufactured by dequantize + columnwise-only requantize
     since the two directions scale along perpendicular axes.
 
+    Returns ``(dbias, dequantized)``. ``dbias`` is the per-group bias gradient,
+    produced only when ``with_dbias``: the grouped quantize kernel accumulates it
+    in the columnwise stage it is already running, so it costs no extra pass.
+    ``dequantized`` is the high-precision tensor materialized for the requantize
+    (``None`` if none was needed), for callers that cannot use the fused dbias
+    (e.g. ``scale_bias``, whose dbias/dscales depend on routing probabilities).
+
     Requires rowwise-only input with unswizzled scales and per-group token counts
     that are multiples of 128 (caller contract: ``split_sizes`` is on device).
 
@@ -150,21 +158,41 @@ def prepare_prequantized_mxfp8_input_for_gemm(
             f"but the op's input quantizer expects {quantizer.dtype}."
         )
 
+    if with_dbias and not with_columnwise:
+        # The fused dbias is accumulated by the quantize kernel's columnwise stage.
+        raise NotImplementedError(
+            "Pre-quantized MXFP8 grouped input cannot produce a fused dbias without "
+            "a columnwise copy."
+        )
+
     # Manufacture columnwise data for the wgrad GEMM: dequantize the rowwise
     # wire data and requantize columnwise-only.
+    dbias = None
+    dequantized = None
     if with_columnwise:
-        hp_x = tex.group_dequantize(grouped_x, TE_DType[dtype])
+        dequantized = tex.group_dequantize(grouped_x, TE_DType[dtype]).rowwise_data.view(
+            grouped_x.logical_shape
+        )
         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,
-        )
+        if with_dbias:
+            colwise_x, dbias = tex.bgrad_group_quantize(
+                dequantized,
+                colwise_quantizer,
+                num_groups,
+                split_sizes,
+                tensor_offsets=tensor_offsets,
+            )
+        else:
+            colwise_x = tex.group_quantize(
+                dequantized,
+                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
 
@@ -201,6 +229,8 @@ def prepare_prequantized_mxfp8_input_for_gemm(
     grouped_x.scale_inv = tmp._rowwise_scale_inv.view(-1)
     grouped_x._with_gemm_swizzled_scales = True
 
+    return dbias, dequantized
+
 
 def maybe_autocast_dtype(
     *,
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index 1c946ef020..6eaf786a2c 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -1691,8 +1691,25 @@ def _fuser_backward_grouped_tensor(
 
         # Flatten grad_output to 2D (total_tokens, out_features)
         # to figure out total tokens.
-        dy_2d = grad_output.reshape(-1, self.out_features)
-        total_tokens = dy_2d.size(0)
+        prequantized_mxfp8_grad = (
+            with_quantized_compute
+            and isinstance(grad_output, GroupedTensor)
+            and isinstance(ctx.grad_output_quantizers[0], MXFP8Quantizer)
+            and isinstance(grad_output.quantizer, MXFP8Quantizer)
+        )
+        if prequantized_mxfp8_grad:
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, out_features) layout; just validate the shape.
+            if grad_output.dim() != 2 or grad_output.size(-1) != self.out_features:
+                raise ValueError(
+                    "GroupedTensor grad output must have shape (total_tokens, "
+                    f"{self.out_features}), but got {tuple(grad_output.size())}."
+                )
+            dy_2d = None
+            total_tokens = grad_output.size(0)
+        else:
+            dy_2d = grad_output.reshape(-1, self.out_features)
+            total_tokens = dy_2d.size(0)
 
         # Build the grad_output GroupedTensor.
         # Optionally get dbias is fusion available with bgrad_group_quantize
@@ -1704,7 +1721,35 @@ def _fuser_backward_grouped_tensor(
             )
             grad_output_quantizer.optimize_for_gemm = True
 
-            if (
+            if prequantized_mxfp8_grad:
+                # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse the
+                # rowwise data for the dgrad GEMM and manufacture the columnwise copy
+                # for wgrad. ``scale_bias`` needs the high-precision grad below, so it
+                # takes the dequantized tensor instead of the fused dbias.
+                grouped_dy = grouped_storage_from_grouped_tensor(grad_output)
+                dbias_packed, dy_2d = prepare_prequantized_mxfp8_input_for_gemm(
+                    grouped_dy,
+                    grad_output_quantizer,
+                    num_groups,
+                    split_sizes,
+                    dtype,
+                    with_columnwise=ctx.weight_requires_grad,
+                    with_dbias=has_bias and not self._scale_bias,
+                    tensor_offsets=base_split_offsets * self.out_features,
+                )
+                if has_bias and self._scale_bias:
+                    if dy_2d is None:
+                        # dbias/dscales below need the dequantized grad, which is
+                        # only materialized when the columnwise copy is built.
+                        raise NotImplementedError(
+                            "Pre-quantized MXFP8 grad output with scale_bias requires "
+                            "weight gradients."
+                        )
+                else:
+                    # Nothing else reads the dequantized grad; drop it so the buffer
+                    # is freed instead of living until backward ends.
+                    dy_2d = None
+            elif (
                 has_bias
                 and not self._scale_bias
                 and isinstance(grad_output_quantizer, MXFP8Quantizer)
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 4d9d3972f6..36c63c56ca 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -1589,7 +1589,16 @@ def fuser_backward(
         # Tensor properties
         fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features)
         fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features)
-        grad_output = grad_output.reshape(-1, fc2_weight_shape[0])
+        if isinstance(grad_output, GroupedTensor):
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, out_features) layout; just validate the shape.
+            if grad_output.dim() != 2 or grad_output.size(-1) != fc2_weight_shape[0]:
+                raise ValueError(
+                    "GroupedTensor grad output must have shape (total_tokens, "
+                    f"{fc2_weight_shape[0]}), but got {tuple(grad_output.size())}."
+                )
+        else:
+            grad_output = grad_output.reshape(-1, fc2_weight_shape[0])
         out_shape = list(grad_output.size())
         num_groups = fc1_op.num_groups
         fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
@@ -1657,12 +1666,43 @@ def fuser_backward(
             isinstance(fc2_grad_output_quantizer, NVFP4Quantizer)
             and isinstance(grad_output_quantizer, NVFP4Quantizer)
         )
+        prequantized_mxfp8_grad = isinstance(fc2_grad_output_quantizer, MXFP8Quantizer)
         if (
-            not output_fc2_dbias
+            (not output_fc2_dbias or prequantized_mxfp8_grad)
             and isinstance(grad_output, GroupedTensor)
             and fc2_grad_output_quantizer_matches
         ):
-            grouped_fc2_dy = grad_output
+            if prequantized_mxfp8_grad:
+                # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse
+                # the rowwise data for the dgrad GEMM and manufacture FC2's
+                # columnwise copy for wgrad. ``scale_bias`` needs the
+                # high-precision grad below, so it takes the dequantized tensor
+                # instead of the fused dbias.
+                grouped_fc2_dy = grouped_storage_from_grouped_tensor(grad_output)
+                fc2_dbias_packed, fc2_dy = prepare_prequantized_mxfp8_input_for_gemm(
+                    grouped_fc2_dy,
+                    fc2_grad_output_quantizer,
+                    num_groups,
+                    split_sizes,
+                    dtype,
+                    with_columnwise=fc2_ctx.weight_requires_grad,
+                    with_dbias=output_fc2_dbias and not scale_bias,
+                    tensor_offsets=base_split_offsets * fc2_weight_shape[0],
+                )
+                if scale_bias:
+                    if fc2_dy is None:
+                        # dbias/dscales below need the dequantized grad, which is
+                        # only materialized when the columnwise copy is built.
+                        raise NotImplementedError(
+                            "Pre-quantized MXFP8 grad output with scale_bias requires "
+                            "FC2 weight gradients."
+                        )
+                else:
+                    # Nothing else reads the dequantized grad; drop it so the
+                    # buffer is freed instead of living until backward ends.
+                    fc2_dy = None
+            else:
+                grouped_fc2_dy = grad_output
         else:
             fc2_dy = maybe_dequantize(grad_output, dtype)
             if output_fc2_dbias and not scale_bias:

From 855f480657cdcc03e31700307c32aa71e914e23e Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Sun, 26 Jul 2026 12:33:40 -0700
Subject: [PATCH 04/18] allow scaled_bias + frozen weights in prequant path

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py             | 21 ++++++-
 transformer_engine/pytorch/ops/_common.py     | 60 ++++++++++++++-----
 .../pytorch/ops/basic/grouped_linear.py       | 15 ++---
 .../pytorch/ops/fused/grouped_mlp.py          | 13 +---
 4 files changed, 71 insertions(+), 38 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index 487d080566..be22dc940c 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -1342,6 +1342,7 @@ def _run(x):
             torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0)
 
     @pytest.mark.parametrize("bias", (False, True))
+    @pytest.mark.parametrize("weight_requires_grad", (False, True))
     def test_grouped_mlp_prequantized_mxfp8_grad(
         self,
         *,
@@ -1351,6 +1352,7 @@ def test_grouped_mlp_prequantized_mxfp8_grad(
         dtype: torch.dtype = torch.bfloat16,
         device: torch.device = "cuda",
         bias: bool,
+        weight_requires_grad: bool,
     ) -> None:
         """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor grad output.
 
@@ -1362,6 +1364,11 @@ def test_grouped_mlp_prequantized_mxfp8_grad(
             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")
+        if not weight_requires_grad:
+            # Independent of pre-quantization: the fused forward saves the input
+            # activations whenever anything requires grad, then asserts they carry
+            # columnwise data -- which is only built when the weights need grads.
+            pytest.skip("Fused grouped MLP does not support frozen weights")
         maybe_skip_quantization(
             "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype
         )
@@ -1401,6 +1408,13 @@ def test_grouped_mlp_prequantized_mxfp8_grad(
                 fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2
             )
 
+        # Frozen experts (weights) with a still-training bias/router is the case
+        # where the dequantized grad is needed but is not a byproduct of wgrad.
+        if not weight_requires_grad:
+            for fc in (fc1, fc2):
+                for group_idx in range(group_size):
+                    getattr(fc, f"weight{group_idx}").requires_grad_(False)
+
         def _run(grad):
             x_in = x.detach().clone().requires_grad_()
             fc2_extra = (split_sizes, probs) if bias else (split_sizes,)
@@ -1410,9 +1424,10 @@ def _run(grad):
             grads = [("dx", x_in.grad)]
             for name, fc in (("fc1", fc1), ("fc2", fc2)):
                 for group_idx in range(group_size):
-                    weight = getattr(fc, f"weight{group_idx}")
-                    grads.append((f"{name}_w{group_idx}", weight.grad.detach().clone()))
-                    weight.grad = None
+                    if weight_requires_grad:
+                        weight = getattr(fc, f"weight{group_idx}")
+                        grads.append((f"{name}_w{group_idx}", weight.grad.detach().clone()))
+                        weight.grad = None
                     if bias:
                         bias_param = getattr(fc, f"bias{group_idx}")
                         grads.append((f"{name}_b{group_idx}", bias_param.grad.detach().clone()))
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index 53e01365cf..280c4651df 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -116,26 +116,57 @@ def prepare_prequantized_mxfp8_input_for_gemm(
     *,
     with_columnwise: bool,
     with_dbias: bool = False,
+    with_dequantized: bool = False,
     tensor_offsets: Optional[torch.Tensor] = None,
 ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
     """Make an already-quantized MXFP8 grouped input GEMM-ready (in place).
 
-    For inputs that arrive rowwise-quantized (e.g. FP8 token dispatch). On return
-    the rowwise scales are GEMM-swizzled; columnwise data/scales are populated iff
-    ``with_columnwise``, manufactured by dequantize + columnwise-only requantize
-    since the two directions scale along perpendicular axes.
+    For inputs that arrive rowwise-quantized (e.g. FP8 token dispatch), where the
+    high-precision tensor no longer exists. The rowwise data feeds the GEMM as-is
+    and its scales are swizzled; the columnwise copy cannot be derived from it
+    (the two directions scale along perpendicular axes) so it is manufactured by
+    dequantize + columnwise-only requantize.
 
-    Returns ``(dbias, dequantized)``. ``dbias`` is the per-group bias gradient,
-    produced only when ``with_dbias``: the grouped quantize kernel accumulates it
-    in the columnwise stage it is already running, so it costs no extra pass.
-    ``dequantized`` is the high-precision tensor materialized for the requantize
-    (``None`` if none was needed), for callers that cannot use the fused dbias
-    (e.g. ``scale_bias``, whose dbias/dscales depend on routing probabilities).
-
-    Requires rowwise-only input with unswizzled scales and per-group token counts
-    that are multiples of 128 (caller contract: ``split_sizes`` is on device).
+    The input must be rowwise-only with unswizzled scales, and each group's token
+    count must be a multiple of 128 so per-group scales start on a swizzle-tile
+    boundary.
 
     TODO: optimize and fuse the round-trips requant
+
+    Parameters
+    ----------
+    grouped_x : GroupedTensorStorage
+        Rowwise-quantized input, updated in place.
+    quantizer : MXFP8Quantizer
+        The op's input quantizer. Supplies the FP8 dtype for the columnwise copy.
+    num_groups : int
+        Number of groups.
+    split_sizes : torch.Tensor
+        Per-group row counts, on device.
+    dtype : torch.dtype
+        High-precision dtype to dequantize to.
+    with_columnwise : bool
+        Whether to build the columnwise copy that the wgrad GEMM consumes.
+    with_dbias : bool, default = ``False``
+        Whether to also produce the per-group bias gradient. The quantize kernel
+        accumulates it in the columnwise stage it is already running, so this
+        costs no extra pass but requires ``with_columnwise``.
+    with_dequantized : bool, default = ``False``
+        Whether to dequantize even when no columnwise copy is needed. It must be
+        requested here rather than recovered later, since the rowwise scales are
+        swizzled before returning and dequantization requires unswizzled ones.
+    tensor_offsets : torch.Tensor, optional
+        Per-group element offsets for the columnwise requantize.
+
+    Returns
+    -------
+    torch.Tensor or None
+        Per-group bias gradient, when ``with_dbias``.
+    torch.Tensor or None
+        Dequantized input, when it was materialized (``with_columnwise``, where it
+        is a byproduct, or ``with_dequantized``). Used by callers that cannot use
+        the fused ``dbias``, e.g. ``scale_bias``, whose dbias/dscales depend on
+        the routing probabilities.
     """
     if grouped_x.rowwise_data is None:
         raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise data.")
@@ -169,10 +200,11 @@ def prepare_prequantized_mxfp8_input_for_gemm(
     # wire data and requantize columnwise-only.
     dbias = None
     dequantized = None
-    if with_columnwise:
+    if with_columnwise or with_dequantized:
         dequantized = tex.group_dequantize(grouped_x, TE_DType[dtype]).rowwise_data.view(
             grouped_x.logical_shape
         )
+    if with_columnwise:
         colwise_quantizer = quantizer.copy()
         colwise_quantizer.set_usage(rowwise=False, columnwise=True)
         colwise_quantizer.optimize_for_gemm = True
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index ad23890545..eb113a5fc4 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -1748,19 +1748,12 @@ def _fuser_backward_grouped_tensor(
                     dtype,
                     with_columnwise=ctx.weight_requires_grad,
                     with_dbias=has_bias and not self._scale_bias,
+                    with_dequantized=has_bias and self._scale_bias,
                     tensor_offsets=base_split_offsets * self.out_features,
                 )
-                if has_bias and self._scale_bias:
-                    if dy_2d is None:
-                        # dbias/dscales below need the dequantized grad, which is
-                        # only materialized when the columnwise copy is built.
-                        raise NotImplementedError(
-                            "Pre-quantized MXFP8 grad output with scale_bias requires "
-                            "weight gradients."
-                        )
-                else:
-                    # Nothing else reads the dequantized grad; drop it so the buffer
-                    # is freed instead of living until backward ends.
+                if not (has_bias and self._scale_bias):
+                    # Only scale_bias reads the dequantized grad; drop it so the
+                    # buffer is freed instead of living until backward ends.
                     dy_2d = None
             elif has_bias and not self._scale_bias and fuse_bgrad:
                 grouped_dy, dbias_packed = tex.bgrad_group_quantize(
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 36c63c56ca..a5645ff488 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -1687,18 +1687,11 @@ def fuser_backward(
                     dtype,
                     with_columnwise=fc2_ctx.weight_requires_grad,
                     with_dbias=output_fc2_dbias and not scale_bias,
+                    with_dequantized=scale_bias,
                     tensor_offsets=base_split_offsets * fc2_weight_shape[0],
                 )
-                if scale_bias:
-                    if fc2_dy is None:
-                        # dbias/dscales below need the dequantized grad, which is
-                        # only materialized when the columnwise copy is built.
-                        raise NotImplementedError(
-                            "Pre-quantized MXFP8 grad output with scale_bias requires "
-                            "FC2 weight gradients."
-                        )
-                else:
-                    # Nothing else reads the dequantized grad; drop it so the
+                if not scale_bias:
+                    # Only scale_bias reads the dequantized grad; drop it so the
                     # buffer is freed instead of living until backward ends.
                     fc2_dy = None
             else:

From 8ef7e7e89cc30bf0a3f5922369d8aa0bf3d9dffe Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Sun, 26 Jul 2026 13:19:05 -0700
Subject: [PATCH 05/18] allow bias + frozen weight path in prequantized input

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py         | 71 +++++++++++++++--------
 transformer_engine/pytorch/ops/_common.py | 23 ++++----
 2 files changed, 60 insertions(+), 34 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index be22dc940c..0d97288e98 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -541,7 +541,8 @@ def _run(x):
         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("bias", (False, True))
+    @pytest.mark.parametrize("bias_mode", ("none", "plain", "scaled"))
+    @pytest.mark.parametrize("weight_requires_grad", (False, True))
     def test_grouped_linear_prequantized_mxfp8_grad(
         self,
         *,
@@ -550,19 +551,30 @@ def test_grouped_linear_prequantized_mxfp8_grad(
         split_alignment: int = 128,
         dtype: torch.dtype = torch.bfloat16,
         device: torch.device = "cuda",
-        bias: bool,
+        bias_mode: str,
+        weight_requires_grad: bool,
     ) -> None:
         """Rowwise-only MXFP8 GroupedTensor grad output (FP8 token dispatch, backward).
 
         Mirrors ``test_grouped_linear_prequantized_mxfp8_input`` on the backward
         side: the rowwise data feeds the dgrad GEMM and the columnwise copy is
-        manufactured for wgrad. With ``bias`` the per-group bias gradient rides
-        the columnwise stage of the quantize kernel.
+        manufactured for wgrad. Covers all three bias gradient sources: none,
+        ``plain`` (fused into the columnwise stage of the quantize kernel, or
+        reduced from the dequantized grad when frozen weights leave no columnwise
+        stage), and ``scaled`` (``scale_bias``, whose dbias/dscales need the
+        dequantized grad because they depend on the routing probabilities).
+
+        With frozen weights TE also requires frozen biases, so bias gradients are
+        not observable there; ``dscales`` still is, since the probabilities are an
+        input rather than a parameter.
         """
         if not mxfp8_available:
             pytest.skip(reason_for_no_mxfp8)
         maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype)
 
+        has_bias = bias_mode != "none"
+        use_scale_bias = bias_mode == "scaled"
+
         # Split sizes (including an empty group)
         split_sizes = [split_alignment * i for i in range(group_size)]
         random.shuffle(split_sizes)
@@ -573,6 +585,7 @@ def test_grouped_linear_prequantized_mxfp8_grad(
 
         x = torch.rand((total_tokens, in_features), dtype=dtype, device=device) - 0.5
         dy_hp = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5
+        probs = torch.rand((total_tokens,), dtype=dtype, device=device)
 
         # Wire-format grad output and its exact dequantization (the reference grad).
         dy_wire = self._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes)
@@ -582,35 +595,47 @@ def test_grouped_linear_prequantized_mxfp8_grad(
 
         recipe = make_recipe("mxfp8")
         op = te.ops.GroupedLinear(
-            group_size, in_features, out_features, bias=bias, device=device, dtype=dtype
+            group_size,
+            in_features,
+            out_features,
+            bias=has_bias,
+            device=device,
+            dtype=dtype,
+            scale_bias=use_scale_bias,
         )
+        if not weight_requires_grad:
+            # TE requires bias.requires_grad to match weight.requires_grad.
+            for param in op.parameters():
+                param.requires_grad_(False)
 
         def _run(grad):
             x_in = x.detach().clone().requires_grad_()
+            probs_in = probs.detach().clone().requires_grad_(use_scale_bias)
+            extra_inputs = (split_sizes, probs_in) if use_scale_bias else (split_sizes,)
             with te.autocast(enabled=True, recipe=recipe):
-                y = op(x_in, split_sizes)
+                y = op(x_in, *extra_inputs)
             # Deliver ``grad`` as the op's grad output, as FP8 dispatch would.
             _InjectGrad.apply(y, grad).backward(torch.ones_like(y))
-            wgrads, bgrads = [], []
-            for group_idx in range(group_size):
-                weight = getattr(op, f"weight{group_idx}")
-                wgrads.append(weight.grad.detach().clone())
-                weight.grad = None
-                if bias:
-                    bias_param = getattr(op, f"bias{group_idx}")
-                    bgrads.append(bias_param.grad.detach().clone())
-                    bias_param.grad = None
-            return x_in.grad, wgrads, bgrads
+            grads = [("dx", x_in.grad)]
+            if use_scale_bias:
+                grads.append(("dprobs", probs_in.grad))
+            if weight_requires_grad:
+                for group_idx in range(group_size):
+                    weight = getattr(op, f"weight{group_idx}")
+                    grads.append((f"w{group_idx}", weight.grad.detach().clone()))
+                    weight.grad = None
+                    if has_bias:
+                        bias_param = getattr(op, f"bias{group_idx}")
+                        grads.append((f"b{group_idx}", bias_param.grad.detach().clone()))
+                        bias_param.grad = None
+            return grads
 
-        dx_ref, wgrads_ref, bgrads_ref = _run(dy_ref)
-        dx_test, wgrads_test, bgrads_test = _run(dy_wire)
+        grads_ref = _run(dy_ref)
+        grads_test = _run(dy_wire)
 
         # Bit-exact match expected (identical quantized grads and kernels).
-        torch.testing.assert_close(dx_test, dx_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)
-        for bgrad_test, bgrad_ref in zip(bgrads_test, bgrads_ref):
-            torch.testing.assert_close(bgrad_test, bgrad_ref, rtol=0, atol=0)
+        for (_, grad_test), (_, grad_ref) in zip(grads_test, grads_ref):
+            torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0)
 
     @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16))
     @pytest.mark.parametrize(
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index 280c4651df..7d2589bee3 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -20,6 +20,7 @@
 from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
 from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage
 from ..quantized_tensor import QuantizedTensorStorage
+from ..triton.grouped_dbias_dscales import compute_grouped_dbias
 from ..utils import canonicalize_dtype
 
 
@@ -148,9 +149,10 @@ def prepare_prequantized_mxfp8_input_for_gemm(
     with_columnwise : bool
         Whether to build the columnwise copy that the wgrad GEMM consumes.
     with_dbias : bool, default = ``False``
-        Whether to also produce the per-group bias gradient. The quantize kernel
-        accumulates it in the columnwise stage it is already running, so this
-        costs no extra pass but requires ``with_columnwise``.
+        Whether to also produce the per-group bias gradient. When a columnwise
+        copy is built the quantize kernel accumulates it in the stage it is
+        already running, so it costs no extra pass; otherwise it is reduced from
+        the dequantized tensor.
     with_dequantized : bool, default = ``False``
         Whether to dequantize even when no columnwise copy is needed. It must be
         requested here rather than recovered later, since the rowwise scales are
@@ -189,18 +191,11 @@ def prepare_prequantized_mxfp8_input_for_gemm(
             f"but the op's input quantizer expects {quantizer.dtype}."
         )
 
-    if with_dbias and not with_columnwise:
-        # The fused dbias is accumulated by the quantize kernel's columnwise stage.
-        raise NotImplementedError(
-            "Pre-quantized MXFP8 grouped input cannot produce a fused dbias without "
-            "a columnwise copy."
-        )
-
     # Manufacture columnwise data for the wgrad GEMM: dequantize the rowwise
     # wire data and requantize columnwise-only.
     dbias = None
     dequantized = None
-    if with_columnwise or with_dequantized:
+    if with_columnwise or with_dbias or with_dequantized:
         dequantized = tex.group_dequantize(grouped_x, TE_DType[dtype]).rowwise_data.view(
             grouped_x.logical_shape
         )
@@ -227,6 +222,12 @@ def prepare_prequantized_mxfp8_input_for_gemm(
             )
         grouped_x.columnwise_data = colwise_x.columnwise_data
         grouped_x.columnwise_scale_inv = colwise_x.columnwise_scale_inv
+    elif with_dbias:
+        # No columnwise stage to accumulate into (e.g. frozen weights need no
+        # wgrad), so reduce the dequantized grad directly.
+        dbias = compute_grouped_dbias(
+            dequantized, tex.splits_to_offsets(split_sizes, 1), num_groups
+        )
 
     # Convert rowwise scales to the GEMM-swizzled layout. The grouped GEMM
     # reads activation scales as one (total_tokens, cols) matrix, so the

From 6015c91b92bee2575cf4b870f85b03d991b8f539 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Wed, 29 Jul 2026 20:32:36 -0700
Subject: [PATCH 06/18] use .copy to create group tensor

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py             |  2 +-
 transformer_engine/pytorch/ops/_common.py     | 37 -------------------
 .../pytorch/ops/basic/grouped_linear.py       |  5 +--
 .../pytorch/ops/fused/grouped_mlp.py          |  5 +--
 4 files changed, 5 insertions(+), 44 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index 0d97288e98..4bfb0ebb10 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -462,7 +462,7 @@ def _make_rowwise_mxfp8_wire_input(
     ) -> "GroupedTensor":
         """Rowwise-only MXFP8 GroupedTensor with compact scales (FP8 dispatch wire format)."""
         wire_quantizer = MXFP8Quantizer(
-            fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False
+            fp8_dtype=TE_DType[torch.float8_e4m3fn], rowwise=True, columnwise=False
         )
         x_wire = tex.group_quantize(
             x_hp, wire_quantizer, group_size, split_sizes.to(dtype=torch.int64)
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index 7d2589bee3..24d2928ab7 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -16,7 +16,6 @@
 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
@@ -72,42 +71,6 @@ 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_input_for_gemm(
     grouped_x: GroupedTensorStorage,
     quantizer: MXFP8Quantizer,
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index eb113a5fc4..5b0ebc857a 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -47,7 +47,6 @@
     get_accumulate_flag_in_param,
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
-    grouped_storage_from_grouped_tensor,
     is_quantized_tensor,
     maybe_dequantize,
     prepare_prequantized_mxfp8_input_for_gemm,
@@ -1336,7 +1335,7 @@ def _fuser_forward_grouped_tensor(
             # rowwise data to the forward GEMM as-is, manufacture the
             # columnwise copy needed by the wgrad GEMM, and swizzle the
             # rowwise scales for the GEMM.
-            grouped_x = grouped_storage_from_grouped_tensor(input_)
+            grouped_x = input_.copy()
             prepare_prequantized_mxfp8_input_for_gemm(
                 grouped_x,
                 input_quantizers[0],
@@ -1739,7 +1738,7 @@ def _fuser_backward_grouped_tensor(
                 # rowwise data for the dgrad GEMM and manufacture the columnwise copy
                 # for wgrad. ``scale_bias`` needs the high-precision grad below, so it
                 # takes the dequantized tensor instead of the fused dbias.
-                grouped_dy = grouped_storage_from_grouped_tensor(grad_output)
+                grouped_dy = grad_output.copy()
                 dbias_packed, dy_2d = prepare_prequantized_mxfp8_input_for_gemm(
                     grouped_dy,
                     grad_output_quantizer,
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index a5645ff488..4cad8ff65b 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -52,7 +52,6 @@
     get_accumulate_flag_in_param,
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
-    grouped_storage_from_grouped_tensor,
     is_quantized_tensor,
     maybe_dequantize,
     prepare_prequantized_mxfp8_input_for_gemm,
@@ -1091,7 +1090,7 @@ def fuser_forward(
             or isinstance(fc1_input_quantizer, NVFP4Quantizer)
             and isinstance(input_quantizer, NVFP4Quantizer)
         ):
-            grouped_fc1_x = grouped_storage_from_grouped_tensor(input_)
+            grouped_fc1_x = input_.copy()
             if (
                 isinstance(fc1_input_quantizer, MXFP8Quantizer)
                 and not grouped_fc1_x._with_gemm_swizzled_scales
@@ -1678,7 +1677,7 @@ def fuser_backward(
                 # columnwise copy for wgrad. ``scale_bias`` needs the
                 # high-precision grad below, so it takes the dequantized tensor
                 # instead of the fused dbias.
-                grouped_fc2_dy = grouped_storage_from_grouped_tensor(grad_output)
+                grouped_fc2_dy = grad_output.copy()
                 fc2_dbias_packed, fc2_dy = prepare_prequantized_mxfp8_input_for_gemm(
                     grouped_fc2_dy,
                     fc2_grad_output_quantizer,

From 12f1dbfe027b638c22c23d3c906c5cc9e5e41612 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Thu, 30 Jul 2026 16:41:10 -0700
Subject: [PATCH 07/18] move the implementation to c++ layer

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 transformer_engine/pytorch/csrc/extensions.h  |   5 +
 .../pytorch/csrc/extensions/cast.cpp          |  55 ++++++
 .../pytorch/csrc/extensions/pybind.cpp        |   7 +
 .../pytorch/csrc/extensions/swizzle.cpp       |  17 ++
 transformer_engine/pytorch/ops/_common.py     | 168 +-----------------
 .../pytorch/ops/basic/grouped_linear.py       |  69 ++++---
 .../pytorch/ops/fused/grouped_mlp.py          |  81 ++++++---
 7 files changed, 184 insertions(+), 218 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 95f02c64f0..30df6d5975 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -367,6 +367,11 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 std::optional<at::Tensor> last_dims,
                                 std::optional<at::Tensor> tensor_offsets);
 
+py::object group_requantize_columnwise_and_swizzle_rowwise_(
+    py::handle grouped_x, py::handle columnwise_quantizer, const size_t num_tensors,
+    std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
+    bool return_dequantized);
+
 std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
                                               std::vector<py::handle> quantizer_list);
 
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index 8b1cd384aa..92b7ff22e7 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -663,6 +663,61 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o
   return py::reinterpret_borrow<py::object>(out_py);
 }
 
+py::object group_requantize_columnwise_and_swizzle_rowwise_(
+    py::handle grouped_x, py::handle columnwise_quantizer, const size_t num_tensors,
+    std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
+    bool return_dequantized) {
+  init_extension();
+
+  NVTE_CHECK(!grouped_x.attr("rowwise_data").is_none(),
+             "Pre-quantized MXFP8 grouped input is missing rowwise data.");
+  NVTE_CHECK(!grouped_x.attr("scale_inv").is_none(),
+             "Pre-quantized MXFP8 grouped input is missing rowwise scales.");
+  NVTE_CHECK(!grouped_x.attr("_with_gemm_swizzled_scales").cast<bool>(),
+             "Pre-quantized MXFP8 grouped input must have unswizzled scales.");
+  NVTE_CHECK(grouped_x.attr("columnwise_data").is_none(),
+             "Pre-quantized MXFP8 grouped input must be rowwise-only.");
+  if (!grouped_x.attr("quantizer").is_none()) {
+    // The GEMM consumes the input's rowwise data verbatim while the wgrad GEMM consumes the
+    // columnwise copy built here, so a dtype mismatch would make the two directions disagree.
+    NVTE_CHECK(grouped_x.attr("quantizer").attr("dtype").cast<DType>() ==
+                   columnwise_quantizer.attr("dtype").cast<DType>(),
+               "Pre-quantized MXFP8 grouped input and the columnwise quantizer disagree on the "
+               "FP8 dtype.");
+  }
+
+  const auto logical_shape = grouped_x.attr("logical_shape").cast<py::tuple>();
+  const auto total_tokens = logical_shape[0].cast<size_t>();
+  const auto hidden_dim = logical_shape[1].cast<size_t>();
+  // Each group's token count must be a multiple of 128 too, so that every group's scales start
+  // on a swizzle-tile boundary. Those counts live on the device (host reads would break CUDA
+  // graph capture), so that half is the caller's contract rather than an assertion.
+  NVTE_CHECK(total_tokens % 128 == 0 && hidden_dim % 128 == 0,
+             "Pre-quantized MXFP8 grouped input requires dims that are multiples of 128, but got (",
+             total_tokens, ", ", hidden_dim, ").");
+
+  // Dequantize first: it reads the rowwise scales, which the swizzle below overwrites.
+  auto dequantized_grouped = group_dequantize(grouped_x, otype);
+  auto dequantized = dequantized_grouped.attr("rowwise_data").cast<at::Tensor>().view(
+      {static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
+
+  // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
+  // data because the two directions scale along perpendicular axes. The caller hands us a
+  // columnwise-only, optimize_for_gemm quantizer, so the kernel emits swizzled columnwise
+  // scales directly and only the rowwise scales need the explicit swizzle below.
+  auto columnwise = group_quantize(dequantized, columnwise_quantizer, num_tensors, first_dims,
+                                   std::nullopt, tensor_offsets, std::nullopt);
+  grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
+  grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
+
+  grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false);
+
+  if (return_dequantized) {
+    return py::cast(dequantized);
+  }
+  return py::none();
+}
+
 namespace {
 
 void multi_tensor_quantize_impl(const std::vector<TensorWrapper> &input_list,
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index c6bf7eb516..3a5049d7b8 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -216,6 +216,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
         py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none());
+  m.def("group_requantize_columnwise_and_swizzle_rowwise_",
+        transformer_engine::pytorch::group_requantize_columnwise_and_swizzle_rowwise_,
+        "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle "
+        "its rowwise scales for GEMM, in place",
+        py::arg("grouped_x"), py::arg("columnwise_quantizer"), py::arg("num_tensors"),
+        py::arg("first_dims"), py::arg("otype"), py::arg("tensor_offsets") = py::none(),
+        py::arg("return_dequantized") = false);
   m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize,
         "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer"));
   m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)",
diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
index c90a7d6d0d..866d858ee3 100644
--- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
@@ -389,6 +389,23 @@ std::optional<SwizzledGroupedScales> maybe_swizzle_grouped_tensor(GroupedTensorW
         tensor_offsets.data_ptr, static_cast<DType>(tensor_offsets.dtype), tensor_offsets.shape);
   }
 
+  // Varying per-tensor dimensions. Leaving these unset declares the grouped tensor uniform,
+  // which selects the uniform-shape swizzle kernel.
+  const auto first_dims = input.get_first_dims();
+  if (first_dims.data_ptr != nullptr) {
+    swizzle_input.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
+                                 first_dims.shape);
+    swizzle_output.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
+                                  first_dims.shape);
+  }
+  const auto last_dims = input.get_last_dims();
+  if (last_dims.data_ptr != nullptr) {
+    swizzle_input.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
+                                last_dims.shape);
+    swizzle_output.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
+                                 last_dims.shape);
+  }
+
   // Per-tensor logical dimensions (uniform-shape grouped tensor).
   const size_t num_tensors = input.num_tensors();
   const auto logical_shape_nvte = input.logical_shape();
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index 24d2928ab7..c2b22e32bf 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -10,16 +10,12 @@
 
 import torch
 
-import transformer_engine_torch as tex
 from transformer_engine_torch import FP8TensorMeta
-from ..constants import MXFP8_BLOCK_SCALING_SIZE, TE_DType
 from ..torch_version import torch_version
 from ..quantization import FP8GlobalStateManager
 from ..tensor.float8_tensor import Float8Tensor
-from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
-from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage
+from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..quantized_tensor import QuantizedTensorStorage
-from ..triton.grouped_dbias_dscales import compute_grouped_dbias
 from ..utils import canonicalize_dtype
 
 
@@ -71,161 +67,13 @@ def maybe_dequantize(
     return tensor
 
 
-def prepare_prequantized_mxfp8_input_for_gemm(
-    grouped_x: GroupedTensorStorage,
-    quantizer: MXFP8Quantizer,
-    num_groups: int,
-    split_sizes: torch.Tensor,
-    dtype: torch.dtype,
-    *,
-    with_columnwise: bool,
-    with_dbias: bool = False,
-    with_dequantized: bool = False,
-    tensor_offsets: Optional[torch.Tensor] = None,
-) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
-    """Make an already-quantized MXFP8 grouped input GEMM-ready (in place).
-
-    For inputs that arrive rowwise-quantized (e.g. FP8 token dispatch), where the
-    high-precision tensor no longer exists. The rowwise data feeds the GEMM as-is
-    and its scales are swizzled; the columnwise copy cannot be derived from it
-    (the two directions scale along perpendicular axes) so it is manufactured by
-    dequantize + columnwise-only requantize.
-
-    The input must be rowwise-only with unswizzled scales, and each group's token
-    count must be a multiple of 128 so per-group scales start on a swizzle-tile
-    boundary.
-
-    TODO: optimize and fuse the round-trips requant
-
-    Parameters
-    ----------
-    grouped_x : GroupedTensorStorage
-        Rowwise-quantized input, updated in place.
-    quantizer : MXFP8Quantizer
-        The op's input quantizer. Supplies the FP8 dtype for the columnwise copy.
-    num_groups : int
-        Number of groups.
-    split_sizes : torch.Tensor
-        Per-group row counts, on device.
-    dtype : torch.dtype
-        High-precision dtype to dequantize to.
-    with_columnwise : bool
-        Whether to build the columnwise copy that the wgrad GEMM consumes.
-    with_dbias : bool, default = ``False``
-        Whether to also produce the per-group bias gradient. When a columnwise
-        copy is built the quantize kernel accumulates it in the stage it is
-        already running, so it costs no extra pass; otherwise it is reduced from
-        the dequantized tensor.
-    with_dequantized : bool, default = ``False``
-        Whether to dequantize even when no columnwise copy is needed. It must be
-        requested here rather than recovered later, since the rowwise scales are
-        swizzled before returning and dequantization requires unswizzled ones.
-    tensor_offsets : torch.Tensor, optional
-        Per-group element offsets for the columnwise requantize.
-
-    Returns
-    -------
-    torch.Tensor or None
-        Per-group bias gradient, when ``with_dbias``.
-    torch.Tensor or None
-        Dequantized input, when it was materialized (``with_columnwise``, where it
-        is a byproduct, or ``with_dequantized``). Used by callers that cannot use
-        the fused ``dbias``, e.g. ``scale_bias``, whose dbias/dscales depend on
-        the routing probabilities.
-    """
-    if grouped_x.rowwise_data is None:
-        raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise data.")
-    if grouped_x.scale_inv is None:
-        raise ValueError("Pre-quantized MXFP8 grouped input is missing rowwise scales.")
-    if grouped_x._with_gemm_swizzled_scales:
-        raise NotImplementedError("Pre-quantized MXFP8 grouped input must have unswizzled scales.")
-    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 unswizzled 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.
-    dbias = None
-    dequantized = None
-    if with_columnwise or with_dbias or with_dequantized:
-        dequantized = tex.group_dequantize(grouped_x, TE_DType[dtype]).rowwise_data.view(
-            grouped_x.logical_shape
-        )
-    if with_columnwise:
-        colwise_quantizer = quantizer.copy()
-        colwise_quantizer.set_usage(rowwise=False, columnwise=True)
-        colwise_quantizer.optimize_for_gemm = True
-        colwise_quantizer.internal = True
-        if with_dbias:
-            colwise_x, dbias = tex.bgrad_group_quantize(
-                dequantized,
-                colwise_quantizer,
-                num_groups,
-                split_sizes,
-                tensor_offsets=tensor_offsets,
-            )
-        else:
-            colwise_x = tex.group_quantize(
-                dequantized,
-                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
-    elif with_dbias:
-        # No columnwise stage to accumulate into (e.g. frozen weights need no
-        # wgrad), so reduce the dequantized grad directly.
-        dbias = compute_grouped_dbias(
-            dequantized, tex.splits_to_offsets(split_sizes, 1), num_groups
-        )
-
-    # 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 unswizzled scales are left untouched.
-    # 128-alignment (see docstring) means the scale array needs no padding.
-    total_tokens, cols = grouped_x.logical_shape
-    if total_tokens % 128 != 0 or cols % 128 != 0:
-        raise ValueError(
-            "Pre-quantized MXFP8 grouped input requires dims that are multiples of 128, "
-            f"but got ({total_tokens}, {cols})."
-        )
-    scale_shape = (total_tokens, cols // MXFP8_BLOCK_SCALING_SIZE)
-    if grouped_x.scale_inv.numel() != math.prod(scale_shape):
-        raise ValueError(
-            f"Pre-quantized MXFP8 grouped input has {grouped_x.scale_inv.numel()} rowwise "
-            f"scales, but expected {math.prod(scale_shape)} for shape {scale_shape}."
-        )
-    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
-
-    return dbias, dequantized
+def make_columnwise_gemm_quantizer(quantizer: MXFP8Quantizer) -> MXFP8Quantizer:
+    """Copy of ``quantizer`` configured to emit only GEMM-swizzled columnwise data."""
+    columnwise_quantizer = quantizer.copy()
+    columnwise_quantizer.set_usage(rowwise=False, columnwise=True)
+    columnwise_quantizer.optimize_for_gemm = True
+    columnwise_quantizer.internal = True
+    return columnwise_quantizer
 
 
 def maybe_autocast_dtype(
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index afe36c86dd..499ab7d4cc 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -14,7 +14,7 @@
 import torch
 
 import transformer_engine_torch as tex
-from ...constants import DType
+from ...constants import DType, TE_DType
 from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor
 from ...distributed import CudaRNGStatesTracker
 from ...module._common import WeightGradStore
@@ -48,8 +48,8 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
+    make_columnwise_gemm_quantizer,
     maybe_dequantize,
-    prepare_prequantized_mxfp8_input_for_gemm,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
 )
@@ -1357,15 +1357,19 @@ def _fuser_forward_grouped_tensor(
             # columnwise copy needed by the wgrad GEMM, and swizzle the
             # rowwise scales for the GEMM.
             grouped_x = input_.copy()
-            prepare_prequantized_mxfp8_input_for_gemm(
-                grouped_x,
-                input_quantizers[0],
-                num_groups,
-                split_sizes,
-                dtype,
-                with_columnwise=weight_requires_grad,
-                tensor_offsets=base_split_offsets * self.in_features,
-            )
+            if weight_requires_grad:
+                tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                    grouped_x,
+                    make_columnwise_gemm_quantizer(input_quantizers[0]),
+                    num_groups,
+                    split_sizes,
+                    TE_DType[dtype],
+                    tensor_offsets=base_split_offsets * self.in_features,
+                )
+            else:
+                # No wgrad, so no columnwise copy is needed. The forward GEMM
+                # still requires swizzled rowwise scales.
+                tex.grouped_swizzle_for_gemm(grouped_x, rowwise=True, columnwise=False)
         elif with_quantized_compute:
             input_quantizer = input_quantizers[0]
             input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
@@ -1757,24 +1761,30 @@ def _fuser_backward_grouped_tensor(
             if prequantized_mxfp8_grad:
                 # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse the
                 # rowwise data for the dgrad GEMM and manufacture the columnwise copy
-                # for wgrad. ``scale_bias`` needs the high-precision grad below, so it
-                # takes the dequantized tensor instead of the fused dbias.
+                # for wgrad. Bias grads are reduced from the dequantized grad below,
+                # which is only kept when there is a bias.
                 grouped_dy = grad_output.copy()
-                dbias_packed, dy_2d = prepare_prequantized_mxfp8_input_for_gemm(
-                    grouped_dy,
-                    grad_output_quantizer,
-                    num_groups,
-                    split_sizes,
-                    dtype,
-                    with_columnwise=ctx.weight_requires_grad,
-                    with_dbias=has_bias and not self._scale_bias,
-                    with_dequantized=has_bias and self._scale_bias,
-                    tensor_offsets=base_split_offsets * self.out_features,
-                )
-                if not (has_bias and self._scale_bias):
-                    # Only scale_bias reads the dequantized grad; drop it so the
-                    # buffer is freed instead of living until backward ends.
-                    dy_2d = None
+                if ctx.weight_requires_grad:
+                    dy_2d = tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                        grouped_dy,
+                        make_columnwise_gemm_quantizer(grad_output_quantizer),
+                        num_groups,
+                        split_sizes,
+                        TE_DType[dtype],
+                        tensor_offsets=base_split_offsets * self.out_features,
+                        return_dequantized=has_bias,
+                    )
+                else:
+                    # No wgrad, so no columnwise copy is needed. Dequantize before
+                    # swizzling: dequantization reads the unswizzled rowwise scales.
+                    dy_2d = (
+                        tex.group_dequantize(grouped_dy, TE_DType[dtype]).rowwise_data.view(
+                            total_tokens, self.out_features
+                        )
+                        if has_bias
+                        else None
+                    )
+                    tex.grouped_swizzle_for_gemm(grouped_dy, rowwise=True, columnwise=False)
             elif has_bias and not self._scale_bias and fuse_bgrad:
                 grouped_dy, dbias_packed = tex.bgrad_group_quantize(
                     dy_2d, grad_output_quantizer, num_groups, split_sizes
@@ -1810,7 +1820,8 @@ def _fuser_backward_grouped_tensor(
                     offsets=base_split_offsets,
                 )
             elif dbias_packed is None:
-                # BF16/FP16 path
+                # BF16/FP16 and pre-quantized MXFP8 paths, neither of which fuses dbias
+                # into a quantize kernel.
                 dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups)
             if self.single_grouped_bias:
                 final_bias_grads = [dbias_packed.to(dtype=dtype)]
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index f5281b9561..e6689976ac 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -16,7 +16,7 @@
 from packaging.version import Version as PkgVersion
 
 import transformer_engine_torch as tex
-from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE
+from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload
 from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor
 from ...distributed_weight import (
@@ -31,7 +31,7 @@
 from ...tensor.grouped_tensor import GroupedTensor
 from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
 from ...tensor.storage.grouped_tensor_storage import GroupedTensorStorage
-from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales
+from ...triton.grouped_dbias_dscales import compute_grouped_dbias, compute_grouped_dbias_dscales
 from ...utils import (
     ceil_div,
     clear_tensor_data,
@@ -53,8 +53,8 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
+    make_columnwise_gemm_quantizer,
     maybe_dequantize,
-    prepare_prequantized_mxfp8_input_for_gemm,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
 )
@@ -1224,15 +1224,21 @@ def fuser_forward(
                 # Rowwise-only MXFP8 input (e.g. FP8 token dispatch):
                 # manufacture the columnwise copy needed by the wgrad GEMM
                 # and swizzle the rowwise scales for the forward GEMM.
-                prepare_prequantized_mxfp8_input_for_gemm(
-                    grouped_fc1_x,
-                    fc1_input_quantizer,
-                    num_groups,
-                    split_sizes,
-                    dtype,
-                    with_columnwise=weight_requires_grad,
-                    tensor_offsets=fc1_x_tensor_offsets,
-                )
+                if weight_requires_grad:
+                    tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                        grouped_fc1_x,
+                        make_columnwise_gemm_quantizer(fc1_input_quantizer),
+                        num_groups,
+                        split_sizes,
+                        TE_DType[dtype],
+                        tensor_offsets=fc1_x_tensor_offsets,
+                    )
+                else:
+                    # No wgrad, so no columnwise copy is needed. The forward GEMM
+                    # still requires swizzled rowwise scales.
+                    tex.grouped_swizzle_for_gemm(
+                        grouped_fc1_x, rowwise=True, columnwise=False
+                    )
         else:
             fc1_x = maybe_dequantize(input_, dtype)
             grouped_fc1_x = _group_quantize_for_grouped_mlp(
@@ -1864,24 +1870,41 @@ def fuser_backward(
             if prequantized_mxfp8_grad:
                 # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse
                 # the rowwise data for the dgrad GEMM and manufacture FC2's
-                # columnwise copy for wgrad. ``scale_bias`` needs the
-                # high-precision grad below, so it takes the dequantized tensor
-                # instead of the fused dbias.
+                # columnwise copy for wgrad. Bias grads are reduced from the
+                # dequantized grad, which is only materialized when one is needed.
                 grouped_fc2_dy = grad_output.copy()
-                fc2_dbias_packed, fc2_dy = prepare_prequantized_mxfp8_input_for_gemm(
-                    grouped_fc2_dy,
-                    fc2_grad_output_quantizer,
-                    num_groups,
-                    split_sizes,
-                    dtype,
-                    with_columnwise=fc2_ctx.weight_requires_grad,
-                    with_dbias=output_fc2_dbias and not scale_bias,
-                    with_dequantized=scale_bias,
-                    tensor_offsets=base_split_offsets * fc2_weight_shape[0],
-                )
-                if not scale_bias:
-                    # Only scale_bias reads the dequantized grad; drop it so the
-                    # buffer is freed instead of living until backward ends.
+                need_dequantized = output_fc2_dbias or scale_bias
+                if fc2_ctx.weight_requires_grad:
+                    fc2_dy = tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                        grouped_fc2_dy,
+                        make_columnwise_gemm_quantizer(fc2_grad_output_quantizer),
+                        num_groups,
+                        split_sizes,
+                        TE_DType[dtype],
+                        tensor_offsets=base_split_offsets * fc2_weight_shape[0],
+                        return_dequantized=need_dequantized,
+                    )
+                else:
+                    # No wgrad, so no columnwise copy is needed. Dequantize before
+                    # swizzling: dequantization reads the unswizzled rowwise scales.
+                    fc2_dy = (
+                        tex.group_dequantize(grouped_fc2_dy, TE_DType[dtype]).rowwise_data.view(
+                            grouped_fc2_dy.logical_shape
+                        )
+                        if need_dequantized
+                        else None
+                    )
+                    tex.grouped_swizzle_for_gemm(
+                        grouped_fc2_dy, rowwise=True, columnwise=False
+                    )
+                if output_fc2_dbias and not scale_bias:
+                    # This path has no quantize kernel to fuse dbias into, and the
+                    # consumer below has no fallback, so reduce it here.
+                    fc2_dbias_packed = compute_grouped_dbias(
+                        fc2_dy, base_split_offsets, num_groups
+                    )
+                    # scale_bias is the only later consumer of the dequantized grad;
+                    # drop it so the buffer is freed rather than held until backward ends.
                     fc2_dy = None
             else:
                 grouped_fc2_dy = grad_output

From e3990ddf9a3ae92772c83d98fbfa1c04243fd65a Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 30 Jul 2026 23:42:23 +0000
Subject: [PATCH 08/18] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/csrc/extensions/cast.cpp |  6 ++++--
 transformer_engine/pytorch/ops/fused/grouped_mlp.py | 12 +++---------
 2 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index 92b7ff22e7..a256b65e8c 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -698,8 +698,10 @@ py::object group_requantize_columnwise_and_swizzle_rowwise_(
 
   // Dequantize first: it reads the rowwise scales, which the swizzle below overwrites.
   auto dequantized_grouped = group_dequantize(grouped_x, otype);
-  auto dequantized = dequantized_grouped.attr("rowwise_data").cast<at::Tensor>().view(
-      {static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
+  auto dequantized =
+      dequantized_grouped.attr("rowwise_data")
+          .cast<at::Tensor>()
+          .view({static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
 
   // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
   // data because the two directions scale along perpendicular axes. The caller hands us a
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index e6689976ac..cf0c459605 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -1236,9 +1236,7 @@ def fuser_forward(
                 else:
                     # No wgrad, so no columnwise copy is needed. The forward GEMM
                     # still requires swizzled rowwise scales.
-                    tex.grouped_swizzle_for_gemm(
-                        grouped_fc1_x, rowwise=True, columnwise=False
-                    )
+                    tex.grouped_swizzle_for_gemm(grouped_fc1_x, rowwise=True, columnwise=False)
         else:
             fc1_x = maybe_dequantize(input_, dtype)
             grouped_fc1_x = _group_quantize_for_grouped_mlp(
@@ -1894,15 +1892,11 @@ def fuser_backward(
                         if need_dequantized
                         else None
                     )
-                    tex.grouped_swizzle_for_gemm(
-                        grouped_fc2_dy, rowwise=True, columnwise=False
-                    )
+                    tex.grouped_swizzle_for_gemm(grouped_fc2_dy, rowwise=True, columnwise=False)
                 if output_fc2_dbias and not scale_bias:
                     # This path has no quantize kernel to fuse dbias into, and the
                     # consumer below has no fallback, so reduce it here.
-                    fc2_dbias_packed = compute_grouped_dbias(
-                        fc2_dy, base_split_offsets, num_groups
-                    )
+                    fc2_dbias_packed = compute_grouped_dbias(fc2_dy, base_split_offsets, num_groups)
                     # scale_bias is the only later consumer of the dequantized grad;
                     # drop it so the buffer is freed rather than held until backward ends.
                     fc2_dy = None

From c1e4663c3068e025c493e8f00886f102037edde4 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Thu, 30 Jul 2026 17:46:32 -0700
Subject: [PATCH 09/18] bug fix for operation ordering

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 transformer_engine/pytorch/csrc/extensions/cast.cpp | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index a256b65e8c..db8099e2e2 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -696,24 +696,27 @@ py::object group_requantize_columnwise_and_swizzle_rowwise_(
              "Pre-quantized MXFP8 grouped input requires dims that are multiples of 128, but got (",
              total_tokens, ", ", hidden_dim, ").");
 
-  // Dequantize first: it reads the rowwise scales, which the swizzle below overwrites.
+  // Dequantize first: it reads the rowwise scales, which the swizzle below replaces.
   auto dequantized_grouped = group_dequantize(grouped_x, otype);
   auto dequantized =
       dequantized_grouped.attr("rowwise_data")
           .cast<at::Tensor>()
           .view({static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
 
+  // Swizzle the rowwise scales before attaching any columnwise data: a rowwise-only swizzle
+  // resets columnwise_scale_inv to None, which would strand the columnwise data below with a
+  // null scale pointer.
+  grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false);
+
   // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
   // data because the two directions scale along perpendicular axes. The caller hands us a
   // columnwise-only, optimize_for_gemm quantizer, so the kernel emits swizzled columnwise
-  // scales directly and only the rowwise scales need the explicit swizzle below.
+  // scales directly.
   auto columnwise = group_quantize(dequantized, columnwise_quantizer, num_tensors, first_dims,
                                    std::nullopt, tensor_offsets, std::nullopt);
   grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
   grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
 
-  grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false);
-
   if (return_dequantized) {
     return py::cast(dequantized);
   }

From a349d0d30f9542e152fa823a44e213dababa8fc2 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Thu, 30 Jul 2026 22:11:29 -0700
Subject: [PATCH 10/18] add comprehensive tests

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 .../test_mxfp8_group_quantize_graph_safe.py   | 237 ++++++++++++++++++
 .../pytorch/csrc/extensions/cast.cpp          |   4 +
 2 files changed, 241 insertions(+)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index d07953ce37..6502a2b314 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -469,3 +469,240 @@ def test_grouped_tensor_mxfp8_with_paged_stashing(
         valid_M=valid_M,
         optimize_for_gemm=optimize_for_gemm,
     )
+
+
+# ---------------------------------------------------------------------------------------------
+# Pre-quantized MXFP8 input (FP8 token dispatch)
+#
+# tex.group_requantize_columnwise_and_swizzle_rowwise_ takes a grouped tensor that arrives
+# ALREADY rowwise-quantized (its high-precision form no longer exists), and makes it GEMM-ready
+# in both directions: the rowwise data passes through verbatim with its scales swizzled, and the
+# columnwise copy is rebuilt via dequantize + columnwise-only requantize.
+#
+# These mirror the edge-case matrices of test_grouped_tensor_mxfp8_versus_reference and
+# test_grouped_tensor_mxfp8_with_paged_stashing so the same shapes, zero-token placements and
+# uneven splits exercise this path.
+# ---------------------------------------------------------------------------------------------
+
+
+def make_prequantized_wire_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor):
+    """Rowwise-only, unswizzled grouped tensor, as FP8 dispatch delivers it."""
+    wire_quantizer = MXFP8Quantizer(
+        fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False
+    )
+    # Must stay unswizzled: the requantize path asserts on it and dequantize needs compact scales.
+    wire_quantizer.optimize_for_gemm = False
+    wire = fused_grouped_quantize(x, split_section_tensor, wire_quantizer)
+    assert wire.columnwise_data is None
+    assert not wire._with_gemm_swizzled_scales
+    return wire
+
+
+def check_prequantized_requantize_versus_reference(
+    x_dtype: torch.dtype,
+    M: int,
+    N: int,
+    split_sections: list[int],
+) -> None:
+    """Run the pre-quantized requantize path and check both directions against a reference.
+
+    The reference is derived from dequantize(wire), not from the original high-precision x: that
+    is the only data a consumer can see after dispatch, and MXFP8 rowwise requantization is
+    idempotent, so it is an exact reference rather than an approximate one.
+    """
+    device = "cuda"
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+
+    # The buffer is always M rows. Paged stashing is just the case where the groups cover fewer
+    # than M of them (valid_M < M) and the tail holds garbage the kernels must leave alone; the
+    # non-paged case is the same code path with sum(split_sections) == M.
+    x = torch.randn((M, N), dtype=x_dtype, device=device)
+    split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device=device)
+    num_groups = len(split_sections)
+    # Rows the groups actually cover. Beyond this the buffers hold whatever the allocator handed
+    # out, so nothing past it may be compared.
+    valid_rows = sum(split_sections)
+
+    wire = make_prequantized_wire_tensor(x, split_section_tensor)
+
+    # Snapshot what must survive verbatim, plus the compact scales the reference swizzles.
+    rowwise_data_before = wire.rowwise_data.clone()
+    wire_splits_before = [
+        (t._rowwise_data.view(dtype=torch.uint8).clone(), t._rowwise_scale_inv.clone())
+        for t in wire.split_into_quantized_tensors()
+    ]
+
+    # Reference high-precision input: everything downstream is derived from this. Only the live
+    # rows are kept -- dequantize allocates M rows but writes only the ones the groups cover.
+    dequantized_ref = (
+        tex.group_dequantize(wire, te.DType.kBFloat16)
+        .rowwise_data.view(M, N)[:valid_rows, :]
+        .clone()
+    )
+
+    # Reference columnwise copy, quantized per group from the dequantized data.
+    colwise_quantizers = [
+        MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=False, columnwise=True)
+        for _ in range(num_groups)
+    ]
+    _, _, colwise_data_ref, colwise_scale_ref = reference_group_quantize(
+        dequantized_ref,
+        colwise_quantizers,
+        split_sections,
+        return_rowwise=False,
+        return_transpose=True,
+    )
+
+    # ---- the code under test ----
+    colwise_quantizer = MXFP8Quantizer(
+        fp8_dtype=te.DType.kFloat8E4M3, rowwise=False, columnwise=True
+    )
+    colwise_quantizer.optimize_for_gemm = True
+    colwise_quantizer.internal = True
+    dequantized = tex.group_requantize_columnwise_and_swizzle_rowwise_(
+        wire,
+        colwise_quantizer,
+        num_groups,
+        split_section_tensor,
+        te.DType.kBFloat16,
+        return_dequantized=True,
+    )
+
+    assert wire.columnwise_data is not None, "columnwise data must be built"
+    assert wire.columnwise_scale_inv is not None, "columnwise scales must survive the swizzle"
+
+    # The returned dequantized tensor is what bias gradients are reduced from. Compare only the
+    # live rows: both this and the reference allocate M rows but write only the covered ones, and
+    # their tails are separate uninitialized allocations.
+    torch.testing.assert_close(
+        dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0
+    )
+
+    # The rowwise DATA must pass through untouched; only its scales are re-laid-out.
+    torch.testing.assert_close(wire.rowwise_data, rowwise_data_before, atol=0.0, rtol=0.0)
+
+    if valid_rows > 0:
+        # A tensor whose groups are all empty has no scales to lay out, so the swizzle is a no-op
+        # and leaves the flag unset; every other case must come back swizzled.
+        assert wire._with_gemm_swizzled_scales, "rowwise scales must be marked swizzled"
+
+    # Per-group comparison, same structure as check_grouped_tensor_mxfp8_versus_reference.
+    outputs = wire.split_into_quantized_tensors()
+    x_splits = torch.split(dequantized_ref, split_sections)
+
+    for i, out in enumerate(outputs):
+        rows_i = split_sections[i]
+        scale_before = wire_splits_before[i][1]
+        colwise_data = out._columnwise_data.view(dtype=torch.uint8)
+        colwise_scale = out._columnwise_scale_inv
+
+        if rows_i == 0:
+            # Buffers for empty groups are never written, so only shape and dtype are meaningful.
+            assert_same_shape_and_dtype(colwise_data, colwise_data_ref[i])
+            assert_same_shape_and_dtype(colwise_scale, colwise_scale_ref[i])
+            continue
+
+        # Rowwise scales: the swizzled form of the compact scales this group arrived with. The
+        # rowwise DATA is covered by the whole-buffer identity check above.
+        torch.testing.assert_close(
+            out._rowwise_scale_inv,
+            swizzle_mxfp8_scale(rows_i, N, scale_before, columnwise=False),
+            atol=0.0,
+            rtol=0.0,
+        )
+
+        # Columnwise: rebuilt from the dequantized data, and swizzled by the quantize kernel
+        # because the caller sets optimize_for_gemm.
+        torch.testing.assert_close(colwise_data, colwise_data_ref[i], atol=0.0, rtol=0.0)
+        valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, True)
+        assert (
+            valid_scale_shape == colwise_scale.shape
+        ), "The columnwise scale shape is not correctly aligned"
+        torch.testing.assert_close(
+            colwise_scale,
+            swizzle_mxfp8_scale(rows_i, N, colwise_scale_ref[i], columnwise=True),
+            atol=0.0,
+            rtol=0.0,
+        )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize(
+    "M, N",
+    [
+        # edge case, zero tokens for all
+        (0, 512),
+        # full tile cases
+        (1024, 256),
+        # larger sizes
+        (8192, 1024),
+        (16384, 8192),
+    ],
+)
+@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str)
+@pytest.mark.parametrize(
+    "edge_cases",
+    [
+        "regular",
+        "zero_tokens_front",
+        "zero_tokens_end",
+        "zero_tokens_middle",
+        "random_uneven_split",
+    ],
+)
+def test_prequantized_requantize_versus_reference(
+    x_dtype: torch.dtype,
+    M: int,
+    N: int,
+    edge_cases: str,
+) -> None:
+    split_sections = generate_split_sections(M, N, edge_cases)
+    check_prequantized_requantize_versus_reference(
+        x_dtype=x_dtype,
+        M=M,
+        N=N,
+        split_sections=split_sections,
+    )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize(
+    "M, N",
+    [
+        # M won't be empty in paged stashing
+        (1024, 256),
+        (8192, 1024),
+        (16384, 8192),
+    ],
+)
+@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str)
+@pytest.mark.parametrize(
+    "edge_cases",
+    [
+        "regular",
+        "zero_tokens_all",
+        "zero_tokens_front",
+        "zero_tokens_end",
+        "zero_tokens_middle",
+        "random_uneven_split",
+    ],
+)
+def test_prequantized_requantize_with_paged_stashing(
+    x_dtype: torch.dtype,
+    M: int,
+    N: int,
+    edge_cases: str,
+) -> None:
+    # Paged stashing: the buffer holds M rows but only valid_M carry live tokens; the rest is
+    # garbage the kernels must not touch.
+    valid_M = 0 if edge_cases == "zero_tokens_all" else M // 2
+    split_sections = generate_split_sections(valid_M, N, edge_cases)
+    assert sum(split_sections) == valid_M
+
+    check_prequantized_requantize_versus_reference(
+        x_dtype=x_dtype,
+        M=M,
+        N=N,
+        split_sections=split_sections,
+    )
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index db8099e2e2..dc2d8f24c3 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -707,6 +707,10 @@ py::object group_requantize_columnwise_and_swizzle_rowwise_(
   // resets columnwise_scale_inv to None, which would strand the columnwise data below with a
   // null scale pointer.
   grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false);
+  // The swizzle hands back a 2D [num_tensors * padded_m, padded_k] scale buffer, but grouped
+  // tensors carry scales as a flat array indexed by element offsets (scale_inv_offsets), so
+  // per-group slicing breaks unless it is flattened back.
+  grouped_x.attr("scale_inv") = grouped_x.attr("scale_inv").attr("reshape")(-1);
 
   // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
   // data because the two directions scale along perpendicular axes. The caller hands us a

From a3219de0ca8623eb3e90f52a66f92134ab499bf7 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 31 Jul 2026 05:12:31 +0000
Subject: [PATCH 11/18] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 .../pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index 6502a2b314..00e7673140 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -487,9 +487,7 @@ def test_grouped_tensor_mxfp8_with_paged_stashing(
 
 def make_prequantized_wire_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor):
     """Rowwise-only, unswizzled grouped tensor, as FP8 dispatch delivers it."""
-    wire_quantizer = MXFP8Quantizer(
-        fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False
-    )
+    wire_quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False)
     # Must stay unswizzled: the requantize path asserts on it and dequantize needs compact scales.
     wire_quantizer.optimize_for_gemm = False
     wire = fused_grouped_quantize(x, split_section_tensor, wire_quantizer)
@@ -575,9 +573,7 @@ def check_prequantized_requantize_versus_reference(
     # The returned dequantized tensor is what bias gradients are reduced from. Compare only the
     # live rows: both this and the reference allocate M rows but write only the covered ones, and
     # their tails are separate uninitialized allocations.
-    torch.testing.assert_close(
-        dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0
-    )
+    torch.testing.assert_close(dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0)
 
     # The rowwise DATA must pass through untouched; only its scales are re-laid-out.
     torch.testing.assert_close(wire.rowwise_data, rowwise_data_before, atol=0.0, rtol=0.0)

From 3bbf778264d99939aba6029f6a565105ff93f43e Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Tue, 4 Aug 2026 17:03:38 -0700
Subject: [PATCH 12/18] rename to group_requantize

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 .../mxfp8/test_mxfp8_group_quantize_graph_safe.py   |  4 ++--
 transformer_engine/pytorch/csrc/extensions.h        |  4 ++--
 transformer_engine/pytorch/csrc/extensions/cast.cpp | 13 +++++++++----
 .../pytorch/csrc/extensions/pybind.cpp              |  6 +++---
 .../pytorch/ops/basic/grouped_linear.py             |  4 ++--
 transformer_engine/pytorch/ops/fused/grouped_mlp.py |  4 ++--
 6 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index 00e7673140..500819453f 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -474,7 +474,7 @@ def test_grouped_tensor_mxfp8_with_paged_stashing(
 # ---------------------------------------------------------------------------------------------
 # Pre-quantized MXFP8 input (FP8 token dispatch)
 #
-# tex.group_requantize_columnwise_and_swizzle_rowwise_ takes a grouped tensor that arrives
+# tex.group_requantize takes a grouped tensor that arrives
 # ALREADY rowwise-quantized (its high-precision form no longer exists), and makes it GEMM-ready
 # in both directions: the rowwise data passes through verbatim with its scales swizzled, and the
 # columnwise copy is rebuilt via dequantize + columnwise-only requantize.
@@ -558,7 +558,7 @@ def check_prequantized_requantize_versus_reference(
     )
     colwise_quantizer.optimize_for_gemm = True
     colwise_quantizer.internal = True
-    dequantized = tex.group_requantize_columnwise_and_swizzle_rowwise_(
+    dequantized = tex.group_requantize(
         wire,
         colwise_quantizer,
         num_groups,
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 30df6d5975..b362fbd2d3 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -367,8 +367,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 std::optional<at::Tensor> last_dims,
                                 std::optional<at::Tensor> tensor_offsets);
 
-py::object group_requantize_columnwise_and_swizzle_rowwise_(
-    py::handle grouped_x, py::handle columnwise_quantizer, const size_t num_tensors,
+py::object group_requantize(
+    py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
     std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
     bool return_dequantized);
 
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index dc2d8f24c3..4eaccdcb42 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -663,12 +663,17 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o
   return py::reinterpret_borrow<py::object>(out_py);
 }
 
-py::object group_requantize_columnwise_and_swizzle_rowwise_(
-    py::handle grouped_x, py::handle columnwise_quantizer, const size_t num_tensors,
+py::object group_requantize(
+    py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
     std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
     bool return_dequantized) {
   init_extension();
 
+  // The rowwise data is reused verbatim and only its scales are swizzled, so a quantizer asking
+  // for rowwise output would silently overwrite it with a fresh quantization.
+  NVTE_CHECK(quantizer.attr("columnwise_usage").cast<bool>() &&
+                 !quantizer.attr("rowwise_usage").cast<bool>(),
+             "group_requantize expects a columnwise-only quantizer.");
   NVTE_CHECK(!grouped_x.attr("rowwise_data").is_none(),
              "Pre-quantized MXFP8 grouped input is missing rowwise data.");
   NVTE_CHECK(!grouped_x.attr("scale_inv").is_none(),
@@ -681,7 +686,7 @@ py::object group_requantize_columnwise_and_swizzle_rowwise_(
     // The GEMM consumes the input's rowwise data verbatim while the wgrad GEMM consumes the
     // columnwise copy built here, so a dtype mismatch would make the two directions disagree.
     NVTE_CHECK(grouped_x.attr("quantizer").attr("dtype").cast<DType>() ==
-                   columnwise_quantizer.attr("dtype").cast<DType>(),
+                   quantizer.attr("dtype").cast<DType>(),
                "Pre-quantized MXFP8 grouped input and the columnwise quantizer disagree on the "
                "FP8 dtype.");
   }
@@ -716,7 +721,7 @@ py::object group_requantize_columnwise_and_swizzle_rowwise_(
   // data because the two directions scale along perpendicular axes. The caller hands us a
   // columnwise-only, optimize_for_gemm quantizer, so the kernel emits swizzled columnwise
   // scales directly.
-  auto columnwise = group_quantize(dequantized, columnwise_quantizer, num_tensors, first_dims,
+  auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims,
                                    std::nullopt, tensor_offsets, std::nullopt);
   grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
   grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index 3a5049d7b8..53924cafbb 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -216,11 +216,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
         py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none());
-  m.def("group_requantize_columnwise_and_swizzle_rowwise_",
-        transformer_engine::pytorch::group_requantize_columnwise_and_swizzle_rowwise_,
+  m.def("group_requantize",
+        transformer_engine::pytorch::group_requantize,
         "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle "
         "its rowwise scales for GEMM, in place",
-        py::arg("grouped_x"), py::arg("columnwise_quantizer"), py::arg("num_tensors"),
+        py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"),
         py::arg("first_dims"), py::arg("otype"), py::arg("tensor_offsets") = py::none(),
         py::arg("return_dequantized") = false);
   m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize,
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index 499ab7d4cc..c5664e875a 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -1358,7 +1358,7 @@ def _fuser_forward_grouped_tensor(
             # rowwise scales for the GEMM.
             grouped_x = input_.copy()
             if weight_requires_grad:
-                tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                tex.group_requantize(
                     grouped_x,
                     make_columnwise_gemm_quantizer(input_quantizers[0]),
                     num_groups,
@@ -1765,7 +1765,7 @@ def _fuser_backward_grouped_tensor(
                 # which is only kept when there is a bias.
                 grouped_dy = grad_output.copy()
                 if ctx.weight_requires_grad:
-                    dy_2d = tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                    dy_2d = tex.group_requantize(
                         grouped_dy,
                         make_columnwise_gemm_quantizer(grad_output_quantizer),
                         num_groups,
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 60aaad622e..293f70e844 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -1225,7 +1225,7 @@ def fuser_forward(
                 # manufacture the columnwise copy needed by the wgrad GEMM
                 # and swizzle the rowwise scales for the forward GEMM.
                 if weight_requires_grad:
-                    tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                    tex.group_requantize(
                         grouped_fc1_x,
                         make_columnwise_gemm_quantizer(fc1_input_quantizer),
                         num_groups,
@@ -1878,7 +1878,7 @@ def fuser_backward(
                 grouped_fc2_dy = grad_output.copy()
                 need_dequantized = output_fc2_dbias or scale_bias
                 if fc2_ctx.weight_requires_grad:
-                    fc2_dy = tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                    fc2_dy = tex.group_requantize(
                         grouped_fc2_dy,
                         make_columnwise_gemm_quantizer(fc2_grad_output_quantizer),
                         num_groups,

From e890104acf92c924275286bf5bcbebc4986c5c64 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 6 Aug 2026 17:14:03 +0000
Subject: [PATCH 13/18] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/csrc/extensions.h          |  7 +++----
 transformer_engine/pytorch/csrc/extensions/cast.cpp   | 11 +++++------
 transformer_engine/pytorch/csrc/extensions/pybind.cpp |  7 +++----
 3 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index b362fbd2d3..9fb4387905 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -367,10 +367,9 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 std::optional<at::Tensor> last_dims,
                                 std::optional<at::Tensor> tensor_offsets);
 
-py::object group_requantize(
-    py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
-    std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
-    bool return_dequantized);
+py::object group_requantize(py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
+                            std::optional<at::Tensor> first_dims, DType otype,
+                            std::optional<at::Tensor> tensor_offsets, bool return_dequantized);
 
 std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
                                               std::vector<py::handle> quantizer_list);
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index 42773cb03a..248b356641 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -664,10 +664,9 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o
   return py::reinterpret_borrow<py::object>(out_py);
 }
 
-py::object group_requantize(
-    py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
-    std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
-    bool return_dequantized) {
+py::object group_requantize(py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
+                            std::optional<at::Tensor> first_dims, DType otype,
+                            std::optional<at::Tensor> tensor_offsets, bool return_dequantized) {
   init_extension();
 
   // The rowwise data is reused verbatim and only its scales are swizzled, so a quantizer asking
@@ -722,8 +721,8 @@ py::object group_requantize(
   // data because the two directions scale along perpendicular axes. The caller hands us a
   // columnwise-only, optimize_for_gemm quantizer, so the kernel emits swizzled columnwise
   // scales directly.
-  auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims,
-                                   std::nullopt, tensor_offsets, std::nullopt);
+  auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt,
+                                   tensor_offsets, std::nullopt);
   grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
   grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
 
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index 53924cafbb..f6301fdc5d 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -216,12 +216,11 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
         py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none());
-  m.def("group_requantize",
-        transformer_engine::pytorch::group_requantize,
+  m.def("group_requantize", transformer_engine::pytorch::group_requantize,
         "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle "
         "its rowwise scales for GEMM, in place",
-        py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"),
-        py::arg("first_dims"), py::arg("otype"), py::arg("tensor_offsets") = py::none(),
+        py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
+        py::arg("otype"), py::arg("tensor_offsets") = py::none(),
         py::arg("return_dequantized") = false);
   m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize,
         "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer"));

From 76662c2eaecbd541d1528e7c16c9220a22134a35 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Thu, 6 Aug 2026 11:14:16 -0700
Subject: [PATCH 14/18] rename func with inplace tag

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py | 4 ++--
 transformer_engine/pytorch/csrc/extensions.h                | 2 +-
 transformer_engine/pytorch/csrc/extensions/cast.cpp         | 4 ++--
 transformer_engine/pytorch/csrc/extensions/pybind.cpp       | 4 ++--
 transformer_engine/pytorch/ops/basic/grouped_linear.py      | 4 ++--
 transformer_engine/pytorch/ops/fused/grouped_mlp.py         | 4 ++--
 6 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index 500819453f..91c4eaa6e8 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -474,7 +474,7 @@ def test_grouped_tensor_mxfp8_with_paged_stashing(
 # ---------------------------------------------------------------------------------------------
 # Pre-quantized MXFP8 input (FP8 token dispatch)
 #
-# tex.group_requantize takes a grouped tensor that arrives
+# tex.group_requantize_inplace takes a grouped tensor that arrives
 # ALREADY rowwise-quantized (its high-precision form no longer exists), and makes it GEMM-ready
 # in both directions: the rowwise data passes through verbatim with its scales swizzled, and the
 # columnwise copy is rebuilt via dequantize + columnwise-only requantize.
@@ -558,7 +558,7 @@ def check_prequantized_requantize_versus_reference(
     )
     colwise_quantizer.optimize_for_gemm = True
     colwise_quantizer.internal = True
-    dequantized = tex.group_requantize(
+    dequantized = tex.group_requantize_inplace(
         wire,
         colwise_quantizer,
         num_groups,
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index b362fbd2d3..70cbe65df3 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -367,7 +367,7 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 std::optional<at::Tensor> last_dims,
                                 std::optional<at::Tensor> tensor_offsets);
 
-py::object group_requantize(
+py::object group_requantize_inplace(
     py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
     std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
     bool return_dequantized);
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index 42773cb03a..0abb4c42d0 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -664,7 +664,7 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o
   return py::reinterpret_borrow<py::object>(out_py);
 }
 
-py::object group_requantize(
+py::object group_requantize_inplace(
     py::handle grouped_x, py::handle quantizer, const size_t num_tensors,
     std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
     bool return_dequantized) {
@@ -674,7 +674,7 @@ py::object group_requantize(
   // for rowwise output would silently overwrite it with a fresh quantization.
   NVTE_CHECK(quantizer.attr("columnwise_usage").cast<bool>() &&
                  !quantizer.attr("rowwise_usage").cast<bool>(),
-             "group_requantize expects a columnwise-only quantizer.");
+             "group_requantize_inplace expects a columnwise-only quantizer.");
   NVTE_CHECK(!grouped_x.attr("rowwise_data").is_none(),
              "Pre-quantized MXFP8 grouped input is missing rowwise data.");
   NVTE_CHECK(!grouped_x.attr("scale_inv").is_none(),
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index 53924cafbb..356147c209 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -216,8 +216,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
         py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none());
-  m.def("group_requantize",
-        transformer_engine::pytorch::group_requantize,
+  m.def("group_requantize_inplace",
+        transformer_engine::pytorch::group_requantize_inplace,
         "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle "
         "its rowwise scales for GEMM, in place",
         py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"),
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index c5664e875a..5ea1083edf 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -1358,7 +1358,7 @@ def _fuser_forward_grouped_tensor(
             # rowwise scales for the GEMM.
             grouped_x = input_.copy()
             if weight_requires_grad:
-                tex.group_requantize(
+                tex.group_requantize_inplace(
                     grouped_x,
                     make_columnwise_gemm_quantizer(input_quantizers[0]),
                     num_groups,
@@ -1765,7 +1765,7 @@ def _fuser_backward_grouped_tensor(
                 # which is only kept when there is a bias.
                 grouped_dy = grad_output.copy()
                 if ctx.weight_requires_grad:
-                    dy_2d = tex.group_requantize(
+                    dy_2d = tex.group_requantize_inplace(
                         grouped_dy,
                         make_columnwise_gemm_quantizer(grad_output_quantizer),
                         num_groups,
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 293f70e844..4581e8c0e9 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -1225,7 +1225,7 @@ def fuser_forward(
                 # manufacture the columnwise copy needed by the wgrad GEMM
                 # and swizzle the rowwise scales for the forward GEMM.
                 if weight_requires_grad:
-                    tex.group_requantize(
+                    tex.group_requantize_inplace(
                         grouped_fc1_x,
                         make_columnwise_gemm_quantizer(fc1_input_quantizer),
                         num_groups,
@@ -1878,7 +1878,7 @@ def fuser_backward(
                 grouped_fc2_dy = grad_output.copy()
                 need_dequantized = output_fc2_dbias or scale_bias
                 if fc2_ctx.weight_requires_grad:
-                    fc2_dy = tex.group_requantize(
+                    fc2_dy = tex.group_requantize_inplace(
                         grouped_fc2_dy,
                         make_columnwise_gemm_quantizer(fc2_grad_output_quantizer),
                         num_groups,

From a056c6dbf23dd17d0f0093ce4b0333648864125e Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 6 Aug 2026 18:20:38 +0000
Subject: [PATCH 15/18] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/csrc/extensions.h        | 5 ++---
 transformer_engine/pytorch/csrc/extensions/cast.cpp | 5 ++---
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 3d3e5f5c04..49dd68ebf2 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -368,9 +368,8 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 std::optional<at::Tensor> tensor_offsets);
 
 py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
-                                    const size_t num_tensors,
-                                    std::optional<at::Tensor> first_dims, DType otype,
-                                    std::optional<at::Tensor> tensor_offsets,
+                                    const size_t num_tensors, std::optional<at::Tensor> first_dims,
+                                    DType otype, std::optional<at::Tensor> tensor_offsets,
                                     bool return_dequantized);
 
 std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index a4249a7742..860cd6bda1 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -665,9 +665,8 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o
 }
 
 py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
-                                    const size_t num_tensors,
-                                    std::optional<at::Tensor> first_dims, DType otype,
-                                    std::optional<at::Tensor> tensor_offsets,
+                                    const size_t num_tensors, std::optional<at::Tensor> first_dims,
+                                    DType otype, std::optional<at::Tensor> tensor_offsets,
                                     bool return_dequantized) {
   init_extension();
 

From ce2e5a31849f05b93ac8bacbe8dd0808b81febf5 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Thu, 6 Aug 2026 16:10:59 -0700
Subject: [PATCH 16/18] refactor the requantize function

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 .../test_mxfp8_group_quantize_graph_safe.py   | 117 +++++++++++++++++-
 .../pytorch/csrc/extensions/cast.cpp          | 100 +++++++++------
 transformer_engine/pytorch/ops/_common.py     |   9 --
 .../pytorch/ops/basic/grouped_linear.py       |  94 +++++---------
 .../pytorch/ops/fused/grouped_mlp.py          | 113 +++++------------
 5 files changed, 241 insertions(+), 192 deletions(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index 91c4eaa6e8..9b8b82c9bf 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -496,6 +496,28 @@ def make_prequantized_wire_tensor(x: torch.Tensor, split_section_tensor: torch.T
     return wire
 
 
+def make_op_quantizer(columnwise: bool):
+    """The op's input quantizer, configured the way the ops layer configures it.
+
+    ``columnwise`` mirrors ``weight_requires_grad``: it tells the helper whether a wgrad GEMM
+    will consume a columnwise copy.
+    """
+    quantizer = MXFP8Quantizer(
+        fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=columnwise
+    )
+    quantizer.optimize_for_gemm = True
+    return quantizer
+
+
+def make_gemm_ready_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor):
+    """Grouped tensor already GEMM-ready in both directions (swizzled scales)."""
+    quantizer = make_op_quantizer(columnwise=True)
+    tensor = fused_grouped_quantize(x, split_section_tensor, quantizer)
+    assert tensor.columnwise_data is not None
+    assert tensor._with_gemm_swizzled_scales
+    return tensor
+
+
 def check_prequantized_requantize_versus_reference(
     x_dtype: torch.dtype,
     M: int,
@@ -553,14 +575,11 @@ def check_prequantized_requantize_versus_reference(
     )
 
     # ---- the code under test ----
-    colwise_quantizer = MXFP8Quantizer(
-        fp8_dtype=te.DType.kFloat8E4M3, rowwise=False, columnwise=True
-    )
-    colwise_quantizer.optimize_for_gemm = True
-    colwise_quantizer.internal = True
+    # The op's quantizer, configured as the ops layer does: columnwise_usage says a wgrad GEMM
+    # will consume the columnwise copy, so the helper builds it and switches rowwise off itself.
     dequantized = tex.group_requantize_inplace(
         wire,
-        colwise_quantizer,
+        make_op_quantizer(columnwise=True),
         num_groups,
         split_section_tensor,
         te.DType.kBFloat16,
@@ -702,3 +721,89 @@ def test_prequantized_requantize_with_paged_stashing(
         N=N,
         split_sections=split_sections,
     )
+
+
+def _requantize_setup(M: int = 1024, N: int = 256):
+    """Common inputs for the state-dispatch tests below."""
+    torch.manual_seed(0)
+    split_sections = [M // 4] * 4
+    x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda")
+    return x, torch.tensor(split_sections, dtype=torch.int64, device="cuda"), len(split_sections)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_passes_through_gemm_ready_input():
+    """A tensor already GEMM-ready in both directions is left untouched."""
+    x, splits, num_groups = _requantize_setup()
+    tensor = make_gemm_ready_tensor(x, splits)
+    rowwise_before = tensor.rowwise_data.clone()
+    columnwise_before = tensor.columnwise_data.clone()
+    scale_before = tensor.scale_inv.clone()
+
+    out = tex.group_requantize_inplace(
+        tensor, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16
+    )
+
+    assert out is None
+    assert tensor._with_gemm_swizzled_scales
+    torch.testing.assert_close(tensor.rowwise_data, rowwise_before, atol=0.0, rtol=0.0)
+    torch.testing.assert_close(tensor.columnwise_data, columnwise_before, atol=0.0, rtol=0.0)
+    torch.testing.assert_close(tensor.scale_inv, scale_before, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_skips_columnwise_when_not_needed():
+    """columnwise_usage=False (frozen weights) swizzles rowwise without building columnwise."""
+    x, splits, num_groups = _requantize_setup()
+    wire = make_prequantized_wire_tensor(x, splits)
+
+    out = tex.group_requantize_inplace(
+        wire, make_op_quantizer(columnwise=False), num_groups, splits, te.DType.kBFloat16
+    )
+
+    assert out is None
+    assert wire._with_gemm_swizzled_scales, "the GEMM still needs swizzled rowwise scales"
+    assert wire.columnwise_data is None, "no wgrad GEMM, so no columnwise copy should be built"
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_rejects_dequantized_from_gemm_ready_input():
+    """Bias grads cannot be served from an already-swizzled input, so this must raise."""
+    x, splits, num_groups = _requantize_setup()
+    tensor = make_gemm_ready_tensor(x, splits)
+
+    with pytest.raises(RuntimeError, match="compact format"):
+        tex.group_requantize_inplace(
+            tensor,
+            make_op_quantizer(columnwise=True),
+            num_groups,
+            splits,
+            te.DType.kBFloat16,
+            return_dequantized=True,
+        )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_rejects_swizzled_without_columnwise():
+    """Swizzled rowwise scales with no columnwise copy: it can no longer be rebuilt."""
+    x, splits, num_groups = _requantize_setup()
+    wire = make_prequantized_wire_tensor(x, splits)
+    # Swizzle in place, leaving the tensor rowwise-only.
+    tex.grouped_swizzle_for_gemm(wire, True, False)
+
+    with pytest.raises(RuntimeError, match="columnwise copy"):
+        tex.group_requantize_inplace(
+            wire, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16
+        )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_rejects_dtype_mismatch():
+    """The helper keeps the input's format; it does not convert between formats."""
+    x, splits, num_groups = _requantize_setup()
+    wire = make_prequantized_wire_tensor(x, splits)
+    mismatched = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E5M2, rowwise=True, columnwise=True)
+    mismatched.optimize_for_gemm = True
+
+    with pytest.raises(RuntimeError, match="dtype"):
+        tex.group_requantize_inplace(wire, mismatched, num_groups, splits, te.DType.kBFloat16)
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index 860cd6bda1..f149a0c570 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -670,27 +670,49 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
                                     bool return_dequantized) {
   init_extension();
 
-  // The rowwise data is reused verbatim and only its scales are swizzled, so a quantizer asking
-  // for rowwise output would silently overwrite it with a fresh quantization.
-  NVTE_CHECK(quantizer.attr("columnwise_usage").cast<bool>() &&
-                 !quantizer.attr("rowwise_usage").cast<bool>(),
-             "group_requantize_inplace expects a columnwise-only quantizer.");
-  NVTE_CHECK(!grouped_x.attr("rowwise_data").is_none(),
-             "Pre-quantized MXFP8 grouped input is missing rowwise data.");
-  NVTE_CHECK(!grouped_x.attr("scale_inv").is_none(),
-             "Pre-quantized MXFP8 grouped input is missing rowwise scales.");
-  NVTE_CHECK(!grouped_x.attr("_with_gemm_swizzled_scales").cast<bool>(),
-             "Pre-quantized MXFP8 grouped input must have unswizzled scales.");
-  NVTE_CHECK(grouped_x.attr("columnwise_data").is_none(),
-             "Pre-quantized MXFP8 grouped input must be rowwise-only.");
-  if (!grouped_x.attr("quantizer").is_none()) {
-    // The GEMM consumes the input's rowwise data verbatim while the wgrad GEMM consumes the
-    // columnwise copy built here, so a dtype mismatch would make the two directions disagree.
-    NVTE_CHECK(grouped_x.attr("quantizer").attr("dtype").cast<DType>() ==
-                   quantizer.attr("dtype").cast<DType>(),
-               "Pre-quantized MXFP8 grouped input and the columnwise quantizer disagree on the "
-               "FP8 dtype.");
-  }
+  const bool has_rowwise =
+      !grouped_x.attr("rowwise_data").is_none() && !grouped_x.attr("scale_inv").is_none();
+  const bool has_columnwise = !grouped_x.attr("columnwise_data").is_none() &&
+                              !grouped_x.attr("columnwise_scale_inv").is_none();
+  const bool swizzled = grouped_x.attr("_with_gemm_swizzled_scales").cast<bool>();
+
+  NVTE_CHECK(has_rowwise,
+             "Grouped input has no rowwise data and scales for the GEMM to consume.");
+
+  // The tensor's own quantization must match what the op expects on every path: even a
+  // pass-through hands its data straight to the GEMM. This keeps the input's format rather than
+  // converting between formats.
+  const auto input_quantizer = grouped_x.attr("quantizer");
+  NVTE_CHECK(!input_quantizer.is_none(), "Grouped input has no quantizer.");
+  NVTE_CHECK(Py_TYPE(input_quantizer.ptr()) == Py_TYPE(quantizer.ptr()),
+             "Grouped input and the op disagree on quantization format.");
+  NVTE_CHECK(input_quantizer.attr("dtype").cast<DType>() == quantizer.attr("dtype").cast<DType>(),
+             "Grouped input and the quantizer disagree on the FP8 dtype.");
+
+  // The columnwise copy is only worth building when a wgrad GEMM will consume it. Read this
+  // before the usage is overridden below.
+  const bool need_columnwise = quantizer.attr("columnwise_usage").cast<bool>();
+
+  if (swizzled) {
+    // Already GEMM-ready. Nothing can be derived from here, since dequantization requires scales
+    // in compact format, so the input must already carry everything that will be consumed. No
+    // quantization kernel runs, which makes this path format-agnostic.
+    NVTE_CHECK(has_columnwise || !need_columnwise,
+               "Grouped input has swizzled scales but no columnwise data for the wgrad GEMM. It "
+               "cannot be rebuilt, because dequantization requires scales in compact format.");
+    NVTE_CHECK(!return_dequantized,
+               "Cannot return a dequantized tensor for an already-swizzled grouped input: "
+               "dequantization requires scales in compact format.");
+    return py::none();
+  }
+
+  NVTE_CHECK(!has_columnwise,
+             "Grouped input already has columnwise data but unswizzled scales; requantizing "
+             "from it is not supported.");
+
+  // Everything below runs quantization kernels, so it is MXFP8-only.
+  NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()),
+             "Requantizing a grouped input is only supported for MXFP8.");
 
   const auto logical_shape = grouped_x.attr("logical_shape").cast<py::tuple>();
   const auto total_tokens = logical_shape[0].cast<size_t>();
@@ -699,15 +721,18 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
   // on a swizzle-tile boundary. Those counts live on the device (host reads would break CUDA
   // graph capture), so that half is the caller's contract rather than an assertion.
   NVTE_CHECK(total_tokens % 128 == 0 && hidden_dim % 128 == 0,
-             "Pre-quantized MXFP8 grouped input requires dims that are multiples of 128, but got (",
+             "Requantizing a grouped input requires dims that are multiples of 128, but got (",
              total_tokens, ", ", hidden_dim, ").");
 
-  // Dequantize first: it reads the rowwise scales, which the swizzle below replaces.
-  auto dequantized_grouped = group_dequantize(grouped_x, otype);
-  auto dequantized =
-      dequantized_grouped.attr("rowwise_data")
-          .cast<at::Tensor>()
-          .view({static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
+  // Dequantize first: it reads the rowwise scales, which the swizzle below replaces. Left
+  // undefined when nothing consumes it, which skips the pass entirely.
+  at::Tensor dequantized;
+  if (need_columnwise || return_dequantized) {
+    dequantized = group_dequantize(grouped_x, otype)
+                      .attr("rowwise_data")
+                      .cast<at::Tensor>()
+                      .view({static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
+  }
 
   // Swizzle the rowwise scales before attaching any columnwise data: a rowwise-only swizzle
   // resets columnwise_scale_inv to None, which would strand the columnwise data below with a
@@ -718,14 +743,17 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
   // per-group slicing breaks unless it is flattened back.
   grouped_x.attr("scale_inv") = grouped_x.attr("scale_inv").attr("reshape")(-1);
 
-  // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
-  // data because the two directions scale along perpendicular axes. The caller hands us a
-  // columnwise-only, optimize_for_gemm quantizer, so the kernel emits swizzled columnwise
-  // scales directly.
-  auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt,
-                                   tensor_offsets, std::nullopt);
-  grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
-  grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
+  if (need_columnwise) {
+    // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
+    // data because the two directions scale along perpendicular axes. Quantizing rowwise as well
+    // would redo work we already have, so that direction is switched off; the caller sets
+    // optimize_for_gemm, which makes the kernel emit swizzled columnwise scales directly.
+    quantizer.attr("set_usage")(py::arg("rowwise") = false, py::arg("columnwise") = true);
+    auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt,
+                                     tensor_offsets, std::nullopt);
+    grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
+    grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
+  }
 
   if (return_dequantized) {
     return py::cast(dequantized);
diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py
index 5bef5143dd..f39115c4c6 100644
--- a/transformer_engine/pytorch/ops/_common.py
+++ b/transformer_engine/pytorch/ops/_common.py
@@ -91,15 +91,6 @@ def maybe_dequantize(
     return tensor
 
 
-def make_columnwise_gemm_quantizer(quantizer: MXFP8Quantizer) -> MXFP8Quantizer:
-    """Copy of ``quantizer`` configured to emit only GEMM-swizzled columnwise data."""
-    columnwise_quantizer = quantizer.copy()
-    columnwise_quantizer.set_usage(rowwise=False, columnwise=True)
-    columnwise_quantizer.optimize_for_gemm = True
-    columnwise_quantizer.internal = True
-    return columnwise_quantizer
-
-
 def maybe_autocast_dtype(
     *,
     device_type: str = "cuda",
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index 5ea1083edf..8eb911ef81 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -48,7 +48,6 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
-    make_columnwise_gemm_quantizer,
     maybe_dequantize,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
@@ -1331,13 +1330,12 @@ def _fuser_forward_grouped_tensor(
 
         # Flatten to 2D so the first dim is the total token count.
         original_shape = list(input_.size())
-        prequantized_mxfp8_input = (
-            with_quantized_compute
-            and isinstance(input_, GroupedTensor)
-            and isinstance(input_quantizers[0], MXFP8Quantizer)
-            and isinstance(input_.quantizer, MXFP8Quantizer)
-        )
-        if prequantized_mxfp8_input:
+        prequantized_input = with_quantized_compute and isinstance(input_, GroupedTensor)
+        if with_quantized_compute:
+            input_quantizer = input_quantizers[0]
+            input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
+            input_quantizer.optimize_for_gemm = True
+        if prequantized_input:
             # GroupedTensor forbids reshape and is already in the canonical
             # (total_tokens, in_features) layout; just validate the shape.
             if input_.dim() != 2 or input_.size(-1) != self.in_features:
@@ -1351,29 +1349,19 @@ def _fuser_forward_grouped_tensor(
             total_tokens = x.size(0)
 
         # Build the input GroupedTensor.
-        if prequantized_mxfp8_input:
-            # Rowwise-only MXFP8 input (e.g. FP8 token dispatch): feed the
-            # rowwise data to the forward GEMM as-is, manufacture the
-            # columnwise copy needed by the wgrad GEMM, and swizzle the
-            # rowwise scales for the GEMM.
+        if prequantized_input:
+            # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data
+            # for the GEMM and let the helper supply whatever else the GEMMs need.
             grouped_x = input_.copy()
-            if weight_requires_grad:
-                tex.group_requantize_inplace(
-                    grouped_x,
-                    make_columnwise_gemm_quantizer(input_quantizers[0]),
-                    num_groups,
-                    split_sizes,
-                    TE_DType[dtype],
-                    tensor_offsets=base_split_offsets * self.in_features,
-                )
-            else:
-                # No wgrad, so no columnwise copy is needed. The forward GEMM
-                # still requires swizzled rowwise scales.
-                tex.grouped_swizzle_for_gemm(grouped_x, rowwise=True, columnwise=False)
+            tex.group_requantize_inplace(
+                grouped_x,
+                input_quantizer,
+                num_groups,
+                split_sizes,
+                TE_DType[dtype],
+                tensor_offsets=base_split_offsets * self.in_features,
+            )
         elif with_quantized_compute:
-            input_quantizer = input_quantizers[0]
-            input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
-            input_quantizer.optimize_for_gemm = True
             grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes)
         else:
             # No quantize: wrap the contiguous high-precision buffer.
@@ -1723,13 +1711,8 @@ def _fuser_backward_grouped_tensor(
 
         # Flatten grad_output to 2D (total_tokens, out_features)
         # to figure out total tokens.
-        prequantized_mxfp8_grad = (
-            with_quantized_compute
-            and isinstance(grad_output, GroupedTensor)
-            and isinstance(ctx.grad_output_quantizers[0], MXFP8Quantizer)
-            and isinstance(grad_output.quantizer, MXFP8Quantizer)
-        )
-        if prequantized_mxfp8_grad:
+        prequantized_grad = with_quantized_compute and isinstance(grad_output, GroupedTensor)
+        if prequantized_grad:
             # GroupedTensor forbids reshape and is already in the canonical
             # (total_tokens, out_features) layout; just validate the shape.
             if grad_output.dim() != 2 or grad_output.size(-1) != self.out_features:
@@ -1758,33 +1741,20 @@ def _fuser_backward_grouped_tensor(
             fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or (
                 isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad
             )
-            if prequantized_mxfp8_grad:
-                # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse the
-                # rowwise data for the dgrad GEMM and manufacture the columnwise copy
-                # for wgrad. Bias grads are reduced from the dequantized grad below,
-                # which is only kept when there is a bias.
+            if prequantized_grad:
+                # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its
+                # rowwise data for the dgrad GEMM. Bias grads are reduced from the dequantized
+                # grad below, which is only kept when there is a bias.
                 grouped_dy = grad_output.copy()
-                if ctx.weight_requires_grad:
-                    dy_2d = tex.group_requantize_inplace(
-                        grouped_dy,
-                        make_columnwise_gemm_quantizer(grad_output_quantizer),
-                        num_groups,
-                        split_sizes,
-                        TE_DType[dtype],
-                        tensor_offsets=base_split_offsets * self.out_features,
-                        return_dequantized=has_bias,
-                    )
-                else:
-                    # No wgrad, so no columnwise copy is needed. Dequantize before
-                    # swizzling: dequantization reads the unswizzled rowwise scales.
-                    dy_2d = (
-                        tex.group_dequantize(grouped_dy, TE_DType[dtype]).rowwise_data.view(
-                            total_tokens, self.out_features
-                        )
-                        if has_bias
-                        else None
-                    )
-                    tex.grouped_swizzle_for_gemm(grouped_dy, rowwise=True, columnwise=False)
+                dy_2d = tex.group_requantize_inplace(
+                    grouped_dy,
+                    grad_output_quantizer,
+                    num_groups,
+                    split_sizes,
+                    TE_DType[dtype],
+                    tensor_offsets=base_split_offsets * self.out_features,
+                    return_dequantized=has_bias,
+                )
             elif has_bias and not self._scale_bias and fuse_bgrad:
                 grouped_dy, dbias_packed = tex.bgrad_group_quantize(
                     dy_2d, grad_output_quantizer, num_groups, split_sizes
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 4581e8c0e9..518e910625 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -53,7 +53,6 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
-    make_columnwise_gemm_quantizer,
     maybe_dequantize,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
@@ -1209,34 +1208,19 @@ def fuser_forward(
         )
         fc1_input_quantizer.optimize_for_gemm = True
         fc1_input_quantizer.internal = True
-        input_quantizer = getattr(input_, "quantizer", None)
-        if isinstance(input_, GroupedTensor) and (
-            isinstance(fc1_input_quantizer, MXFP8Quantizer)
-            and isinstance(input_quantizer, MXFP8Quantizer)
-            or isinstance(fc1_input_quantizer, NVFP4Quantizer)
-            and isinstance(input_quantizer, NVFP4Quantizer)
-        ):
+        if isinstance(input_, GroupedTensor):
+            # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data
+            # for the GEMM and let the helper supply whatever else the GEMMs need. An input that
+            # is already GEMM-ready in both directions passes through untouched.
             grouped_fc1_x = input_.copy()
-            if (
-                isinstance(fc1_input_quantizer, MXFP8Quantizer)
-                and not grouped_fc1_x._with_gemm_swizzled_scales
-            ):
-                # Rowwise-only MXFP8 input (e.g. FP8 token dispatch):
-                # manufacture the columnwise copy needed by the wgrad GEMM
-                # and swizzle the rowwise scales for the forward GEMM.
-                if weight_requires_grad:
-                    tex.group_requantize_inplace(
-                        grouped_fc1_x,
-                        make_columnwise_gemm_quantizer(fc1_input_quantizer),
-                        num_groups,
-                        split_sizes,
-                        TE_DType[dtype],
-                        tensor_offsets=fc1_x_tensor_offsets,
-                    )
-                else:
-                    # No wgrad, so no columnwise copy is needed. The forward GEMM
-                    # still requires swizzled rowwise scales.
-                    tex.grouped_swizzle_for_gemm(grouped_fc1_x, rowwise=True, columnwise=False)
+            tex.group_requantize_inplace(
+                grouped_fc1_x,
+                fc1_input_quantizer,
+                num_groups,
+                split_sizes,
+                TE_DType[dtype],
+                tensor_offsets=fc1_x_tensor_offsets,
+            )
         else:
             fc1_x = maybe_dequantize(input_, dtype)
             grouped_fc1_x = _group_quantize_for_grouped_mlp(
@@ -1856,57 +1840,28 @@ def fuser_backward(
         output_fc2_dbias = fc2_op.has_bias
         fc2_dbias_packed = None
         fc2_dy = None
-        grad_output_quantizer = getattr(grad_output, "quantizer", None)
-        fc2_grad_output_quantizer_matches = (
-            isinstance(fc2_grad_output_quantizer, MXFP8Quantizer)
-            and isinstance(grad_output_quantizer, MXFP8Quantizer)
-        ) or (
-            isinstance(fc2_grad_output_quantizer, NVFP4Quantizer)
-            and isinstance(grad_output_quantizer, NVFP4Quantizer)
-        )
-        prequantized_mxfp8_grad = isinstance(fc2_grad_output_quantizer, MXFP8Quantizer)
-        if (
-            (not output_fc2_dbias or prequantized_mxfp8_grad)
-            and isinstance(grad_output, GroupedTensor)
-            and fc2_grad_output_quantizer_matches
-        ):
-            if prequantized_mxfp8_grad:
-                # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse
-                # the rowwise data for the dgrad GEMM and manufacture FC2's
-                # columnwise copy for wgrad. Bias grads are reduced from the
-                # dequantized grad, which is only materialized when one is needed.
-                grouped_fc2_dy = grad_output.copy()
-                need_dequantized = output_fc2_dbias or scale_bias
-                if fc2_ctx.weight_requires_grad:
-                    fc2_dy = tex.group_requantize_inplace(
-                        grouped_fc2_dy,
-                        make_columnwise_gemm_quantizer(fc2_grad_output_quantizer),
-                        num_groups,
-                        split_sizes,
-                        TE_DType[dtype],
-                        tensor_offsets=base_split_offsets * fc2_weight_shape[0],
-                        return_dequantized=need_dequantized,
-                    )
-                else:
-                    # No wgrad, so no columnwise copy is needed. Dequantize before
-                    # swizzling: dequantization reads the unswizzled rowwise scales.
-                    fc2_dy = (
-                        tex.group_dequantize(grouped_fc2_dy, TE_DType[dtype]).rowwise_data.view(
-                            grouped_fc2_dy.logical_shape
-                        )
-                        if need_dequantized
-                        else None
-                    )
-                    tex.grouped_swizzle_for_gemm(grouped_fc2_dy, rowwise=True, columnwise=False)
-                if output_fc2_dbias and not scale_bias:
-                    # This path has no quantize kernel to fuse dbias into, and the
-                    # consumer below has no fallback, so reduce it here.
-                    fc2_dbias_packed = compute_grouped_dbias(fc2_dy, base_split_offsets, num_groups)
-                    # scale_bias is the only later consumer of the dequantized grad;
-                    # drop it so the buffer is freed rather than held until backward ends.
-                    fc2_dy = None
-            else:
-                grouped_fc2_dy = grad_output
+        if isinstance(grad_output, GroupedTensor):
+            # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise
+            # data for the dgrad GEMM. Bias grads are reduced from the dequantized grad, which is
+            # only materialized when one is needed. A grad that is already GEMM-ready in both
+            # directions passes through untouched.
+            grouped_fc2_dy = grad_output.copy()
+            fc2_dy = tex.group_requantize_inplace(
+                grouped_fc2_dy,
+                fc2_grad_output_quantizer,
+                num_groups,
+                split_sizes,
+                TE_DType[dtype],
+                tensor_offsets=base_split_offsets * fc2_weight_shape[0],
+                return_dequantized=output_fc2_dbias or scale_bias,
+            )
+            if output_fc2_dbias and not scale_bias:
+                # This path has no quantize kernel to fuse dbias into, and the consumer below
+                # has no fallback, so reduce it here.
+                fc2_dbias_packed = compute_grouped_dbias(fc2_dy, base_split_offsets, num_groups)
+                # scale_bias is the only later consumer of the dequantized grad; drop it so the
+                # buffer is freed rather than held until backward ends.
+                fc2_dy = None
         else:
             fc2_dy = maybe_dequantize(grad_output, dtype)
             if output_fc2_dbias and not scale_bias:

From 226405d9c5dbf25ac62cf5e941ebcb1d1831b798 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Thu, 6 Aug 2026 23:55:51 +0000
Subject: [PATCH 17/18] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py | 4 +---
 transformer_engine/pytorch/csrc/extensions/cast.cpp         | 3 +--
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index 9b8b82c9bf..36f9c1e91a 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -502,9 +502,7 @@ def make_op_quantizer(columnwise: bool):
     ``columnwise`` mirrors ``weight_requires_grad``: it tells the helper whether a wgrad GEMM
     will consume a columnwise copy.
     """
-    quantizer = MXFP8Quantizer(
-        fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=columnwise
-    )
+    quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=columnwise)
     quantizer.optimize_for_gemm = True
     return quantizer
 
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index f149a0c570..5ce0261c82 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -676,8 +676,7 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
                               !grouped_x.attr("columnwise_scale_inv").is_none();
   const bool swizzled = grouped_x.attr("_with_gemm_swizzled_scales").cast<bool>();
 
-  NVTE_CHECK(has_rowwise,
-             "Grouped input has no rowwise data and scales for the GEMM to consume.");
+  NVTE_CHECK(has_rowwise, "Grouped input has no rowwise data and scales for the GEMM to consume.");
 
   // The tensor's own quantization must match what the op expects on every path: even a
   // pass-through hands its data straight to the GEMM. This keeps the input's format rather than

From 3032804aafcb3f574368c0800fb4f8842f526607 Mon Sep 17 00:00:00 2001
From: YangFei1990 <feiw@nvidia.com>
Date: Thu, 6 Aug 2026 17:34:20 -0700
Subject: [PATCH 18/18] minor fix for test

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index 36f9c1e91a..94861ea433 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -789,7 +789,7 @@ def test_prequantized_requantize_rejects_swizzled_without_columnwise():
     # Swizzle in place, leaving the tensor rowwise-only.
     tex.grouped_swizzle_for_gemm(wire, True, False)
 
-    with pytest.raises(RuntimeError, match="columnwise copy"):
+    with pytest.raises(RuntimeError, match="cannot be rebuilt"):
         tex.group_requantize_inplace(
             wire, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16
         )