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

[hwanseoc]

Description

The backward counterpart of NVIDIA#3641: pass operands to the cuDNN frontend grouped dGLU, dgrad (quant), and wgrad wrappers in their natural layouts in the fused MXFP8 grouped MLP backward.

It needs a cuDNN frontend with NVIDIA/cudnn-frontend#1422 (merged) and is stacked on NVIDIA#3641. This fork PR targets the NVIDIA#3641 branch, so the diff shows only the backward change.

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

• dGLU: when the dGLU wrapper 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.

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

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

• Unchanged: NVFP4, SReLU, and older frontends keep the existing path.

• Test: test_grouped_mlp_canonical_layouts_match_legacy now also checks the layout of every operand in the dGLU and FC1 dgrad calls, with bitwise-identical results. The fused op docstring notes the backward.

Measurements

Forward + backward host p50 per step, in µs, with the cuDNN frontend from NVIDIA/cudnn-frontend#1422:

Configuration	Host p50
NVIDIA#3641 (canonical forward only)	1419
This PR	1329

Setup:

• 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).
• Method: SM100, 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 extended test, this matches main. 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 [PyTorch] Pass canonical layouts to cuDNN grouped GEMM in fused grouped MLP NVIDIA/TransformerEngine#3641.

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