[PyTorch] Fix NaN expert weight gradients at num_groups == 1 with SReLU

[GarlGuo]

Description

At num_groups == 1 the fused grouped MLP may quantize and scale-swizzle the token buffer as one dense tensor over tensor.shape[0] rows and read it back the same way. That is only coherent if the cuDNN kernels also write their intermediates densely over tensor.shape[0], which is what use_single_group_runtime_offsets asks of them — they then derive M from the runtime tensor shapes rather than from the offsets.

TE selects the dense consumer without checking whether the producer will be dense. The two disagree for ScaledSReLU at any version, and for every activation on a cuDNN frontend older than 1.27.0, because _cudnn_frontend_supports_single_group_runtime_offsets returns False there and the kernels pack their output to sum(split_sizes) rows instead.

The columnwise swizzled MXFP8 scale layout is [k/128][m/128][32][4][4], with the m-tile as an inner stride, so an m-extent mismatch misindexes every k-tile past the first: consumer k-tile j reads producer k-tile j * T_c / T_p. The head rows — the ones that do belong to the group — pick up the wrong scale factors, and high k-tiles read scale memory the producer never wrote. The data is read correctly; only the scales move, so the error is order 1 and often NaN.

Crucially this is value-independent: zeroing the padded tail does not help, because the defect is in the indexing rather than in the padding's contents. A caller cannot work around it.

Callers may legally pass a token buffer padded past sum(split_sizes); an MoE expert-parallel fixed-capacity receive buffer always does, and test_grouped_linear_cuda_graph_safe documents the padded tail as "intentionally outside every group" and "uninitialized".

Fix to #3484

Minimal reproduction

Single GPU, no Megatron, no distributed. NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 must be set before the TE import.

import os
os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1"  # must precede the TE import
import torch
import transformer_engine.pytorch as te
from transformer_engine.pytorch.quantization import MXFP8BlockScaling

HIDDEN, LIVE, PAD = 1024, 512, 256


def wgrad(total_rows, fill):
    torch.manual_seed(0)
    fc1 = te.ops.GroupedLinear(1, HIDDEN, HIDDEN, bias=False, device="cuda",
                               dtype=torch.bfloat16, accumulate_into_main_grad=True)
    fc2 = te.ops.GroupedLinear(1, HIDDEN, HIDDEN, bias=False, device="cuda",
                               dtype=torch.bfloat16, accumulate_into_main_grad=True)
    model = te.ops.Sequential(fc1, te.ops.ScaledSReLU(), fc2)
    for op in (fc1, fc2):
        w = op.get_parameter("weight0")
        w.main_grad = torch.zeros(w.shape, dtype=torch.float32, device="cuda")

    gen = torch.Generator(device="cuda").manual_seed(1)
    live_x = torch.randn(LIVE, HIDDEN, generator=gen, device="cuda", dtype=torch.bfloat16)
    live_dy = torch.randn(LIVE, HIDDEN, generator=gen, device="cuda", dtype=torch.bfloat16)

    x = torch.full((total_rows, HIDDEN), fill, device="cuda", dtype=torch.bfloat16)
    dy = torch.full((total_rows, HIDDEN), fill, device="cuda", dtype=torch.bfloat16)
    x[:LIVE], dy[:LIVE] = live_x, live_dy
    x = x.detach().requires_grad_(True)

    # split_sizes covers the LIVE rows only; the rest of the buffer is padding.
    splits = torch.tensor([LIVE], device="cuda", dtype=torch.int32)
    probs = torch.ones(total_rows, device="cuda", dtype=torch.float32)

    with te.autocast(enabled=True, recipe=MXFP8BlockScaling()):
        out = model(x, splits, probs, splits)
    out.backward(dy)
    return fc1.get_parameter("weight0").main_grad.double()


reference = wgrad(LIVE, 0.0)                      # no padding
for label, fill in (("padding zeroed  ", 0.0), ("padding garbage ", 3.0)):
    got = wgrad(LIVE + PAD, fill)
    rel = ((got - reference).norm() / reference.norm()).item()
    bad = int((~torch.isfinite(got)).sum())
    print(f"{label} rel_err={rel:.4e}  non_finite={bad}/{got.numel()}")

On main:

padding zeroed   rel_err=nan  non_finite=411648/1048576
padding garbage  rel_err=nan  non_finite=411648/1048576

With this PR:

padding zeroed   rel_err=0.0000e+00  non_finite=0/1048576
padding garbage  rel_err=0.0000e+00  non_finite=0/1048576

Note the two main rows are identical: the corruption does not depend on what the padding holds, so a caller cannot avoid it by zeroing.

Why the fix is where it is

The two halves of the "treat this as one dense group" decision are currently gated differently. The request to the kernels is already guarded, because it can be refused:

if supports_single_group_runtime_offsets:
    fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1

but its siblings — use_single_group_weight_swizzle, use_single_discrete_weight, use_single_group_dense_fc2, use_single_group_dense_dgrad, the dense quantize and the plain wgrad GEMM — are gated on num_groups == 1 alone. So when the specialization is unavailable, TE still reads everything densely while the kernels pack.

This PR applies the guard TE already writes, at the sites that were missing it:

use_dense_single_group = num_groups == 1 and (
    _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
)

That predicate is called directly rather than wrapped, matching the three other call sites in this file, and the flag is named for the decision it drives rather than for a property of the data. It deliberately does not claim the single group spans the buffer — split_sizes is device-resident, so that is unknowable here — only that the producer and the consumer will agree on the row extent.

They have to move together: quantization layout and kernel indexing must agree, so enabling a subset pairs a dense consumer with a packed producer and yields NaN rather than a slowdown. Hence one flag threaded explicitly through the grouped-quantize helpers and _compute_grad_params.

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

Please list the changes introduced in this PR:

• Gate the num_groups == 1 dense shortcut family on _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)) as well as on the group count, so TE only reads operands densely when the cuDNN kernels also write them densely.
• Thread that one decision explicitly through _group_quantize_for_grouped_mlp, _group_quantize_with_amax_for_grouped_mlp and _compute_grad_params rather than recomputing it per site, so all of the shortcuts move together.

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

tests/pytorch/test_grouped_mlp.py is unmodified and at baseline: 6 failed, 34 passed, the 6 being pre-existing nvfp4_rht and cuBLAS-version failures also present on main. The fix is verified out-of-tree by the reproducer above; happy to add an in-tree regression test if you would like one in this PR.

[greptile-apps]

Greptile Summary

This PR makes the fused grouped-MLP single-group optimization conditional on cuDNN runtime-offset support, keeping quantization layouts and their GEMM consumers consistent.

• Threads one use_dense_single_group decision through forward and backward quantization, scale swizzling, FC2, dgrad, and wgrad paths.
• Falls back to grouped storage and offset handling for ScaledSReLU and older cuDNN frontends.
• Lacks an in-tree regression test for the padded single-group SReLU gradient failure.

Confidence Score: 4/5

The implementation appears safe to merge, with the non-blocking concern that its corrected padded single-group SReLU behavior is not protected by an automated regression test.

The shared predicate is consistently applied across the affected forward and backward layout paths, and no correctness regression was established; the remaining finding concerns future regression detection.

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

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Consistently gates single-group dense producer and consumer paths on cuDNN runtime-offset support; focused regression coverage is still missing.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[Single-group fused MLP] --> B{Runtime-offset specialization supported?}
  B -->|Yes| C[Dense quantization and scale layout]
  C --> D[Dense FC2, dgrad, and wgrad consumers]
  B -->|No| E[Grouped quantization using split sizes and offsets]
  E --> F[Generic grouped GEMM consumers]

Loading

_{Reviews (1): Last reviewed commit: "[PyTorch] Take the num_groups == 1 dense..." | Re-trigger Greptile}

[greptile-apps]

Regression coverage is missing

This capability gate fixes a specific numerical failure, but no automated test combines a padded single-group buffer, ScaledSReLU, MXFP8 or NVFP4, and backward gradient validation. The existing single-group MXFP8 test uses unpadded ScaledSwiGLU, while padded fused-MLP coverage uses GLU activations and four groups. Without a focused regression test, the dense-versus-packed layout mismatch and resulting NaN gradients could return without failing CI.

Knowledge Base Used:

• Verification and CI matrix
• PyTorch tensor and operator abstractions

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!