[tingyangk]

Description

Adds an unquantized (BF16) fused path for the MoE grouped MLP GroupedLinear → ScaledSwiGLU/ScaledClampedQGeGLU → GroupedLinear. Until now NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 fused only MXFP8/NVFP4. Without a block-scaled recipe the MLP ran as three separate ops.

The implementation uses cuDNN Frontend CuTeDSL grouped-GEMM kernels for forward, dgrad/dGLU, and wgrad on Blackwell. It requires cuDNN Frontend ≥ 1.32.0 with canonical BF16 layout support.

Thanks @hwanseoc for the help on N-major dgrad weights and direct weight-gradient allocation optimizations (hwanseoc#4).

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

Benchmarks

Unfused vs. fused, Qwen3.5-397B-A17B per-GPU shape

Benchmarked on one B300 using the per-GPU workload for EP32: 16 local experts, hidden size 4096, expert FFN size 1024 with SwiGLU, and 2,560 rows per expert (40,960 tokens total).

The benchmark uses BF16 eager execution with TE from this branch, cuDNN Frontend 1.32.0 (develop at 87531690), and NGC PyTorch 26.09. Times are CUDA-event GPU times: the median of 7 trials × 50 iterations, with 3 interleaved rounds per configuration.

Weight layout	Path	Forward (ms)	Backward (ms)	Train step (ms)	Step TFLOP/s	Step speedup
Per-expert	Unfused	0.889	1.683	2.556	1,210	—
Per-expert	Fused	0.820	1.580	2.400	1,288	1.065×
Packed	Unfused	0.881	1.677	2.557	1,209	—
Packed	Fused	0.820	1.574	2.393	1,292	1.068×

At the target shape, fusion improves forward performance by 1.075–1.084× and backward performance by approximately 1.065×. End-to-end training-step throughput improves by 1.065–1.068×. Packed and per-expert weights show nearly identical fused performance.

Token-load sweep

The following sweep uses per-expert weights and measures the complete training step:

Rows per expert	Total tokens	Unfused (ms)	Fused (ms)	Speedup
2,560	40,960	2.567	2.428	1.057×
4,096	65,536	3.992	3.753	1.064×
8,192	131,072	7.872	7.383	1.066×

The fused speedup remains stable as token load increases, reaching 1.066× at 8,192 rows per expert.

Testing

Added coverage for:

• Fusion and capability matching
• SwiGLU and GeGLU forward/backward
• FC2 scale_bias and separate scale gradients
• Packed, discrete, and distributed weights
• Main-gradient accumulation and delayed wgrad
• FP16/BF16 autocast behavior
• Zero-token batches and wgrad semantics
• Strided FP32 activation probabilities and FC2 scales
• Determinism rejection for atomic scale gradients

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

[greptile-apps]

Confidence Score: 3/5

[Medium impact] The change affects BF16 grouped MoE MLP training on a limited set of supported systems.

The PR still needs to preserve supported input and routing-scale shapes before merging.

Findings

1.  Batch dimensions disappear ▶
2.  Routing gradients have wrong shape ▶

Summary

Adds a BF16 fused grouped MLP path for Blackwell using cuDNN Frontend canonical layouts.

• Blackwell MoE layers can fuse BF16 projections with SwiGLU or supported GeGLU.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[GroupedLinear + scaled GLU + GroupedLinear] --> B{BF16 fusion supported?}
    B -->|No| C[Unfused operations]
    B -->|Yes| D[FC1 + GLU]
    D --> E[FC2]
    E --> F[Backward: dGLU and input gradients]
    F --> G[Shared weight-gradient helper]

Loading

_{Reviews (3) · Last reviewed commit: "Review feedback fixes" · Reviewed by Greptile}

[greptile-apps]

Batch dimensions disappear

The BF16 path flattens input_ before recording in_shape. A supported input shaped [2, 256, hidden] therefore returns [512, out_features] instead of [2, 256, out_features]. Backward also returns a flat input gradient, which can fail because its shape does not match the input. The unfused GroupedLinear preserves the leading dimensions.

Save the original shape before flattening, and restore it for both the output and input gradient.

Knowledge Base Used: PyTorch tensor and operator abstractions

[tingyangk]

Guards to input/scale tensors are added.

[greptile-apps]

Routing gradients have wrong shape

The BF16 path flattens the activation’s scales and returns grad_scales without restoring their original shape. For a [512, hidden] input with trainable scales shaped [2, 256], forward succeeds but backward returns a [512] scale gradient, which autograd rejects. The unfused activation accepts this scales shape and preserves it in the gradient.

Save the scales’ original shape and reshape their gradient before returning it.

Knowledge Base Used: PyTorch tensor and operator abstractions

[tingyangk]

Guards to input/scale tensors are added.

[timmoon10]

Overall LGTM, pending CI. It may be worth dealing with Greptile's comments by erroring out if the input tensors is not 2D and the scale tensor is not 1D.

[timmoon10]

/te-ci pytorch