[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]

Confidence Score: 5/5

[Medium risk] Adds grouped amax computation for row-scaled quantization.

The PR appears safe to merge based on the reviewed changes and the fixes for both previous findings.

Summary

The PR adds a grouped NVFP4 amax API and CUDA path that computes per-expert rowwise and columnwise maxima from packed BF16 input, with C++ correctness tests and a disabled benchmark. The current code checks column alignment and skips null per-expert amax buffers, addressing both previous findings.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Packed BF16 input] --> B[Grouped amax API]
  C[Expert row counts] --> B
  B --> D[Zero allocated amax vectors]
  D --> E[Grouped CUDA amax kernel]
  E --> F[Per-expert rowwise and columnwise amax]

Loading

_{Reviews (7) · Last reviewed commit: "Merge branch 'main' into pr/nvfp4-row-sc..."}

[timmoon10]

LGTM, pending CI.

[timmoon10]

/te-ci core

[timmoon10]

/te-ci core

[timmoon10]

/te-ci core

[timmoon10]

This function name does not mention row-scaled NVFP4 and the API is natural for both
tensor-scaled NVFP4 and row-scaled NVFP4. We only support row-scaled NVFP4 currently, so we should add an internal check to error out if tensor-scaled NVFP4 tensors are provided. The docs mention that this is specific to row-scaled NVFP4, but that's very flimsy.

[timmoon10]

Also, this API is redundant with nvte_group_amax:

TransformerEngine/transformer_engine/common/include/transformer_engine/multi_tensor.h

Line 313 in 5d7bbd8

void nvte_group_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections,

[timmoon10]

#3649 deprecates this function in favor of nvte_group_amax.