[tdophung]

Description

Applies to NVFP4 1D quantization the optimization #3459 made for MXFP8 cast-only quantization: load the input straight from global memory into registers (G2R) instead of staging it through shared memory with TMA. The new kernel serves BF16 input for rowwise and rowwise + columnwise output, with round-to-nearest or stochastic rounding and E4M3 or UE5M3 scales. Requests it does not cover keep using quantize_transpose_tuned_1D: row-scaled NVFP4, fast math, GEMM-swizzled scales, the columnwise RHT path, and tensors past the kernel's grid limit.

The CuTeDSL counterpart of this kernel is a separate draft PR, #3809.

Points for discussion:

• Stochastic rounding random stream. The new kernel draws Philox once per 16-element output block, keyed by the block's coordinates and direction. Each draw is exactly what curanddx returns for that key and counter. For a given seed, the random bits therefore differ from tuned_1D's, whose stream is keyed by thread id. SR output stays unbiased and statistically equivalent, but it is not bitwise reproducible against main. Round-to-nearest output is byte-identical to main.
• Amax kernel rewrite. The first pass of every NVFP4 quantize, which FP8 current scaling also uses, is now a persistent kernel. Its result is bit-identical to the old kernel. It is up to 1.10x faster at large sizes and 3-12% (0.1-0.6 us) slower at 2048x2048 and below, so its end-to-end effect needs checking.

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

• quantize_transpose_nvfp4_register_1D.cuh: a register-resident BF16 -> NVFP4 1D quantize kernel, dispatched ahead of quantize_transpose_tuned_1D for every request it supports.
  • The rowwise-only kernel reads each 16-element block with one 32-byte load.
  • The rowwise + columnwise kernel reads each row twice, once per direction, with the second read served from L2. Only the transposed output is staged in shared memory.
  • Each per-block coefficient depends only on the block's scale byte, so the kernel computes the divides once per CTA into a lookup table.
• nvfp4_block_philox.cuh: per-block Philox for stochastic rounding. The number of rounds follows NVTE_BUILD_NUM_PHILOX_ROUNDS; builds with fewer than 3 rounds keep tuned_1D for stochastic rounding.
• recipe/current_scaling.cu: the amax kernel becomes persistent. Each lane keeps four 32-byte loads in flight, 16-bit inputs use packed abs-max on SM 8.9+, and any alignment is supported.
• util/ptx.cuh, util/ptx_arch_spec.cuh: new PTX helpers.
• Tests:
  • A rowwise_only parameter for test_cast_nvfp4_transpose.cu, which previously always requested both directions.
  • New test_cast_nvfp4_sr_streams.cu, which checks three things: the generator matches curanddx, every block gets its own subsequence, and the rowwise and columnwise draws are independent.

Results (GB200, ncu kernel time, BF16 randn, E4M3 scales, cold L2)

Speedup over main. Shapes run from 1024x1024 to 32768x16384, and the larger speedups are at the smaller shapes.

Mode	Speedup vs main
Rowwise, round-to-nearest	1.04-1.43x
Rowwise, stochastic rounding	1.12-1.36x
Rowwise + columnwise, round-to-nearest	1.12-1.38x
Rowwise + columnwise, stochastic rounding	1.23-1.45x

Testing

• C++ test_operator with --gtest_filter='OperatorTest/FusedCastTransposeNVFP4*:OperatorTestUE5M3/FusedCastTransposeNVFP4*:*StochasticRoundingStreams*': 198/198 passed.
• Output comparison against main on 580 cases:
  • Coverage: E4M3 and UE5M3 scales, missing amax, the noop flag, special values, and ragged and tall shapes, plus 49152x32768 and 69632x32768.
  • Round-to-nearest output, the scales and the amaxes are byte-identical to main.
  • Stochastic rounding codes differ only by rounding to the adjacent value, as expected from the different random stream.

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