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

[YangFei1990]

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

[phu0ngng]

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

[YangFei1990]

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]

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

[phu0ngng]

Should we validate grouped_x.rowwise_scale as well?

[phu0ngng]

compact -> unswizzled.

[phu0ngng]

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]

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

[YangFei1990]

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?

[greptile-apps]

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]

Loading

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

[YangFei1990]

/te-ci pytorch