[Pytorch] Add support for row-wise quanted input for grouped gemm by YangFei1990 · Pull Request #3244 · NVIDIA/TransformerEngine

YangFei1990 · 2026-07-23T18:27:58Z

Description

To support the NCCL EP FP8 dispatch, this PR introduce a capability to allow grouped gemm to take input that is already quantized in row-wise.

For forward, we directly pass the row-wise quantized for forward grouped gemm and dequant + requant with col-wise to save for backward.
For backward, we pass the row-wise quantized for dgrad computation and dequant + requant with col-wise for wgrad

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Signed-off-by: YangFei1990 <feiw@nvidia.com>

phu0ngng · 2026-07-23T20:22:01Z

+            and isinstance(input_quantizers[0], MXFP8Quantizer)
+            and isinstance(input_.quantizer, MXFP8Quantizer)


Should we also assert input_quantizers[0] == input_.quantizer?

I think == comparison is too strict. It will compare all value fields, I think things like columnwise_usage or optimize_for_gemm could mismatch? Maybe we should explicit assert dype to be E4M3? But let me know if the dispatch output tensor's quantizer can match all value fields.

phu0ngng · 2026-07-23T20:31:34Z

+    with_columnwise: bool,
+    tensor_offsets: Optional[torch.Tensor] = None,
+) -> None:
+    """Prepare a rowwise-only MXFP8 grouped input for grouped GEMM (in place).


Looks like this helper actually "prepares" colwise besides rowwise. I think we should make the function name and docstring clearer and more consistent.

phu0ngng · 2026-07-23T20:32:32Z

+        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."


compact -> unswizzled.

phu0ngng · 2026-07-23T20:40:30Z

+    scale_shape = (
+        round_up_to_nearest_multiple(total_tokens, 128),
+        round_up_to_nearest_multiple(cols // 32, 4),
+    )


Do we need to do this while we already have 128 per-expert-alignment in Dispatch? The total tokens should already be divisible by 128.

For hidden size, there will be an assertion in Dispatch to expect % 512 = 0.

Here, I think we can add assertions instead.

phu0ngng · 2026-07-24T23:19:51Z

+
+def prepare_prequantized_mxfp8_grouped_input(
+    grouped_x: GroupedTensorStorage,
+    quantizer: MXFP8Quantizer,


Actually, is there a GroupedQuantizer in TE/PyT? If so, we probably need to use it instead.

I do not see a GroupedQuantizer in Pytorch. I see what used is tex.group_quantize. Let me know if I missed it. Also can you help me understand the benefit of GroupedQuantizer and how it compared with current Pytorch approach?

Signed-off-by: YangFei1990 <feiw@nvidia.com>

greptile-apps · 2026-07-26T04:39:50Z

Greptile Summary

Adds row-wise prequantized MXFP8 support to grouped GEMM paths.

Repackages grouped tensors into saveable storage and manufactures columnwise data by dequantizing and requantizing when wgrad needs it.
Reuses row-wise data for forward and dgrad GEMMs while supporting ordinary and scaled-bias gradient reductions.
Extends grouped linear and fused grouped MLP tests for prequantized inputs and gradients.

Confidence Score: 3/5

The PR is not yet safe to merge because fused grouped MLP still aborts with frozen expert weights when input or routing gradients are required.

The basic grouped-linear fixes now supply dequantized gradients for scaled bias and directly reduce ordinary dbias without a wgrad stage, but the fused forward still omits columnwise activation data when weights are frozen and then asserts that this data exists while saving tensors, preventing the targeted frozen-expert backward flow from running.

Files Needing Attention: transformer_engine/pytorch/ops/fused/grouped_mlp.py, transformer_engine/pytorch/utils.py, tests/pytorch/test_grouped_mlp.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/_common.py	Adds shared conversion and preparation helpers for row-wise prequantized MXFP8 grouped tensors, including columnwise requantization and direct dbias reduction.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Integrates prequantized MXFP8 grouped inputs and gradients into basic grouped-linear forward and backward paths.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Integrates the same prequantized flow into fused grouped MLP, but its existing activation-saving invariant still blocks the frozen-weight configuration targeted by the prior scaled-bias report.
tests/pytorch/test_grouped_mlp.py	Adds numerical coverage for row-wise prequantized grouped inputs and gradients, while explicitly skipping fused backward with frozen weights.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Row-wise MXFP8 GroupedTensor] --> B[Repack as GroupedTensorStorage]
  B --> C[Swizzle row-wise scales]
  B --> D{Columnwise copy needed?}
  D -->|Yes| E[Dequantize]
  E --> F[Columnwise-only requantize]
  C --> G[Forward or dgrad grouped GEMM]
  F --> H[Wgrad grouped GEMM]
  E --> I[Bias and routing-scale reductions]

_{Reviews (4): Last reviewed commit: "allow bias + frozen weight path in prequ..." | Re-trigger Greptile}

Signed-off-by: YangFei1990 <feiw@nvidia.com>

YangFei1990 · 2026-07-26T20:24:46Z

/te-ci pytorch

zhongbozhu · 2026-07-27T23:36:19Z

+        requires_grad=False,
+        with_gemm_swizzled_scales=False,
+    )
+    tex.swizzle_scales_for_gemm_(tmp)


It feels like we are swizzling as if the activation is a dense tensor, which is only valid is the activation has tokens-per-expert padded to 128 multiple or higher (like 256).

How are we handling the swizzle of columnwise scales?

Yes this is the precondition, that we require the input dimension (both row and col) to be 128 multiple or higher. While I built in this way, please let me know if we should handle padding here.

@zhongbozhu I think the group_quantize + swizzle fusion generates the swizzled scales for columnwise.

zhongbozhu · 2026-07-27T23:36:50Z

+    )
+
+
+def prepare_prequantized_mxfp8_input_for_gemm(


This function needs to be heavily unit tested like here:

TransformerEngine/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py

Line 333 in 98524e5

@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)

zhongbozhu · 2026-07-27T23:37:35Z

+    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(


Does group dequantize work with already swizzled scales? Would be better if we put an assert here.

We have an assert above in line 177

zhongbozhu · 2026-07-27T23:38:51Z

+        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,


Why is colwise data None if we have if with_columnwise: above?

When we have colwise_quantizer.optimize_for_gemm = True, the colwise data should be already swizzled in the quantize kernel. We need to manually swizzle for rowwise data because it is generated from NCCL EP dispatch.

zhongbozhu · 2026-07-27T23:45:20Z

        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):


Overall I think it makes sense to limit the scope of mxfp8 alltoall to the TE fused ops module because it assumes using grouped kernels already so it's fine if we enforce shape alignment on top of it and we can safely ignore the zero padding of mxfp8 scaling factors. However, was NCCL-EP designed to work with TE fused ops only? Sounds like we should insert more assertions elsewhere if TE fused ops module is not used.

cc @timmoon10

NCCL EP itself is designed to be compatible with both fused and unfused path. We would like to make it general and in mcore I built many support for both paths. When mcore is ready for the grouped tensor path (which I think you are working on?), we can adopt NCCL EP to that path as well.

For mxfp8 specifically, the major connected interface is dispatch (combine takes bf16 input). TE EP dispatch will return GroupedTensor with only rowwise quantized, I think it should work with both paths?

cc @phu0ngng

vthumbe1503 · 2026-07-27T22:30:16Z

+    ) -> "GroupedTensor":
+        """Rowwise-only MXFP8 GroupedTensor with compact scales (FP8 dispatch wire format)."""
+        wire_quantizer = MXFP8Quantizer(
+            fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False


Recommendation is to use TE_DType instead of tex.DType when constructing quantizers now, could we please do that here?

vthumbe1503 · 2026-07-28T01:15:29Z

+                # 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)


There is already copy method defined in GroupedTensorStorage class to achieve this functionality of creating a shallow storage copy. Could we please use that and remove this method?

Suggested change

grouped_fc2_dy = grouped_storage_from_grouped_tensor(grad_output)

grouped_fc2_dy = grad_output.copy()

Same applies in all other places where this method is used,

vthumbe1503 · 2026-07-28T01:36:29Z

+        )
+    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


I think the better way should be to directly use the grouped_swizzle API which handles variable dims as well.

This would work only if all per expert dims are padded to 128.

vthumbe1503 · 2026-07-28T01:43:34Z

+    )
+
+
+def prepare_prequantized_mxfp8_input_for_gemm(


I am uncomfortable with the complexity of this function.

I think we should define two functions, one for the prepare_rowwise_mxfp8_input for forward pass and preparing_colwise_mxfp8_input for the backward pass.

To be honest I dont think prepare_rowwise_mxfp8_input is even needed, Simply calling grouped_swizzle API on the activation should be enough which can be done directly before the GEMM in grouped_linear.py and grouped_mlp.py. However, the perf of the variable swizzle kernel is not being well studied. So it would be good to evaluate that before making the change.

For colwise, I am assuming the rowwise dequant + quant is well studied in terms of convergence. And so having a sperate function call for that is ideal which also takes care of the dbias computation

vthumbe1503 · 2026-07-28T01:47:33Z

+        requires_grad=False,
+        with_gemm_swizzled_scales=False,
+    )
+    tex.swizzle_scales_for_gemm_(tmp)


@zhongbozhu I think the group_quantize + swizzle fusion generates the swizzled scales for columnwise.

add support for row-wise quanted input for grouped gemm

11e9007

Signed-off-by: YangFei1990 <feiw@nvidia.com>

YangFei1990 requested a review from phu0ngng July 23, 2026 18:28

phu0ngng reviewed Jul 23, 2026

View reviewed changes

phu0ngng requested a review from vthumbe1503 July 23, 2026 20:42

phu0ngng reviewed Jul 24, 2026

View reviewed changes

YangFei1990 added 2 commits July 25, 2026 12:05

doc change

d234cd8

Signed-off-by: YangFei1990 <feiw@nvidia.com>

implement for backward

376b94c

Signed-off-by: YangFei1990 <feiw@nvidia.com>

YangFei1990 marked this pull request as ready for review July 26, 2026 04:37

YangFei1990 requested a review from timmoon10 as a code owner July 26, 2026 04:37

greptile-apps Bot reviewed Jul 26, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/basic/grouped_linear.py Outdated

YangFei1990 added 2 commits July 25, 2026 21:49

merge from main

782835c

Signed-off-by: YangFei1990 <feiw@nvidia.com>

allow scaled_bias + frozen weights in prequant path

855f480

Signed-off-by: YangFei1990 <feiw@nvidia.com>

greptile-apps Bot reviewed Jul 26, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/basic/grouped_linear.py

allow bias + frozen weight path in prequantized input

8ef7e7e

Signed-off-by: YangFei1990 <feiw@nvidia.com>

zhongbozhu reviewed Jul 27, 2026

View reviewed changes

zhongbozhu suggested changes Jul 27, 2026

View reviewed changes

vthumbe1503 reviewed Jul 28, 2026

View reviewed changes

		and isinstance(input_quantizers[0], MXFP8Quantizer)
		and isinstance(input_.quantizer, MXFP8Quantizer)

	grouped_fc2_dy = grouped_storage_from_grouped_tensor(grad_output)
	grouped_fc2_dy = grad_output.copy()

		)


		def prepare_prequantized_mxfp8_input_for_gemm(

		)


		def prepare_prequantized_mxfp8_input_for_gemm(

Uh oh!

Conversation

YangFei1990 commented Jul 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Checklist:

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

greptile-apps Bot commented Jul 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 3/5

Important Files Changed

Flowchart

Uh oh!

Uh oh!

Uh oh!

YangFei1990 commented Jul 26, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

YangFei1990 commented Jul 23, 2026 •

edited

Loading

greptile-apps Bot commented Jul 26, 2026 •

edited

Loading