[PyTorch] Pass canonical layouts to cuDNN grouped GEMM in fused grouped MLP

[hwanseoc]

Description

Reduce host overhead in the fused MXFP8 grouped MLP (GroupedMLP_CuTeGEMMGLU), forward and backward, by passing operands to the cuDNN frontend grouped GEMM wrappers in their natural layouts.

cuDNN frontend develop lets the grouped GEMM wrappers take A as (sum(m), k), B as (num_groups, n, k), 1-D probabilities and the packed scale buffers as-is, and return 2-D outputs: NVIDIA/cudnn-frontend#1233 for the forward GLU and quant wrappers, NVIDIA/cudnn-frontend#1422 for the backward dGLU, dgrad (quant) and wgrad wrappers. Today TE builds 3-D/6-D views of every operand on each call, and the frontend maps them back.

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

• Forward: when the GLU and quant wrappers advertise supports_canonical_layouts, the MXFP8 forward passes canonical A/B/probability/scale operands and consumes the 2-D C/D and packed SFD outputs directly, without per-call view/permute chains.

• FC2 amax: FC2 skips its unused output amax when the quant wrapper advertises supports_optional_amax.

• Backward dGLU: when the dGLU wrapper also advertises supports_canonical_layouts, the MXFP8 backward passes:

  • 2-D grad output and activation input;
  • 1-D probabilities and dprob;
  • packed scale buffers as-is;
  • the FC2 column-wise weight as (num_groups, k, n) with b_major="n".

  The 2-D dGLU outputs are used directly.

• Backward FC1 dgrad: uses the 2-D dGLU outputs, the FC1 column-wise weight with b_major="n", and a 2-D grad input.

• Backward wgrad: passes flat scale buffers when the wgrad wrapper supports them.

• Unchanged: NVFP4, SReLU, and frontends without these attributes keep the existing path. The forward and backward are gated separately, so a frontend with only Reduce grouped GLU/quant host overhead and route grouped SwiGLU through unified GLU cudnn-frontend#1233 still gets the forward change.

• Test: a new test, test_grouped_mlp_canonical_layouts_match_legacy, checks that every operand of the FC1 GLU, FC2 quant, FC2 dGLU and FC1 dgrad calls reaches the frontend in its natural layout, and that outputs and gradients are bitwise identical to the legacy views.

• Docs: the fused op docstring notes the layout behavior.

Measurements

Host p50 in µs, median of three interleaved rounds, on this branch rebased on main at 459dedc.

Per forward:

Configuration	Host p50	End-to-end p50
main	458.0	470.4
This PR	360.7	374.2

Per forward + backward step:

Configuration	Host p50
main	1500
Canonical forward only	1416
This PR (forward and backward)	1346

Setup:

• Hardware and stack: SM100, TE built from this branch (sm_100), cuDNN frontend with Reduce grouped GLU/quant host overhead and route grouped SwiGLU through unified GLU cudnn-frontend#1233 and Accept canonical layouts in grouped dGLU, dgrad quant, and wgrad cudnn-frontend#1422, CuTeDSL 4.7.1, torch 2.14.0a0.
• Workload: an isolated 4096-row, 8-expert MLP with hidden size 2048 and intermediate size 512, MXFP8 weights, BF16 input and probabilities, ScaledSwiGLU(glu_interleave_size=32), gradients enabled.
• Method: pinned CPU, 300 synchronized samples per process, configurations interleaved over three rounds.

Testing

• tests/pytorch/test_grouped_mlp.py with NVTE_CUTEDSL_FUSED_GROUPED_MLP=1, rebased on main at 459dedc: 555 passed, 4463 skipped, 2 failed. Apart from the new test, this matches main on the same build. The two failures, test_single_grouped_weight_eval_preserves_columnwise_usage and test_dprob_is_bit_exact_across_runs[scaled_swiglu-GroupedMLP_CuTeGEMMGLU], also fail on main; the second needs a newer cuDNN frontend than the one used here.
• The forward output, input and probability gradients, and both weight gradients are bitwise identical to main.

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

@coderabbitai ignore

[zhongbozhu]

This will help resolve this issue: #2897

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no blocking issue was found in the latest changes.

Summary

The PR passes natural MXFP8 layouts to supporting cuDNN wrappers while keeping the legacy path for older wrappers.

• The fused MXFP8 grouped MLP sends natural tensor layouts to cuDNN.

_{Reviews (4) · Last reviewed commit: "[PyTorch] Keep the typed cuDNN weight vi..." · Reviewed by Greptile}

[timmoon10]

LGTM, pending CI

[timmoon10]

/te-ci pytorch

[zhongbozhu]

@hwanseoc @timmoon10 is it possible to get rid of this torch.view as well? To further reduce CPU overhead.

[hwanseoc]

Done for the weight views in 76581d8. The typed (groups, n, k) view of each weight buffer is now built once and kept on the grouped weight. It is rebuilt only if the buffer's data pointer changes. This covers all four sites: FC1 and FC2 in the forward, and the FC2 dGLU and FC1 dgrad in the backward. Each site drops from two torch ops (1.6 µs) to 0.24 µs.

The remaining view(dtype=...) calls are on buffers that are new every step: the quantized activations and the freshly swizzled weight scales. Each costs about 0.6 µs. TE stores quantized data as uint8 and cuDNN takes the FP8/FP4 dtype, so dropping those needs one of two changes:

• the TE quantizers allocate typed storage, or
• the cuDNN frontend accepts raw uint8 MXFP8 buffers with a dtype hint.

I can do the cuDNN side as a follow-up.