Add NVFP4 RHT Support for SM120 and SM121

[new-TonyWang]

Summary

• use a reference-equivalent ATen matrix multiplication for the unfused transposed RHT path on SM120/SM121
• preserve the existing nvte_hadamard_transform path on other architectures
• derive the ATen dtype from the TE input dtype rather than hard-coding BF16
• keep the existing native non-fused NVFP4 quantization after the RHT step

Motivation

The fused RHT+quantize kernel is not eligible on SM120/SM121. The existing unfused BF16 MMA Hadamard kernel is mathematically equivalent to the Python reference, but its accumulation order can differ by one BF16 ULP. Values near FP4 rounding thresholds can therefore produce different exact FP4 codes or E4M3 scales.

Using the same block-matmul formulation as the reference on SM120/SM121 eliminates that numerical-path mismatch while leaving other architectures unchanged.

Validation

On NVIDIA GB10 (SM121):

• tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py: 200 passed
• random-sign and no-random-sign cases pass
• C++ and Python allocator paths pass
• contiguous, non-contiguous, padded, and large matrix cases pass

Standalone Hadamard matmul checks were also run for FP32, TF32, FP16, and BF16; all were deterministic across repeated runs and produced finite outputs.

[greptile-apps]

Greptile Summary

The PR adds a dedicated SM120/SM121 fused NVFP4 columnwise RHT path while retaining existing dispatch behavior on other architectures.

• Performs the 16-point BF16 Hadamard transform with warp MMA and directly emits columnwise FP4 data and compact scales.
• Computes and globally reduces post-RHT amax before quantization.
• Adds exact single-GPU and distributed SM12x regression coverage.
• Adds an environment switch for forcing the unfused path in reference comparisons.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the post-RHT amax now follows the same BF16 MMA arithmetic as the fused transform, and distributed reduction occurs before quantization consumes the resulting read-only amax.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh	Adds the SM12x warp-MMA RHT implementation, columnwise-only dispatch support, and associated shared-memory/TMA handling without leaving a blocking issue.
transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu	Routes SM120/121 requests through the new no-TMEM NVFP4 path while preserving the existing path for other architectures.
transformer_engine/pytorch/csrc/quantizer.cpp	Expands RHT fusion eligibility to SM120/121 and preserves the required post-quantization scale swizzle behavior.
tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py	Adds exact fused-versus-native and distributed post-RHT amax regression tests for SM120/121.
tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py	Updates SM12x exact comparisons to use the native unfused arithmetic path as the reference.

Sequence Diagram

sequenceDiagram
  participant Q as NVFP4Quantizer
  participant A as Hadamard amax kernel
  participant D as Distributed MAX
  participant F as SM12x fused RHT+quantize
  Q->>A: Compute rowwise and post-RHT columnwise amax
  A-->>Q: Local amax values
  Q->>D: All-reduce amax when enabled
  D-->>Q: Global amax values
  Q->>F: Input, RHT matrix, and reduced amax
  F->>F: "BF16 A @ H via warp MMA"
  F->>F: Derive block scales and FP4 codes
  F-->>Q: Quantized rowwise/columnwise outputs

Loading

_{Reviews (9): Last reviewed commit: "perf(nvfp4): quantize SM12x RHT fragment..." | Re-trigger Greptile}

[new-TonyWang]

Good catch. The SM120/SM121 ATen path was initially reusing the columnwise amax produced earlier by nvte_hadamard_transform_amax, so the transformed data and its amax could come from different accumulation orders. Fixed in cc828f5: the unfused helper now receives compute_amax and, after matmul_out, recomputes the columnwise post-RHT amax directly from that ATen output before nvte_quantize_v2. The supplied/precomputed amax path remains unchanged when compute_amax == false.

[new-TonyWang]

Addressed the distributed-reduction issue in c7a3b2a. On SM120/SM121, the ATen-derived rowwise and post-RHT columnwise local amax values are now written during the existing amax-computation phase, before reduce_amaxes(). The unfused quantization helper no longer overwrites amax after reduction, so distributed ranks consume the group-reduced value while quantizing the ATen-generated RHT data.

[new-TonyWang]

Added SM120/SM121 regression coverage in ea87075. The new tests (1) compare the stored post-RHT columnwise amax exactly against the ATen/reference RHT output with random-sign masking both disabled and enabled, and (2) exercise two-rank NCCL amax reduction with different rank-local maxima to verify the ATen-derived value is reduced globally before quantization. The distributed case skips unless two SM120/SM121 GPUs are available.

[new-TonyWang]

Implemented solution A in 7a0a0d0: the SM120/SM121 unfused path now materializes the ATen RHT intermediate once during post-RHT amax computation, derives local amax from that exact buffer, carries it across distributed MAX reduction, and reuses it for nvte_quantize_v2 instead of running transpose+matmul a second time.

Validation on Spark GB10 (SM121):

• 302 passed, 1 skipped: test_nvfp4_rht_sm12x.py + test_nvfp4_rht_quantize_exact.py (the skipped case requires 2 GPUs).
• A/B median end-to-end columnwise RHT quantization latency, same build/environment:
  • 128x128: 55.647 us -> 45.799 us (-17.7%)
  • 1024x1024: 107.837 us -> 69.376 us (-35.7%)
  • 4096x4096: 2216.001 us -> 1429.687 us (-35.5%)

The post-RHT amax remains derived before distributed reduction and from the same BF16 RHT data that is subsequently quantized.

[new-TonyWang]

Implemented the SM120/SM121 no-TMEM fused RHT+NVFP4 quantize path in 434f242.

Key points:

• Reuses the existing 1D TMA quantize/transpose pipeline; the H16 matrix is loaded once per CTA via a separate TMA transaction/mbarrier.
• Computes the 16-point columnwise RHT in FP32 registers, rounds to BF16 at the same reduction boundaries selected by ATen on SM12x, and derives block scales/codes from those exact values.
• Uses the tuned 8-value BF16-to-FP4 conversion helpers, including stochastic rounding.
• Keeps SM100/110 on the existing UMMA/TMEM kernel and limits the new path to non-grouped SM120/121 dispatch.
• Rowwise-only output remains on the ordinary 1D quantizer; columnwise-only and dual-output modes use the fused RHT path.
• Makes RHT application a compile-time host-template option so unrelated quantize translation units do not instantiate the RHT kernel specialization.

Spark GB10 (SM121) validation:

• focused exact suite across small/large eligible shapes, rowwise-only/columnwise-only/dual output, C++/Python allocation, and sign-mask on/off: 70 passed
• dedicated SM12x regression suite comparing packed FP4 codes, scale bytes, and amax against ATen RHT -> tuned quantizer: 10 passed, 1 skipped (2-GPU NCCL case)

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.

[new-TonyWang]

SM120/121 fused RHT update: replaced the per-thread scalar 16-point Hadamard reduction with BF16 warp MMA (WMMA), while matching ATen accumulation/rounding order for both small and large GEMM regimes.

Spark GB10 (SM121) validation:

• test_nvfp4_rht_sm12x.py: 10 passed, 1 skipped
• selected exact-reference matrix (128–2048, rowwise+columnwise modes): 96 passed
• full CUDA build: passed

Square BF16 columnwise-only benchmark (post-RHT amax/sign mask), median us:

• 1K: 100.653
• 2K: 337.343
• 4K: 1560.165
• 8K: 6262.939
• 16K: 25726.336
• 32K: 129565.420

Compared with the preceding scalar-fused implementation, changes are -13.39%, +1.68%, +2.09%, -6.65%, -2.37%, and -1.86% respectively (positive means faster). Thus MMA improves 2K/4K slightly, but the current shared-memory staging and two-pass K=8 compatibility path regress the other sizes; this is correctness-first and not yet a net performance win.

Commit: f9e28f4

[new-TonyWang]

Updated the SM120/121 numerical contract to match Transformer Engine's native Hadamard implementation rather than the architecture-dependent ATen/cuBLAS result.

Key changes (36e1179f):

• fused path now issues one BF16 m16n16k16 warp MMA per tile (removed both masked K=8 passes)
• restored nvte_hadamard_transform and nvte_hadamard_transform_amax for the unfused/reference path on SM120/121
• exact SM12x tests compare fused output against TE's native unfused K=16 path
• added NVTE_NVFP4_DISABLE_RHT_CAST_FUSION=1 to make that native reference path directly testable on otherwise fusion-eligible shapes

Spark GB10 / SM121 validation:

• SM12x native-K16 regression tests: 18 passed, 1 skipped
• selected exact-reference matrix (128 through 2048, rowwise/columnwise modes): 96 passed
• pre-commit: passed
• full CUDA build: passed

Single-K16 median latency (us):

• 1K: 97.980
• 2K: 331.205
• 4K: 1640.145
• 8K: 6439.363
• 16K: 25605.623
• 32K: 129225.983

Relative to the two-masked-K8 implementation: 1K +2.66%, 2K +1.82%, 4K -5.13%, 8K -2.82%, 16K +0.47%, 32K +0.26% (positive means faster). Variance/memory traffic dominates at the larger sizes, but the redundant masked MMA work is now removed and the arithmetic exactly matches TE native K16 behavior.

[new-TonyWang]

Implemented all four follow-up SM120/121 RHT staging optimizations in 9923d1cd:

1. Removed FP32 accumulator staging: the fused path now calls TE's native mma_m16_n16_k16_b16_b16_b16_noacc helper, which performs the single K=16 MMA and FP32-to-BF16 conversion directly in registers.
2. Direct fragment consumption: each warp reduces per-row amax across the WMMA 4-lane ownership group, computes the E4M3 scale, converts its BF16 fragment directly to packed FP4, and writes the final packed pairs to the TMA output buffer. The full BF16 rht_result_sh round trip is gone.
3. Cached H operand: each warp loads the 16x16 Hadamard B fragment once before its tile loop.
4. Removed explicit A transpose staging: a col-major wmma.load.a consumes the existing TMA input layout directly. rht_mma_a_sh is gone.

RHT-only extra dynamic shared memory is reduced from about 14.5 KiB/block (H + BF16 result + A staging + FP32 accumulator staging) to 512 B/block (H only).

Spark GB10 / SM121 validation:

• full CUDA build: passed
• native-K16 SM12x regression: 18 passed, 1 skipped
• selected exact matrix including partial 64-row tail (128 through 2048, all quantize modes): 96 passed
• pre-commit: passed

Two final benchmark runs (median us; run-to-run variance is visible):

• 1K: 86.823 / 87.525
• 2K: 317.838 / 325.033
• 4K: 1622.316 / 1595.686
• 8K: 6453.001 / 6359.634
• 16K: 25816.952 / 25641.775
• 32K: 129122.548 / 129209.605

Using the second run versus the previous staged single-K16 version gives: 1K +10.67%, 2K +1.86%, 4K +2.71%, 8K +1.24%, 16K -0.14%, 32K +0.01% (positive means faster). Large shapes remain dominated by global-memory/TMA traffic.