[Common] Fix NVFP4 stochastic rounding on architectures without cvt.rs

[davidkny22]

Description

NVFP4 quantization with stochastic rounding returns all-zero data on sm_120 and sm_121.

The four e2m1 stochastic rounding helpers gate on ARCH_HAS_STOCHASTIC_ROUNDING, which is ArchSpecific<100> and ArchSpecific<103>, matching where the cvt.rs instruction exists. The other branch is NVTE_DEVICE_ERROR, which is printf plus assert(0). Release builds define NDEBUG, so the assert is compiled out, the helper returns its zero-initialized output, and each thread prints:

.../common/util/ptx.cuh:953 in function mul_cvt_bf16_to_fp4_8x_stochastic_rounding (thread (0,0,0), block (1,0,0)): FP4 cvt PTX instructions are architecture-specific. Try recompiling with sm_XXXa instead of sm_XXX.

NVFP4BlockScaling sets stochastic_rounding on gradient quantization by default and _compute_nvfp4_support() reports NVFP4 for every device at compute capability 10.0 or above, so on these parts the default recipe trains on zeroed gradients while the NVFP4 GEMMs themselves work. The message is also not actionable here: the build already used NVTE_CUDA_ARCHS=121a, and cvt.rs does not exist at any feature level on this family.

This replaces that branch with a software e2m1 stochastic rounder. It takes 8 random bits per element, which is what cvt.rs.satfinite.e2m1x4.f32 takes from its rbits operand, adds u * step where step is the grid step at that magnitude, and truncates onto the grid, so a value rounds up with probability equal to its fractional position between its two neighbours. The if constexpr is unchanged, so architectures with cvt.rs keep the asm path and nothing else in the helpers moves.

The edges follow the PTX ISA text for .satfinite, under cvt, Floating Point Notes: for .e2m1x4 destinations, "NaN results are converted to positive MAX_NORM", and an input above MAX_NORM in absolute value gives "sign-preserved MAX_NORM of the destination format". So NaN gives +6 and overflow gives ±6.

A fallback rather than refusing stochastic rounding on these architectures: _compute_nvfp4_support() already advertises NVFP4 here and the default recipe turns SR on, so a capability gate would replace a silent wrong answer with a hard failure of the default recipe on hardware that otherwise runs it. The fallback keeps recipe parity.

The fallback is not bit identical to cvt.rs and cannot be. cvt.rs adds rbits to the mantissa bits the conversion discards and rounds on the carry out; this adds a scaled uniform to the value and truncates. Both are stochastic rounding at the same per-element entropy, and their round-up probabilities agree to the 1/256 that 8 bits resolve. There is no cvt.rs behaviour on sm_120 or sm_121 to differ from, and stochastic rounding is not reproducible across architectures in any case.

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

• transformer_engine/common/util/ptx.cuh: add stochastic_round_fp4_e2m1 and use it in the non-cvt.rs branch of mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding, mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding and mul_cvt_bf16_to_fp4_8x_stochastic_rounding.
• transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu: the same in cvt_fp32_to_fp4_4x_with_stochastic_rounding.
• tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py: test that every stochastically rounded value sits on one of the two E2M1 grid values that bracket the exact one, and that two draws differ.

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

No documentation changes needed. On the last box, see below.

Testing done

GB10 (DGX Spark), sm_121a, CUDA 13.1, driver 580.95.05, main @ 0ee1320, built with NVTE_CUDA_ARCHS=121a. clang-format 18.1.6, cpplint 1.6.0 and black 24.4.2 are clean on the changed files, and there are no compiler warnings in the build.

The new test is not gated on architecture, so it exercises cvt.rs where the instruction exists and this branch here. All four parametrizations fail without the patch and pass with it, on separate unpatched and patched builds of the same tree, twice each. Together they reach three of the four helpers: cvt_fp32_to_fp4_4x_with_stochastic_rounding for FP32 input, mul_cvt_bf16_to_fp4_8x_stochastic_rounding for BF16 with 1D scaling, and mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding for BF16 with 2D. I could not reach mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding from any NVFP4 entry point on this hardware; it is the FP32 twin of the same columnwise loop and is fixed the same way.

tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py is 6 failed, 22 passed, 4 skipped with the patch, in 37 seconds. Without it the same file does not finish: I stopped it at 61 minutes, since the error path prints once per thread and this file quantizes 8192x8192 tensors 50 times per case.

The 6 that fail are test_group_stochastic_rounding_quantization_versus_reference with use_tex_split_quantize=True, and they are not this path. They fail at group_row_cast_col_hadamard_transform_cast_fusion.cu:1276 in function group_row_col_rht_gemm_ntt_w_sfc: CUDA Error: invalid argument, which is the grouped RHT dispatch reported in #3219. tex.split_quantize with RHT-enabled NVFP4 quantizers raises the same error on an unpatched build of this tree.

On one 256x256 BF16 tensor, dequantized cosine against the input is 0.000000 before and 0.990767 after, against 0.995495 for round to nearest on the same tensor. Stochastic rounding sitting slightly below round to nearest on a single draw is the expected direction.

If there is anything else you want checked on sm_121, I can run it.

cc @Oleg-Goncharov @ptrendx

[greptile-apps]

Greptile Summary

Adds a software E2M1 stochastic-rounding fallback for architectures without cvt.rs, applies it across FP32/BF16 NVFP4 conversion helpers, and adds PyTorch coverage for adjacent-grid rounding and redraw behavior.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete correctness or security defects identified in the changed fallback paths or tests.

The software path matches the E2M1 grid and satfinite edge behavior, preserves per-element random-bit allocation and scale mapping, and packs values in the same logical order as the existing hardware conversion paths.

Important Files Changed

Filename	Overview
transformer_engine/common/util/ptx.cuh	Implements the shared software E2M1 stochastic rounder and uses it in BF16/FP32 four- and eight-element fallback conversion paths while preserving lane, scale, and nibble ordering.
transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu	Replaces the architecture error path with the shared software rounder for vector blockwise FP32 conversion.
tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py	Adds FP32/BF16 and 1D/2D tests asserting stochastic outputs remain on adjacent E2M1 grid points and successive draws differ.

_{Reviews (1): Last reviewed commit: "Fix NVFP4 stochastic rounding on archite..." | Re-trigger Greptile}