[Common] row-scaled nvfp4 path: add single-launch group fused amax

[cael-ling]

Description

Adds a single-kernel row/col amax pass for grouped (MoE) row-scaled NVFP4. It replaces the per-expert loop of compute_rowwise_amax + compute_columnwise_amax with one launch over a packed (sum_M, K) input, exposed as nvte_group_nvfp4_compute_amax. The kernel is the grouped generalization of #3454 - the single-tensor compute_fused_amax_kernel (same tiling / SMEM pipeline / acquire-barrier reduction). A single input TMA map covers all experts; per-expert amax vectors are selected via a prefix sum of row counts, so a 128-row chunk always lands in one expert.

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

New gtest tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu:

• NVFP4GroupedRowScaledAmaxTestSuite.MatchesPerExpertOracle — grouped amax vs a CPU per-row/per-col max(|x|) reference over the same BF16 bytes, plus cross-launch determinism, across 6 expert/shape configs.
• DISABLED_LoopVsGrouped — microbench comparing the same kernel called per-expert (loop)
  vs once (grouped).

cmake -Btests/cpp/build tests/cpp && cmake --build tests/cpp/build -j
./tests/cpp/build/operator/test_operator --gtest_filter='*NVFP4GroupedRowScaledAmax*'
./tests/cpp/build/operator/test_operator --gtest_also_run_disabled_tests --gtest_filter='*GroupedRowScaledAmaxBench*'

Performance

Same kernel called per-expert (loop) vs a single grouped launch, GB200, K=4096. Speedup is loop / grouped.

experts	M=128 loop→grouped (ms)	speedup	M=512 loop→grouped (ms)	speedup
2	0.0179 → 0.0117	1.52x	0.0180 → 0.0124	1.45x
4	0.0288 → 0.0120	2.39x	0.0271 → 0.0138	1.97x
8	0.0513 → 0.0124	4.15x	0.0514 → 0.0165	3.11x
16	0.0950 → 0.0140	6.81x	0.0960 → 0.0238	4.04x
32	0.1852 → 0.0173	10.70x	0.2009 → 0.0370	5.43x
64	0.3631 → 0.0232	15.65x	0.3960 → 0.0607	6.53x

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

[greptile-apps]

Greptile Summary

Adds a public grouped NVFP4 amax operation that computes per-expert rowwise and columnwise maxima from a packed BF16 input in one CUDA launch.

• Introduces a Blackwell-specific grouped fused-amax kernel with per-expert output selection and optional rowwise/columnwise buffers.
• Validates 128-element shape alignment and grouped split coverage at the native boundary.
• Adds correctness, determinism, heterogeneous-output, and disabled performance-benchmark coverage.
• The current head is identical to the previously reviewed SHA, and both earlier findings are resolved.

Confidence Score: 5/5

The PR appears safe to merge because no new changes or outstanding findings remain at the current head.

The current head exactly matches the previous review SHA. Both prior findings—truncated trailing columns and null optional amax-buffer writes—are resolved, and no new repository-rule violations were identified.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh	Implements the grouped TMA-based fused-amax kernels, validation, optional-buffer handling, and launch configuration.
transformer_engine/common/cast/cast_grouped.cu	Adds the native C API implementation and converts opaque tensor handles for grouped dispatch.
transformer_engine/common/include/transformer_engine/cast.h	Declares and documents the new ABI-facing grouped NVFP4 amax operation.
tests/cpp/operator/test_row_scaled_nvfp4_grouped_amax.cu	Adds grouped correctness, determinism, heterogeneous columnwise-output, and performance-benchmark coverage.
tests/cpp/operator/CMakeLists.txt	Registers the grouped NVFP4 amax test source in the native operator test target.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Packed BF16 input<br/>sum_M × K] --> B[nvte_group_nvfp4_compute_amax]
  S[Per-expert row splits] --> B
  B --> Z[Zero requested amax buffers]
  Z --> K[Grouped fused amax kernel]
  K --> E[Select expert from row-prefix range]
  E --> R[Per-expert rowwise amax]
  E --> C[Per-expert columnwise amax]

Loading

_{Reviews (4): Last reviewed commit: "[Common] Enforce K alignment and per-exp..." | Re-trigger Greptile}