Support FP32 output and dynamic M in row-scaled FP4 grouped GEMM

[zianglih]

Summary

@HumansAnd

Extend the SM100 grouped GEMM quant path added in #251 so row-scaled NVFP4 integrations can use a unified cuDNN backend for the full shape and output-dtype matrix exercised by NVIDIA/TransformerEngine#3042.

Two frontend limitations blocked that integration:

• The API rejected FP4 with 16-element scale vectors and FP32 output even though the underlying kernel already supports the FP32 epilogue.
• The discrete-weight compile path specialized row_scale_tensor to the first runtime M. Its cache key intentionally treats M as dynamic, so reusing the cached kernel with another token count failed tensor-shape validation.

Logical K=64 is already supported by the Blackwell kernel. This change adds exact production-path coverage for K=64, N=256 rather than introducing padding or another backend.

Changes

• Allow FP32 output for FP4 grouped GEMM with 16-element scale vectors.
• Compile the discrete row-scale tensor with the same symbolic M used by A, D, and probability tensors.
• Add exact discrete-pointer, dynamic-scheduler tests for K=64/N=256 with:
  • BF16 and FP32 output
  • bias and no bias
  • one empty expert
  • row-scaled E2M1 data with E4M3 scale factors
• Add a cache regression that reuses one compiled kernel across M=1536 and M=3072.
• Correct the test reference amax for an empty expert to match the kernel's -inf initialization.
• Skip SM107 in the new row-scale tests because the Rubin grouped GEMM quant backend does not support row-scale fusion.
• Remove the obsolete FP32 limitation from the grouped GEMM quant documentation.

Validation

• python3 -m py_compile python/cudnn/gemm/cutedsl/grouped/quant/api.py test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py
• pre-commit run --files python/cudnn/gemm/cutedsl/grouped/quant/api.py test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md -> passed
• B200 --queue hell devbox using the exact editable cuDNN Frontend 1.27.0 checkout, CUTLASS DSL 4.6.0, CUDA 13.0, and cuDNN 9.22:
  • python3 -m pytest -q test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py -k 'discrete_wrapper_fp4_row_scale_small_k or discrete_wrapper_cache_dynamic_m_row_scale'
  • 5 passed, 315 deselected
• Existing discrete row-scale regression coverage:
  • python3 -m pytest -q test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py::test_grouped_gemm_quant_discrete_compile_execute_fp4_row_scale test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py::test_grouped_gemm_quant_discrete_wrapper_fp4_row_scale
  • 2 passed
• Prior TransformerEngine focused row-scaled NVFP4 integration coverage:
  • 14 passed
• Prior TransformerEngine full PyTorch sanity suite:
  • 17266 passed, 33922 skipped, 4 warnings

Notes

• No fallback GEMM backend is added; the integration remains entirely on the cuDNN path.
• pre-commit run --all-files was also attempted. Current hooks rewrote pre-existing vendored, benchmark, and notebook files outside this change; those unrelated rewrites were restored. All hooks for the changed files pass.

Summary by CodeRabbit

• New Features

  • Added support for FP4 grouped GEMM quantization with vector scale size 16 and float32 output.
  • Enabled discrete grouped GEMM configurations with dynamic dimensions, row scaling, optional bias, and dynamic scheduling.

• Bug Fixes

  • Improved handling of empty groups during grouped GEMM reference calculations.

• Tests

  • Expanded coverage for grouped GEMM quantization configurations and caching on supported hardware.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change enables additional discrete grouped GEMM quantization configurations. It adds validation for weight layouts, conditional row-scale tensor creation, expanded cache and wrapper test coverage, fused bias cases, and empty-group reference handling.

Changes

Grouped GEMM quantization

Layer / File(s)	Summary
FP4 and row-scale API support python/cudnn/gemm/cutedsl/grouped/quant/api.py, docs/fe-oss-apis/gemm_fusions/grouped_gemm_quant.md	The API removes the FP4/FP32 rejection, validates discrete b_major, and creates fake row-scale tensors only when row scaling is configured. The documentation removes the unsupported-configuration constraint.
Discrete cache and wrapper coverage test/python/fe_api/grouped_gemm/test_grouped_gemm_quant.py, test/python/fe_api/grouped_gemm/test_grouped_gemm_quant_utils.py	Tests cover configurable dtypes, scheduling, dimensions, group sizes, row scaling, fused bias, cache reuse, FP4 small-k cases, unsupported-case handling, and empty-group amax references.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Suggested reviewers: anerudhan

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely describes the primary changes to FP32 output and dynamic M support for row-scaled FP4 grouped GEMM.
Description check	✅ Passed	The description covers the changes, rationale, compatibility intent, and detailed validation, but omits several template headings and checklist items.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[zianglih]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

Backend pipeline not launched

Reason: PR base branch 'main' does not match 'develop'

[Anerudhan]

Thanks @zianglih for the contribution.
I will update the PR with CI results and if any change is needed.

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

Backend pipeline not launched

Reason: PR base branch 'main' does not match 'develop'

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[coderabbitai]

Note

GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer.

[Anerudhan]

@cudnn-ci-bot run

[Anerudhan]

@cudnn-ci-bot run oss

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-461-0444d32
Pipeline: 61316531
Targets: oss
Could not open: manual:oss (start them from the pipeline page)