Use cuDNN for row-scaled NVFP4 grouped GEMM

[zianglih]

Summary

@HumansAnd

• route row-scaled NVFP4 TN forward GEMMs through cuDNN Frontend grouped GEMM quant on SM100 for both general_gemm (one group) and general_grouped_gemm, keeping their numerical backend unified
• retain the existing per-GEMM emulation on SM107 because the Rubin cuDNN Frontend kernel explicitly does not support row-scale fusion
• support compact and pre-swizzled weight scales, FP32/BF16 inputs and outputs, bias, 4over6, and small or irregular logical shapes
• pad each logical M split to 256 only inside the cuDNN adapter, unpad its output, and keep the public NVFP4 quantization alignment at 128

Required cuDNN Frontend dependencies

This integration requires both:

• Add row-scale support to grouped GEMM quant cudnn-frontend#251, which adds the row_scale_tensor grouped GEMM API and is available in cuDNN Frontend v1.25.0 and later.
• Support FP32 output and dynamic M in row-scaled FP4 grouped GEMM cudnn-frontend#461, which enables FP32 output for FP4 with 16-element scale vectors and makes discrete row scaling dynamic in M so cached kernels support changing token counts. It also adds exact K=64/N=256 production-path coverage. This PR was merged into cuDNN Frontend develop at ab9efe15.

The integration therefore requires a cuDNN Frontend build containing NVIDIA/cudnn-frontend#461. This PR remains draft until a released nvidia-cudnn-frontend version consumed by TransformerEngine includes it.

Motivation

This integration supports the row-scaled NVFP4 work from #2931 while avoiding different numerical backends for dense and grouped forward GEMMs on SM100.

Validation

• rebased onto upstream main at dd965956; current head b6e69356
• git diff --check upstream/main...HEAD
• pre-commit run --all-files
• fresh --queue hell B200 devbox, CUDA 12.9, PyTorch 2.9.1, cuDNN Frontend 1.27.0 built from Support FP32 output and dynamic M in row-scaled FP4 grouped GEMM cudnn-frontend#461 merge commit ab9efe15:
  • exact grouped-versus-singleton Cartesian matrix, including compact/pre-swizzled scales, FP32/BF16 operands and outputs, bias, list/single outputs, 4over6 modes, and small/irregular shapes: 2816 passed
  • existing row-scaled NVFP4 padded/unpadded GroupedLinear bitwise matrix: 16 passed, 8 skipped by existing unsupported-configuration guards
  • row-scaled NVFP4 BF16 Linear forward/backward sanity: 1 passed

[zianglih]

need to rebase and refactor according to recent cute dsl integration

[zianglih]

Rebase/refactor milestone complete.

• rebased the two existing commits onto current main (6377ca16) and force-pushed with a lease
• added separate signed-off refactor commit 96ae951f
• replaced the conflict-era setup with the current fused-MoE conventions: shared discrete-weight packing helper, materialized cuDNN SFA layout, dynamic scheduler, separate weight alpha_tensor and activation row_scale_tensor, and no fallback
• final B200 focused suite: 8 passed; adjacent TE contracts: 2 passed; existing upstream MXFP8 cuDNN control: 1 passed; targeted production-file error lint: passed

The draft remains blocked on the cuDNN Frontend side. NVIDIA/cudnn-frontend#251 provides row_scale_tensor, but current v1.26.0/main still needs explicit Int32 conversions in the MoE scheduler with its pinned CUTLASS DSL 4.5.x. The TE branch deliberately does not carry that workaround.

[zianglih]

Full PR diff audit is complete in 6a71ed3e (separate commit).

• Reduced the PR from 301 additions / 83 deletions to 287 additions / 105 deletions, lowering the net delta by 36 lines.
• Simplified the cuDNN SFA materialization and global-scale construction, reused TE's cached FP32 ones tensor, removed redundant test parameters, and moved row-scaled dispatch ahead of unused cuBLAS setup.
• Kept the adapter aligned with the current fused grouped-MLP conventions: direct cuDNN API use, shared discrete-weight pointer/scale packing, separate alpha_tensor and row_scale_tensor, and no fallback.
• Tightened grouped-vs-per-GEMM comparison to TE's canonical NVFP4 tolerance.

Validation on the exact committed sources:

• focused row-scaled suite: 8/8 passed
• adjacent existing controls: 3/3 passed
• uneven-split Fp8Padding -> GroupedLinear -> Fp8Unpadding no-grad integration: passed with one cuDNN call and exact Int32 padded offsets
• py_compile, git diff --check, pre-commit run --all-files, and targeted production-module pylint: passed

The temporary installed-package scheduler typing patch used for B200 validation was restored afterward. The TE branch still contains no workaround and explicitly depends on the cuDNN Frontend row-scale feature from NVIDIA/cudnn-frontend#251 plus the scheduler correction.

[zianglih]

Rebase and current B200 blocker audit are complete.

• Force-with-lease updated the branch to f2b3276e on upstream main (5a9bcd98); range-diff shows all four PR patches are unchanged.
• On a fresh --queue hell B200, the historical dynamic-scheduler compiler failure did not reproduce. Both cuDNN Frontend 1.25.0 and 1.26.0 passed all 8 focused cases with their pinned CUTLASS DSL 4.5.0, using separate empty compiler caches and no scheduler patch.
• The full TE sanity suite reported 17,122 passed / 33,922 skipped / 144 failed. All 144 failures are row-scaled GroupedLinear small-model cases hitting the PR's explicit K/N alignment rejection at K=64; the matrix also contains FP32 outputs outside this cuDNN path.

The PR remains draft. The old external compiler blocker is no longer reproduced, but the no-fallback support contract still needs an explicit compatibility decision. I did not restore fallback or hide the regression with test skips.

[zianglih]

wait for cudnn-fe release:

• Support FP32 output and dynamic M in row-scaled FP4 grouped GEMM cudnn-frontend#461