Add row-scale support to grouped GEMM quant

[zianglih]

Summary

@HumansAnd

Add an optional row_scale_tensor input to the SM100 grouped GEMM quant path. The new tensor is a contiguous FP32 vector with one scale per valid output row. When provided, the epilogue multiplies FP32 accumulators by alpha[expert] * row_scale[m] before output conversion.

This is intended for integrations that need an additional per-row accumulator scale, such as row-scaled NVFP4. The immediate motivation is TransformerEngine's row-scaled NVFP4 work in NVIDIA/TransformerEngine#2931. The kernel treats row_scale_tensor as an epilogue multiplier and does not interpret how the caller produced those values.

The existing bias and probability semantics are preserved: the row scale applies only to the accumulator path, while fused bias keeps the existing bias * prob behavior.

Changes

• Add sample_row_scale / row_scale_tensor plumbing to GroupedGemmQuantSm100.
• Add scheduler tensor mapping for the row-scale input.
• Apply row-scale multiplication in the grouped GEMM quant epilogue.
• Include the row-scale tensor in wrapper validation and compile-cache signatures.
• Extend grouped GEMM quant tests and Python reference coverage for compile/execute, dense wrapper, and discrete wrapper paths.

Validation

Current upstream-facing branch:

• python3 -m py_compile python/cudnn/grouped_gemm/grouped_gemm_quant/api.py python/cudnn/grouped_gemm/grouped_gemm_quant/grouped_gemm_quant.py python/cudnn/grouped_gemm/moe_sched_extension.py test/python/fe_api/test_grouped_gemm_quant.py test/python/fe_api/test_grouped_gemm_quant_utils.py
• git diff --check upstream/develop...HEAD
• B200 devbox: python3 -m pytest -q test/python/fe_api/test_grouped_gemm_quant.py -k row_scale --tb=short -> 5 passed, 309 deselected

Notes:

• pre-commit run --all-files is unavailable in this checkout because .pre-commit-config.yaml is not present.

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-251-dee5400
Pipeline: 52592384

[Anerudhan]

Thanks for the PR @zianglih, I have launched an internal review. I will let you know on its progress.

[Anerudhan]

@cudnn-ci-bot run

[cudnn-ci-bot]

🚀 Running mirror pipeline

Branch: cudnn-gh/pr-251-55ae251
Pipeline: 52881929

[saltyminty]

Approved but please see above comment about dynamic shapes.

[saltyminty]

I believe the fake tensor needs to also be constructed with symbolic valid_m as the 0th dimension, similar to prob_tensor above.

[zianglih]

Follow-up from validating the row-scale integration in NVIDIA/TransformerEngine#3042 on B200:

The row_scale_tensor API is present in v1.26.0 and the row-scale wrapper contract is correct, but the MoE scheduler does not compile under the release's pinned nvidia-cutlass-dsl[cu13]==4.5.0 without explicit Int32 conversions. Both dynamic and static scheduler attempts hit a DSL type-structure error. The working cuDNN-side changes are explicit Int32(...) conversions for:

• scheduler state updates in _advance_to_expert_containing
• tile-index initialization and division/modulo results in _convert_local_idx_to_tile_coords
• offset loads/subtraction results in get_expert_offset_and_tokens

With those conversions applied temporarily, the cuDNN discrete FP4 row-scale control and all 8 focused TE integration cases pass. The TE PR intentionally includes no scheduler workaround or fallback and remains draft pending a cuDNN-side fix.