Skip to content

Use cuDNN for row-scaled NVFP4 grouped GEMM#3042

Draft
zianglih wants to merge 4 commits into
NVIDIA:mainfrom
zianglih:codex/cudnn-row-scale-nvfp4-grouped-gemm
Draft

Use cuDNN for row-scaled NVFP4 grouped GEMM#3042
zianglih wants to merge 4 commits into
NVIDIA:mainfrom
zianglih:codex/cudnn-row-scale-nvfp4-grouped-gemm

Conversation

@zianglih

@zianglih zianglih commented May 26, 2026

Copy link
Copy Markdown
Contributor

Summary

  • route row-scaled NVFP4 general_grouped_gemm forward calls directly through cuDNN Frontend's grouped GEMM quant kernel
  • fit the adapter to the current fused grouped-MLP structure by reusing TE's discrete-weight scale/pointer packing helper and cuDNN tensor-layout conventions
  • keep the tensor-scaled weight factor in alpha_tensor and the per-row activation factors in row_scale_tensor; cuDNN combines them in the accumulator epilogue
  • align row-scaled NVFP4 MoE splits to 256 rows and reject unsupported layouts/configurations explicitly, with no fallback

Required cuDNN Frontend dependency

This PR explicitly requires the row-scale grouped GEMM feature delivered by NVIDIA/cudnn-frontend#251: a cudnn.grouped_gemm_quant_wrapper_sm100(...) API that accepts row_scale_tensor.

The API is present in cuDNN Frontend v1.26.0. Validation with its pinned CUTLASS DSL 4.5.x also exposed missing explicit Int32 conversions in the cuDNN Frontend MoE scheduler. The TE tests pass with those cuDNN-side conversions applied temporarily. This PR intentionally carries neither a scheduler workaround nor a fallback, so it remains draft until the dependency used by TE contains both the row-scale API and the scheduler correction.

Motivation

This integration supports the row-scaled NVFP4 work from #2931.

Validation

  • python3 -m py_compile on the three changed Python files
  • git diff --check upstream/main
  • pre-commit run --all-files
  • B200, CUDA 13.3, PyTorch 2.13 nightly, cuDNN 9.23, cuDNN Frontend 1.26.0:
    • focused row-scaled grouped GEMM suite: 8 passed (bias/no-bias, list/single output, exact separate scales and offsets, canonical NVFP4 numerical tolerance, and four unsupported cases)
    • adjacent existing controls: 3 passed (row-scaled general_gemm, discrete NVFP4 pointer/scale packing, and MXFP8 cuDNN grouped GEMM)
    • real no-grad module integration: uneven splits [17, 33, 250] passed through Fp8Padding -> GroupedLinear -> Fp8Unpadding, with one cuDNN call and exact padded offsets
    • targeted pylint --errors-only on both changed production modules: passed

@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label May 26, 2026
@zianglih

Copy link
Copy Markdown
Contributor Author

need to rebase and refactor according to recent cute dsl integration

zianglih added 3 commits July 13, 2026 12:59
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
Signed-off-by: Ziang Li <ziangli@umich.edu>
@ziang-and
ziang-and force-pushed the codex/cudnn-row-scale-nvfp4-grouped-gemm branch from a0032dd to 96ae951 Compare July 13, 2026 21:51
@zianglih

Copy link
Copy Markdown
Contributor Author

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.

Signed-off-by: Ziang Li <ziangli@umich.edu>
@zianglih

Copy link
Copy Markdown
Contributor Author

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant