[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: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 for grouped GEMM.

Reuses dispatched rowwise data while reconstructing columnwise storage needed for weight gradients.
Adds native grouped dequantize/requantize and scale-swizzle bindings.
Extends grouped-linear, fused grouped-MLP, and MXFP8 edge-case tests.

Confidence Score: 4/5

The PR is not yet safe to merge because fused grouped-MLP forward still asserts when expert weights are frozen but another input requires gradients.

The frozen-weight fused path creates only rowwise GEMM storage, while forward context preparation still passes that tensor to a helper that requires columnwise storage; the corresponding test configuration is skipped.

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

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/basic/grouped_linear.py	Integrates rowwise-prequantized MXFP8 inputs and gradients into grouped-linear forward and backward paths.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Extends the fused grouped MLP to consume rowwise-prequantized MXFP8 activations and gradients.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Implements grouped dequantization, columnwise requantization, and rowwise-scale swizzling for prequantized tensors.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp	Propagates per-group dimensions through MXFP8 grouped-scale swizzling.
tests/pytorch/test_grouped_mlp.py	Adds grouped-linear and fused-MLP comparisons between prequantized MXFP8 tensors and dequantized references.

Sequence Diagram

sequenceDiagram
    participant Dispatch as FP8 token dispatch
    participant Op as Grouped Linear / MLP
    participant Native as TE native bindings
    participant GEMM as Grouped GEMM
    Dispatch->>Op: Rowwise-only MXFP8 GroupedTensor
    Op->>Native: Swizzle rowwise scales
    alt Weight gradients required
        Native->>Native: Dequantize and columnwise-requantize
    end
    Op->>GEMM: GEMM-ready grouped tensor
    GEMM-->>Op: Forward or backward result

_{Reviews (8): Last reviewed commit: "Merge branch 'main' into mxfp8_input_gro..." | 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.

what you are doing is essentially this right?

group_quantize_colwise_only_swizzle_fused( group_dequantize( mxfp8_input_compact_scales ) )

It would be better if we have a fused kernel for this process, is this already planned?

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

sounds good

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.

Thanks, would be better if this kernel is also fused.

Nit: In case we don't have the kernel yet, I would still prefer having a tex. API calling unfused kernels in C++, and replace its implementation with a fused kernel when it's ready.

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.

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

for more information, see https://pre-commit.ci

greptile-apps · 2026-07-30T23:45:13Z

Want your agent to iterate on Greptile's feedback? Try greploops.

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

for more information, see https://pre-commit.ci

YangFei1990 · 2026-07-31T05:47:56Z

/te-ci L1 pytorch

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 Outdated

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

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py

vthumbe1503 reviewed Jul 28, 2026

View reviewed changes

YangFei1990 added 3 commits July 29, 2026 20:32

use .copy to create group tensor

6015c91

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

merge from main

49ba134

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

move the implementation to c++ layer

12f1dbf

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

YangFei1990 requested a review from ksivaman as a code owner July 30, 2026 23:41

YangFei1990 and others added 2 commits July 30, 2026 16:41

Merge branch 'main' into mxfp8_input_groupedgemm

14fedf0

[pre-commit.ci] auto fixes from pre-commit.com hooks

e3990dd

for more information, see https://pre-commit.ci

bug fix for operation ordering

c1e4663

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

greptile-apps Bot reviewed Jul 31, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/ops/fused/grouped_mlp.py

YangFei1990 and others added 3 commits July 30, 2026 22:11

add comprehensive tests

a349d0d

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

[pre-commit.ci] auto fixes from pre-commit.com hooks

a3219de

for more information, see https://pre-commit.ci

Merge branch 'main' into mxfp8_input_groupedgemm

7c850ef

	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!

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: 4/5

Important Files Changed

Sequence Diagram

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!

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!

greptile-apps Bot commented Jul 30, 2026

Uh oh!

Uh oh!

YangFei1990 commented Jul 31, 2026

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