Add NVFP4 four-over-six (row-wise) quantization to the training prototype

[wolfcomos]

What

Adds the NVFP4 four-over-six adaptive block-scaling recipe to
torchao/prototype/moe_training/nvfp4_training, including the row-wise
(per-row global scale) activation variant.

Four-over-six encodes every FP4 block twice — the standard map-to-6 encoding
and a map-to-4 encoding whose E4M3 block scale is expanded by exactly 1.5x —
and stores the candidate (codes and scale) with the lower MAE/MSE
dequantization error; ties select map-to-6. The global scale bound is reduced
to 256 by default so the 1.5x expansion cannot clip E4M3 (256·1.5 = 384 < 448).
The recipe applies to forward-GEMM operands only (never gradients) and targets
RL / post-training flows (no RHT, no stochastic rounding).

New in nvfp4_training/four_over_six.py:

• four_over_six_quantize(x, global_amax, *, block="1x16"|"16x16", err_mode="mae"|"mse", e4m3_scale_bound=256|448)
  — pure-PyTorch quantizer, transcribed operation-for-operation from the
  reference CUDA kernels (quantize_4over6_nvfp4.cuh). A scalar global_amax
  selects per-tensor scaling; a (rows,) vector selects the row-wise variant
  (1x16 blocks only — a row-scaled tensor has no columnwise form).
• four_over_six_mm / four_over_six_linear — training autograd function:
  • fprop: activations 1x16 four-over-six (optionally row-scaled), weights
    16x16 four-over-six, via scaled_mm (BlockWise1x16 + TensorWise).
  • backward (per-tensor): standard-NVFP4 RTNE gradients against the saved
    columnwise four-over-six activation/weight codes.
  • backward (row-scaled): bf16 GEMMs on saved operands — the quantized wgrad
    operand cannot exist because per-row scales do not transpose.
• NVFP4FourOverSixLinear — drop-in nn.Linear.

Two load-bearing numeric details (documented in the module docstring):

• the four-over-six block-scale association is (block_amax / 6) * S_enc,
  which rounds differently from the standard path's
  block_amax * (S_enc * (1/6)) on a fraction of blocks;
• dividing by a python scalar in torch lowers to a reciprocal multiply
  (double rounding); the two divisions on the selection path use tensor
  denominators to get true correctly-rounded FP32 division.

Validation (GB200, SM100)

• Bitwise parity with the reference CUDA kernels: 384/384 comparisons
  (packed FP4 codes, E4M3 scale bytes, FP32 global amaxes; atol=rtol=0)
  across shapes {128x128, 256x512, 2048x1024, 304x304} × {bf16, fp32} ×
  {MAE, MSE} × {bound 256, 448} × {1x16, 16x16} × {per-tensor, row-scaled} ×
  {rowwise, columnwise}, plus a 64-seed stress of the tightest case.
• Linear-level (same operands through both stacks): fwd/dgrad bitwise at
  512x1024x768; elsewhere ≥ 92 dB SQNR — GEMM reduction-order only, ~76 dB
  below the recipe's own quantization error (~17 dB vs bf16).
• test/prototype/moe_training/nvfp4_training/test_four_over_six.py:
  32 passed, 4 skipped (includes an optional bitwise-parity test that runs
  when TransformerEngine is importable).

Notes

• Pure-PyTorch implementation (compile-friendly); a fused Triton fast path is
  future work.
• Companion torchtitan converter PR: Add NVFP4 four-over-six converter torchtitan#8
  (NVFP4FourOverSixLinearConverter).

CuTe DSL kernel (SM100)

Commit 2 adds a CuTe DSL fast path: one kernel behind
torchao::four_over_six_quantize_cutedsl (custom op + fake impl, so
four_over_six_mm stays traceable), dispatched from inside
four_over_six_quantize when the input is eligible — CUDA bf16/fp32,
contiguous, C % 64 == 0 — and silently falling through to the pure-PyTorch
body otherwise. The PyTorch body is unchanged and remains the fallback and
the bitwise oracle; the public API and the titan-facing contract are
untouched.

Design. One (128, 64) input tile per CTA (TMA G2S in, TMA S2G packed
codes out, u32-vectorized scale stores), 128 threads with one tile row per
thread, so lanes 0-15 / 16-31 of each warp hold 16 consecutive rows — exactly
the 16-lane segments the 16x16 (2D) mode reduces over. Every load-bearing
rounding is pinned with inline PTX so no compiler lowering can change a bit:

• real div.rn.f32 for every division — the (block_amax / 6) * S_enc
  association, S_enc/S_dec, the encode reciprocals
  min(1/(scale·S_dec), FLT_MAX), and the error denominator 6·bound
  (never a reciprocal multiply);
• cvt.rn.satfinite.e4m3x2.f32 scale casts with exact
  cvt.rn.f16x2.e4m3x2 decode-back;
• cvt.rn.satfinite.e2m1x2.f32 FP4 casts (NaN → +6) and the exact
  cvt.rn.f16x2.e2m1x2 error-path dequant;
• NaN-dropping max.f32 block amaxes (all-NaN group → amax 0),
  NaN-dropping min.f32 for the 448 caps;
• per-group dequant error accumulated strictly sequentially in element
  order 0..15 with scalar FP32 RN adds; 16x16 tiles reduce their 16 row
  errors with the exact width-16 shuffle-down tree
  (((e0+e8)+(e4+e12))+((e2+e10)+(e6+e14))) + (((e1+e9)+(e5+e13))+((e3+e11)+(e7+e15)))
  broadcast from the segment base; strict err4 < err6 selection (ties and
  NaN errors pick map-to-6).

SASS/PTX audit: 142 div.rn.f32 in the 1x16 variant (the exact expected
count), identical FFMA counts between MAE and MSE variants (all from the
div.rn expansions — no mul+add contraction in the error path; the MSE
diff*diff is an inline mul.rn.f32).

Bitwise validation (GB200).

• DSL vs pure-PyTorch oracle: 1920/1920 bitwise (atol=rtol=0, uint8
  views of codes and scales) over {mae, mse} × {256, 448} × {1x16, 16x16} ×
  {per-tensor, row-scaled (1x16)} × {bf16, fp32} ×
  shapes {128x256, 256x512, 64x1024, 2048x7168, 128x320} × 16 seeds.
• Special values, all bitwise vs the oracle: zeros, ±Inf injections,
  subnormals, near-448 block amaxes, wide dynamic range (2^-20..2^20),
  amax == 0 rows/tensors, amax == inf.
• TransformerEngine parity: the in-tree 12-case TE test passes with the DSL
  path asserted directly against quantize_4over6_nvfp4.cuh output
  (12/12 bitwise, TE devel container).
• One documented exception: NaN inputs follow the TE kernel semantics
  (NaN-dropping fmaxf block amax; NaN encodes to +6), which the pure-torch
  body cannot reproduce because torch.amax propagates NaN into the block
  scales of NaN-containing groups. NaN-free groups of a NaN-containing tensor
  remain bitwise. Pinned by test_cutedsl_nan_semantics.

Bench (benchmarks/prototype/moe_training/nvfp4_training/bench_cutedsl_four_over_six_quantize.py,
GB200, SM clock 1200 MHz under load — application-clocks cap, 1965 MHz max
boost — HBM 4000 MHz; bf16 input, mae/256; bitwise-checked before timing):

shape	block	row_scaled	cutedsl µs	torch µs	speedup	cutedsl GB/s	torch GB/s
(8192, 2048)	1x16	no	127.0	3770.4	29.7x	338.5	11.4
(8192, 2048)	1x16	yes	127.0	3782.7	29.8x	338.5	11.4
(8192, 2048)	16x16	no	131.1	3719.2	28.4x	328.0	11.6
(8192, 7168)	1x16	no	409.6	11724.8	28.6x	367.3	12.8
(8192, 7168)	1x16	yes	413.7	11776.0	28.5x	363.7	12.8
(8192, 7168)	16x16	no	423.9	11459.1	27.0x	354.9	13.1
(32768, 2048)	1x16	no	464.9	13295.6	28.6x	369.9	12.9
(32768, 2048)	1x16	yes	470.0	13354.0	28.4x	365.9	12.9
(32768, 2048)	16x16	no	481.3	12988.5	27.0x	357.3	13.2
(32768, 7168)	1x16	no	1572.2	44493.9	28.3x	382.8	13.5
(32768, 7168)	1x16	yes	1592.3	44663.8	28.1x	378.0	13.5
(32768, 7168)	16x16	no	1631.2	43358.7	26.6x	369.0	13.9
(131072, 2048)	1x16	no	1793.1	50679.8	28.3x	383.6	13.6
(131072, 2048)	1x16	yes	1815.6	50868.2	28.0x	378.9	13.5
(131072, 2048)	16x16	no	1860.6	49388.0	26.5x	369.7	13.9
(131072, 7168)	1x16	no	6222.9	174026.0	28.0x	386.9	13.8
(131072, 7168)	1x16	yes	6300.8	174654.0	27.7x	382.1	13.8
(131072, 7168)	16x16	no	6460.5	169477.0	26.2x	372.7	14.2

The v1 kernel is correctness-first (single-stage, no pipelining) and is
compute-bound on the exact div.rn chains; a pipelined/ILP pass can raise
GB/s further if the quantizer ever shows up in profiles.

Follow-up fusions (not wired here). The columnwise pass is still
caller-side (x.t().contiguous()); a fused row+col kernel is v2. On the GEMM
side, cuDNN frontend ships prebuilt CuTe-DSL NVFP4 wrappers
(gemm_amax_wrapper_sm100, gemm_swiglu_wrapper_sm100, and grouped-MoE
variants, e4m3 vec16 scales) verified working on GB200 — a future fusion can
fold the next layer's amax (and the MLP activation) into the GEMM and feed
this quantizer directly.

🤖 Generated with Claude Code