[Common][PyTorch] Add row-scaled NVFP4 support to grouped-tensor GEMM by cael-ling · Pull Request #3661 · NVIDIA/TransformerEngine

cael-ling · 2026-10-09T09:20:19Z

Description

Depends on the merge of #3649. This PR is stacked on top of it, please review only the incremental changes.

This PR does three things:

Produce swizzled scales in the cast. When optimize_for_gemm is set, the row-scaled grouped NVFP4 cast bakes the cuBLAS-swizzled scale layout directly (for both the rowwise and columnwise/transpose scales). No separate swizzle pass is needed.
Consume row-scaled operands in the GEMM. general_grouped_gemm_for_grouped_tensor neutralizes the per-row/per-column FP32 outer scale, runs the per-tensor grouped GEMM into an FP32 scratch, and applies the outer scale afterward, mirroring the dense general_gemm path.
Enable columnwise for row-scaled grouped tensors. The prior restriction forbidding columnwise usage for row-scaled grouped NVFP4 is lifted, and the columnwise amax is sized per column so wgrad can consume the transpose. The cast kernel already emitted the columnwise transpose, per-column amax, and swizzled scales.

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

Changes

common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh: add a WITH_GEMM_SWIZZLED_SCALES template path to the row-scaled grouped cast body, both kernels, and both host launchers; the scale-factor epilogue writes cuBLAS-swizzled indices (compact scale column count) for rowwise and columnwise when enabled.
pytorch/cpp_extensions/gemm.py: consume row-scaled operands in general_grouped_gemm_for_grouped_tensor via outer-scale neutralization + FP32 scratch + post-scale, covering discrete-out (wgrad) and grouped-out paths.
pytorch/csrc/quantizer.cpp and pytorch/tensor/storage/grouped_tensor_storage.py: allow columnwise usage for row-scaled grouped NVFP4, size the columnwise amax per column, and split it per-expert per-column.
tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py: add fprop (TN), CUDA-graph, and wgrad (NT) parity tests against a per-expert dense row-scaled GEMM over byte-identical operands.

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

Single-launch grouped row-scaled NVFP4 cast + fused amax, plus graph-safe amax/cast variants driven by device first_dims/tensor_offsets for capture-once/replay MoE. Wired into PyTorch (ep.nvfp4_group_row_scaled_cast) with C++ and PyTorch capture/replay tests. Signed-off-by: Cael Ling <caell@nvidia.com>

Restrict the FUSED_NVFP4_ROW_SCALED selection to shapes the grouped kernels accept (128-aligned last dim, <= 64 experts) so unsupported requests fall back to the per-expert path. Return early for empty routing, make the graph-safe binding's input contiguous, and reject with_rht/with_2d_quantization/with_amax_reduction on both row-scaled entries. Add regression tests for the four guarded paths. Signed-off-by: Cael Ling <caell@nvidia.com>

greptile-apps · 2026-10-09T09:34:40Z

Confidence Score: 3/5

[Medium impact] The changes reach grouped NVFP4 casting and GEMM workflows in PyTorch.

Fix the unsupported-GPU test gate and finish support for differing output widths and live-row counts before merging.

Findings

Tests run on unsupported GPUs ▶
Different output widths fail ▶

Summary

Adds row-scaled NVFP4 support to grouped-tensor GEMM.

Grouped GEMM now uses row-scaled NVFP4 inputs for expert products.
Grouped NVFP4 casts write GEMM-ready scale layouts when optimization is enabled.
Row-scaled grouped NVFP4 tensors now keep columnwise data and amax values per expert.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[Grouped NVFP4 cast] --> B[Packed values and swizzled block scales]
    A --> C[Row and column outer scales]
    B --> D[GEMM with neutral outer scales]
    D --> E[FP32 scratch]
    C --> F[Apply outer scales]
    E --> F
    F --> G[Copy or accumulate into output]

_{Reviews (12) · Last reviewed commit: "Merge branch 'main' into pr/nvfp4-row-sc..." · Reviewed by Greptile}

…uting The grouped row-scaled cast treats its input as 2D [rows, last_dim] indexed by split counts, so route higher-rank inputs and partial splits (split sum below the row count) to the per-expert path instead of the grouped kernel. Make first_dims and tensor_offsets contiguous in the graph-safe cast so a strided routing view is not misread. Add fallback and strided-routing tests. Signed-off-by: Cael Ling <caell@nvidia.com>

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

greptile-apps · 2026-10-09T12:24:53Z

+        last_dim = out.logical_shape[-1]
+        flat = scratch.rowwise_data.view(-1, last_dim)
+        if post_scale_rows is not None:
+            # Per output row; handles uneven per-expert row counts directly.
+            flat.mul_(post_scale_rows.view(-1, 1))


Different output widths fail

The post-scale step treats the entire grouped output as rows with one common width. For listed weights shaped [128, 128] and [256, 128], with two row-scaled activation groups of 128 rows each, the output has logical_shape [128, 384]. That makes flat have 128 rows, while post_scale_rows has 256 values. The multiply raises after GEMM, even though GEMM supports these per-group shapes.

The same assumption fails when an activation's capacity exceeds the live output row count. Apply scales using each group's actual shape and offsets instead of multiplying the whole buffer.

The bulk NVFP4 split_quantize allocation sizes its amax buffer from the first quantizer's row-scale setting, so a group mixing row-scaled and tensor-scaled quantizers writes per-row amax into a per-tensor buffer and corrupts neighbors. Detect the mixed case and fall back to the unfused per-split path. Signed-off-by: Cael Ling <caell@nvidia.com>

Bake the cuBLAS-swizzled scale layout directly into the row-scaled grouped NVFP4 cast when optimize_for_gemm is set, mirroring the per-tensor and RHT cast-fusion paths. The grouped-tensor NVFP4 GEMM requires swizzled scales, and the standalone grouped swizzle supports only MXFP8, so the non-RHT row-scaled path had no swizzle producer. Consume row-scaled operands in general_grouped_gemm_for_grouped_tensor by neutralizing the per-row/per-column outer scale, running the per-tensor grouped GEMM into an FP32 scratch, and applying the outer scale afterward. This covers both the rowwise direction (fprop) and the columnwise direction (wgrad). Allow columnwise usage for row-scaled grouped NVFP4 tensors, lifting the prior restriction, and size the columnwise amax per column so wgrad can consume the transpose. The cast kernel already produced the columnwise transpose, amax, and swizzled scales. Add GEMM-level fprop, CUDA-graph, and wgrad parity tests against a per-expert dense row-scaled GEMM over byte-identical operands. Signed-off-by: Cael Ling <caell@nvidia.com>

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

Reject explicit beta and listed row-scaled operands for the row-scaled NVFP4 grouped path instead of silently mishandling them. Apply the discrete-out post-scale with per-expert row/column offsets so uneven group sizes are handled, and guard the grouped-out column post-scale on a uniform per-expert output row count. Add a discrete-out wgrad test. Signed-off-by: Cael Ling <caell@nvidia.com>

The grouped-out column post-scale split the output into equal-sized groups, which is wrong for uneven per-expert row counts and fails outright when the total is not divisible. Apply each group's column scale over its actual rows using host-side shapes, falling back to device first_dims only for varying- dimension tensors (never built under graph capture). Add an uneven TN test. Signed-off-by: Cael Ling <caell@nvidia.com>

The grouped single-buffer output path read per-expert row counts via first_dims.tolist(), which forces a device-to-host sync and breaks CUDA graph capture (the MXFP8 EP dispatch builds first_dims-backed outputs under capture). Expand the per-expert column scales with a device-side repeat_interleave whose output_size avoids the sync. Also reject a varying per-expert output width, which the shared-width expansion cannot represent. Signed-off-by: Cael Ling <caell@nvidia.com>

Restrict the grouped single-buffer post-scale to the live row count from logical_shape so a capacity-padded output does not misalign the per-row and per-group scales, and keep rejecting a varying per-expert output width (grouped GEMM outputs share the non-contracted dimension, so no TE caller produces one). Raise the test skip threshold to cuBLAS 13.4, matching CUBLAS_NVFP4_GROUPED_GEMM_VERSION, so the tests skip instead of failing the GEMM version check on 13.3. Signed-off-by: Cael Ling <caell@nvidia.com>

Add the kNVTEGroupedRowScaledNVFP4 grouped tensor param and dispatch row-scaled NVFP4 outputs in nvte_group_quantize to the graph-safe amax and cast kernels. Handle row-scaled quantizers in tex.group_quantize. Remove the row-scaled-specific C APIs, the host-split row-scaled cast kernel, tex.nvfp4_group_row_scaled_cast_graph_safe, the ep.py helper, and the fused row-scaled split_quantize path. Add a top-level static_assert for unregistered types in TypeInfo. Signed-off-by: Cael Ling <caell@nvidia.com>

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

Signed-off-by: Cael Ling <caell@nvidia.com>

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

…led-grouped-tensor-gemm Signed-off-by: Cael Ling <caell@nvidia.com> # Conflicts: # transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh

Signed-off-by: Cael Ling <caell@nvidia.com>

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

Dispatch nvte_group_amax and nvte_group_amax_graph_safe on the output tensor format so that row-scaled NVFP4 outputs get per-row amaxes, and deprecate nvte_group_nvfp4_compute_amax in favor of nvte_group_amax. Rename the row-scaled amax helpers to group_row_amax(_graph_safe). In split_quantize, fall back to unfused allocation and quantization whenever any NVFP4 quantizer is row-scaled. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Tim Moon <tmoon@nvidia.com>

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

@greptile-apps

Review suggestion from @greptile-apps Signed-off-by: Tim Moon <tmoon@nvidia.com>

Signed-off-by: Tim Moon <tmoon@nvidia.com>

…led-grouped-tensor-gemm Signed-off-by: Cael Ling <caell@nvidia.com>

Signed-off-by: Cael Ling <caell@nvidia.com>

timmoon10

Functionally this looks reasonable, although I have minor complaints about bad habits in AI-written code.

timmoon10 · 2026-11-05T20:04:29Z

+                        camax_end = nvfp4_columnwise_amax_offsets[i + 1]
+                        amax_columnwise = self.columnwise_amax[camax_start:camax_end]
+                    else:
+                        amax_columnwise = self.columnwise_amax[i : i + 1]


With the change in https://github.com/NVIDIA/TransformerEngine/pull/3661/changes#r4485369977, we can make this general and not have any logic specific to row-scaled NVFP4:

Suggested change

amax_columnwise = self.columnwise_amax[i : i + 1]

if self.amax is not None:

amax_start, amax_end = nvfp4_rowwise_amax_offsets[i:i+2]

amax_rowwise = self.amax[amax_start:amax_end]

if self.columnwise_amax is not None:

amax_start, amax_end = nvfp4_columnwise_amax_offsets[i:i+2]

amax_columnwise = self.columnwise_amax[amax_start:amax_end]

Done. The slicing is one path now, with no row-scaled special case.

- Move row-scaled handling out of general_grouped_gemm_for_grouped_tensor into a dedicated _grouped_gemm_row_scaled_nvfp4 helper. - Compute amax offsets for all NVFP4 grouped tensors and slice amax through a single code path. - Index Vec elements directly instead of reinterpret_cast. - Rename the grouped-tensor GEMM test gate and split skip reasons. - Check that first_dims is a CUDA tensor with num_tensors entries. Signed-off-by: Cael Ling <caell@nvidia.com>

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

cael-ling · 2026-11-06T14:07:24Z

Hello @timmoon10, Thanks for the review! Addressed all your comments in 67e7ea6 and replied inline. All related tests pass on GB200.

timmoon10 · 2026-11-07T00:51:59Z

/te-ci pytorch

timmoon10

LGTM, pending CI

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 · 2026-11-07T02:31:26Z

/te-ci pytorch

greptile-apps · 2026-11-07T02:34:14Z

+if torch.cuda.get_device_capability() < (10, 0):
+    _grouped_tensor_gemm_supported = False
+    _reason_for_no_grouped_tensor_gemm = "NVFP4 grouped-tensor GEMM requires Blackwell+."


Tests run on unsupported GPUs

The new gate accepts SM120 and SM121 with cuBLAS 13.4+, but grouped GEMM explicitly rejects GPUs above SM110. recipe_available also accepts these GPUs, so the new tests run and fail instead of skipping. Restrict this gate to the GPU range supported by grouped GEMM.

cael-ling added 2 commits October 7, 2026 23:57

cael-ling requested review from Oleg-Goncharov, ksivaman and ptrendx as code owners October 9, 2026 09:20

github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Oct 9, 2026

greptile-apps Bot reviewed Oct 9, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/cpp_extensions/gemm.py Resolved

cael-ling and others added 2 commits October 9, 2026 05:02

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

07de769

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

cael-ling force-pushed the pr/nvfp4-row-scaled-grouped-tensor-gemm branch from dc305f7 to 272fad7 Compare October 9, 2026 12:15

greptile-apps Bot reviewed Oct 9, 2026

View reviewed changes

cael-ling and others added 6 commits October 9, 2026 05:49

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

26ab746

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

cael-ling force-pushed the pr/nvfp4-row-scaled-grouped-tensor-gemm branch from 272fad7 to 07ceda5 Compare October 9, 2026 13:11

greptile-apps Bot reviewed Oct 9, 2026

View reviewed changes

Comment thread tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py Resolved

cael-ling mentioned this pull request Oct 13, 2026

Add NVFP4 per-token quantization recipe #3045

Open

1 of 13 tasks

timmoon10 self-requested a review October 14, 2026 03:01

cael-ling and others added 6 commits October 21, 2026 22:07

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

3076ff7

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

Require 128-row experts for uniform row-scaled NVFP4 group_quantize

de54786

Signed-off-by: Cael Ling <caell@nvidia.com>

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

3d656d2

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

Merge branch 'pr/nvfp4-row-scaled-grouped-cast' into pr/nvfp4-row-sca…

20c73bb

…led-grouped-tensor-gemm Signed-off-by: Cael Ling <caell@nvidia.com> # Conflicts: # transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh

Use tex.group_quantize for row-scaled grouped GEMM tests

1d9a318

Signed-off-by: Cael Ling <caell@nvidia.com>

cael-ling and others added 2 commits October 22, 2026 04:02

Trim comments in row-scaled grouped GEMM

035a322

Signed-off-by: Cael Ling <caell@nvidia.com>

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

055277e

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

greptile-apps Bot reviewed Oct 22, 2026

View reviewed changes

Comment thread transformer_engine/pytorch/csrc/extensions/cast.cpp Resolved

timmoon10 and others added 6 commits October 26, 2026 15:23

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

55f69ca

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

Fix missing include

32c57f2

Review suggestion from @greptile-apps Signed-off-by: Tim Moon <tmoon@nvidia.com>

Add missing header include

07f45a1

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Merge branch 'pr/nvfp4-row-scaled-grouped-cast' into pr/nvfp4-row-sca…

0682b76

…led-grouped-tensor-gemm Signed-off-by: Cael Ling <caell@nvidia.com>

Merge branch 'main' into pr/nvfp4-row-scaled-grouped-tensor-gemm

e55648e

Signed-off-by: Cael Ling <caell@nvidia.com>

cael-ling mentioned this pull request Oct 29, 2026

[PyTorch][common] Support row-scaled NVFP4 in GroupedLinear #3776

Draft

1 of 7 tasks

timmoon10 reviewed Nov 5, 2026

View reviewed changes

cael-ling and others added 2 commits November 6, 2026 05:36

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

f3aedc9

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

cael-ling requested a review from timmoon10 November 6, 2026 14:07

timmoon10 previously approved these changes Nov 7, 2026

View reviewed changes

Merge branch 'main' into pr/nvfp4-row-scaled-grouped-tensor-gemm

adfbe2d

Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

timmoon10 dismissed their stale review via adfbe2d November 7, 2026 02:30

greptile-apps Bot reviewed Nov 7, 2026

View reviewed changes

-                        amax_columnwise = self.columnwise_amax[i : i + 1]
+                if self.amax is not None:
+                    amax_start, amax_end = nvfp4_rowwise_amax_offsets[i:i+2]
+                    amax_rowwise = self.amax[amax_start:amax_end]
+                if self.columnwise_amax is not None:
+                    amax_start, amax_end = nvfp4_columnwise_amax_offsets[i:i+2]
+                    amax_columnwise = self.columnwise_amax[amax_start:amax_end]

Conversation

cael-ling commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

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

Uh oh!

Confidence Score: 3/5

Findings

Summary

Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

greptile-apps Bot Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

timmoon10 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

timmoon10 Nov 5, 2026

Choose a reason for hiding this comment

Uh oh!

cael-ling Nov 6, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cael-ling commented Nov 6, 2026

Uh oh!

timmoon10 commented Nov 7, 2026

Uh oh!

timmoon10 left a comment

Choose a reason for hiding this comment

Uh oh!

timmoon10 commented Nov 7, 2026

Uh oh!

greptile-apps Bot Nov 7, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

cael-ling commented Oct 9, 2026 •

edited

Loading

greptile-apps Bot commented Oct 9, 2026 •

edited

Loading

timmoon10 left a comment •

edited

Loading