[Common] Use software stochastic rounding in the RHT cast fusion kernels

[0z5a]

Description

#3281 added a software stochastic rounder (ptx::stochastic_round_fp4_e2m1) for architectures without cvt.rs.satfinite.e2m1x4.f32, and used it in ptx_arch_spec.cuh and the blockwise FP4 transpose. The RHT cast-fusion kernels were left out. Their StochasticNumericConverterBase (five copies under common/hadamard_transform/) still fall back to NVTE_DEVICE_ERROR when ARCH_HAS_STOCHASTIC_ROUNDING is false.

These kernels run on SM100–110 (NVFP4Quantizer::is_eligible_for_rht_cast_fusion; SM120/121 take the no-TMEM path). Among those targets, cvt.rs exists only on sm_100a/sm_103a (ARCH_HAS_STOCHASTIC_ROUNDING also lists 107a), so in practice this affects SM110: Jetson AGX Thor, which is SM101 with CUDA 12.x toolkits. NVTE_DEVICE_ERROR is printf + assert(0), and release builds define NDEBUG. The converter therefore prints from every thread, leaves its output unwritten, and nothing reaches the host. With the default NVFP4BlockScaling() recipe (RHT + stochastic rounding on gradients), wgrad comes back all zeros and dgrad has zero cosine to a BF16 reference. Details are in #2255 (comment).

This PR uses the #3281 rounder in that branch. The cvt.rs path is unchanged: an SM100 build compiles these sources for 100a/103a/107a, where ARCH_HAS_STOCHASTIC_ROUNDING is true and the new code is discarded by if constexpr.

Related: #2255, #3281, #2556

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

• common/util/ptx.cuh: add ptx::stochastic_round_fp4_e2m1x4, which rounds four values with stochastic_round_fp4_e2m1 and packs them like cvt.rs.satfinite.e2m1x4.f32. Value i uses byte i of rbits and goes to nibble i, the same convention as the [Common] Fix NVFP4 stochastic rounding on architectures without cvt.rs #3281 fallbacks in ptx_arch_spec.cuh.
• hadamard_transform/{,group_,row_cast_col_,group_row_cast_col_,graph_safe_group_row_cast_col_}hadamard_transform_cast_fusion.cu: in StochasticNumericConverterBase, call it when cvt.rs is unavailable. NVTE_DEVICE_ERROR stays only for builds without FP4 types (CUDA < 12.8), because these sources are compiled there too.
• tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py: add a (2560, 2048) case to test_group_stochastic_rounding_quantization_versus_reference. With 8 splits the 320-row sections are not 128-aligned, so split_quantize takes the columnwise-only grouped kernel (group_rht_gemm_device). No existing SR test reached that kernel.

Validation

Jetson AGX Thor (CC 11.0): JetPack 7.2, CUDA 13.2, cuDNN 9.20, torch 2.13.0+cu130. Built with NVTE_CUDA_ARCHS=110 (110 + 110f) and NVTE_WITH_NCCL_EP=0, from 9d9bd9e5 + this patch (+ the #3696 include fix needed to build there). The three commits between 9d9bd9e5 and this PR's base do not touch transformer_engine/common.

te.Linear(2048, 2048) and te.GroupedLinear(4, 2048, 2048) (per-tensor and use_grouped_tensor=True), BF16, M=512, default NVFP4BlockScaling(), forward + backward, compared against the same layers in plain PyTorch BF16. Without this PR the default recipe gave dgrad and wgrad cosine 0.0 and all-zero wgrad on te.Linear. With it:

	kernel that ran	dgrad cosine	wgrad cosine	mean |wgrad| (ref)	device printf lines
te.Linear, SR on (default)	row_col_rht_gemm_device	0.985	0.986	18.14 (18.05)	0
te.Linear, NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING=1	row_col_rht_gemm_device	0.989	0.991	18.04 (18.05)	0
te.GroupedLinear, SR on	group_row_col_rht_gemm_device	0.984	0.986 (each of 4)	9.03–9.06 (8.98–9.02)	0
te.GroupedLinear(use_grouped_tensor=True), SR on	group_row_col_rht_gemm_device_graph_safe	0.984	0.986 (each of 4)	9.03–9.06 (8.98–9.02)	0

The forward cosine is 0.989 in every case. A single SR draw has a slightly lower cosine than RN; SR is unbiased, not lower-error per draw.

Test suites, same commands as qa/L0_pytorch_unittest/test.sh, compared test by test against the run before this PR:

suite	before	after
tests/pytorch/nvfp4	9299 passed, 15 failed, 6871 skipped (2730 s)	9314 passed, 0 failed, 6871 skipped (415 s)
test_hybrid_quantization.py	309 passed, 4 failed	310 passed, 3 failed
test_grouped_linear.py	1850 passed, 0 failed	1850 passed, 0 failed

• tests/pytorch/nvfp4: the 15 tests that change are exactly the ones that failed before. 14 are in test_nvfp4_sr_quantize.py ("Stochastic rounding failed - error larger than the round to nearest"), and the 15th is test_nvfp4_rng_lifetime.py::test_nvfp4_stochastic_rounding_allocation_independence. No test changed in the other direction. The earlier run was slower because of the device printf flood.
• test_hybrid_quantization.py: TestHybridMixedFormatQuantizedParams::test_mixed_format_training_loop now passes. Its loss was flat before, because its NVFP4 backward quantizer enables SR. The 3 remaining failures are TestAttentionFactoryNativeRecipeParity MXFP8 cases. They fail because cuDNN 9.20 is below the 9.21 that MXFP8 FusedAttention requires, and are unrelated to this change.
• None of the logs contains the NVTE_DEVICE_ERROR message.

Which of the five converters these runs reach (CUDA kernels recorded per test with torch.profiler):

file	kernel	reached with SR on by
row_cast_col_…	row_col_rht_gemm_device	test_nvfp4_sr_quantize.py, the te.Linear case above
group_row_cast_col_…	group_row_col_rht_gemm_device	test_group_stochastic_rounding_quantization_versus_reference, the te.GroupedLinear case
graph_safe_group_row_cast_col_…	group_row_col_rht_gemm_device_graph_safe	the use_grouped_tensor=True case only; the graph-safe tests in tests/pytorch/nvfp4 run it with SR off
group_…	group_rht_gemm_device	the new (2560, 2048) case with 8 splits; test_nvfp4_group_quantize.py runs it with SR off
hadamard_transform_cast_fusion.cu	rht_gemm_device	not exercised. In this repo only the JAX extension calls nvte_hadamard_transform_cast_fusion_columnwise, and I have no JAX build on this machine. The change there is identical to the other four files.

The new test case passes in all 4 parametrizations. In [True-8-…-2560-2048] the profiler shows that group_rht_gemm_device is the only RHT kernel. The full-suite numbers above were collected before I added this case. I also checked this kernel with SR through the same helper, using split sizes 4×64, 2×192 and 8×64: all pass. I did not re-run the new case against a build without the fix. Like the other grouped SR cases, it asserts that the mean of 50 SR draws beats RN, which an unwritten output cannot satisfy.

The build produces no new warnings. Its warning set is the same as before, except that one existing warning in row_cast_col_… moved down eight lines. I also compiled row_cast_col_… and hadamard_transform_cast_fusion.cu with the same nvcc command for sm_100a (the cvt.rs path) and sm_90. Both succeed. The only warnings are ones the unpatched sources also produce, such as StochasticNumericConverter being unused on sm_90. cpplint 1.6.0 and clang-format 18.1.6 report nothing on the changed files.

Not validated on SM100 (no hardware); that path is unchanged by construction.

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 (the support matrix update for CC 11.0 is a separate docs PR)
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works (the existing SR tests in tests/pytorch/nvfp4 fail on SM110 without this change; the new case covers the one remaining kernel the PyTorch path reaches)
• New and existing unit tests pass locally with my changes

🤖 Generated with Claude Code

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no actionable issues were found.

Summary

Adds software stochastic rounding to five RHT cast-fusion converters when the hardware instruction is unavailable.

• RHT cast-fusion kernels round FP4 values without hardware cvt.rs.

No actionable issues found. Builds and GPU tests were not run during this review.

_{Reviews (1) · Last reviewed commit: "[Common] Use software stochastic roundin..." · Reviewed by Greptile}