MXFP4 weight QAT recipes on MXFP8 and FP8 block-scaling hosts

[xiuhu17]

Overview

Weight-only MXFP4 quantization-aware training for DeepSeek-v4 MoE deployment: weights are projected onto the MXFP4 grid (E2M1 payloads, 1x32 power-of-two UE8M0 scales) before the host recipe quantizes them, so training sees exactly the deployment weight while activations and gradients keep the base recipe untouched.

Two recipes subclass their hosts and reuse the full TE pipeline (quantizers, QTensor layouts, caching, backward_override):

• MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling)
• MXFP4QATFloat8BlockScaling(Float8BlockScaling) (128x128 weight tiles)

Design

One hook in quantize_weight applies mxfp4_fake_quantize (the fused bf16 -> mxfp4(row) -> dequantize -> bf16 projection) right before the host weight quantizer; everything downstream is stock TE. The projection output is exactly representable in bf16 (E2M1 needs 2 mantissa bits, scales are powers of two), so the handoff is lossless by construction:

• MXFP8 rowwise encoding of the projected weight: bitwise lossless (raw payload/scale verified)
• 128x128 blockwise encoding: bitwise lossless within each tile's E4M3 headroom (full payload grid exact through scale spread 2^14, enumerated)
• MXFP8 columnwise (32x1, dgrad under backward_override=None): bounded re-quantization

backward_override weight semantics (verified by dgrad discriminators through real GEMMs): None -> row+col built from the same projected weight, dgrad consumes col fp8; dequantized -> dgrad dequantizes the stored rowwise encoding (bit-exact projected weight in bf16); high_precision -> dgrad uses the original unprojected parameter. Weight gradients flow through an identity STE to .grad/fused main_grad.

Scale contract matches the TileKernels deployment path: floor 2^-126, cap 2^125, integer-bit amax/exponent derivation and non-FTZ PTX so results are bit-identical under --use_fast_math builds.

Implementations

Three bit-identical implementations dispatched via NVTE_MXFP4_QAT_IMPL (auto/cuda/cute_dsl/torch): the CUDA kernel (common/cast/mxfp4/fake_quantize_mxfp4.cuh, tex.mxfp4_fake_quantize), a CuTe DSL kernel, and a PyTorch reference. B200, 8192x8192 bf16: CuTe DSL 59us @ 4.5 TB/s, CUDA 67us @ 4.0 TB/s, torch 3052us.

Correctness

28-test suite, all passing on B200, including a four-way bitwise matrix (TileKernels torch reference vs CuTe DSL vs CUDA vs torch reference) over every edge domain (+-0 with sign, 2^-127/2^-128/2^-133 subnormals, satfinite cap, RTNE midpoints, scale-threshold ulp triples, inf/NaN block poisoning, misaligned views) and every quant/dequant hop (fake-quant bf16+fp32, mxfp8 row raw encode/dequant, mxfp8 col, blockwise 128x128 row+col, bf16+fp32 dequant targets), plus all 65536 bf16 bit patterns exhaustively, fp32 raw-bit fuzz, a fast-math dual-build bit-parity gate, an identity-STE gradient check, and a 24-config e2e matrix (Linear+GroupedLinear x 2 recipes x 3 overrides x fused/unfused wgrad).

Cross-backend findings pinned in tests: TileKernels' e4m3 casts use a hardcoded 1e-4 amax floor (its e2m1 path uses the true 6*2^-126 deployment clamp); TE's mxfp8 columnwise representation canonicalizes -0 to +0.

Guards

Surfaces that bypass the weight hook reject QAT recipes loudly (te.ops, Userbuffers fused ops, fused grouped MLP, quantized primary weights); FSDP2/GTP backward rematerialization applies the projection; cached weight workspaces are invalidated on base<->QAT recipe switches; fp16 weights and fp16 activation/dequantize dtypes are rejected (the grid exceeds fp16 range).

Includes the ptx.cuh exp2f UE8M0 code-0/255 fix (also submitted upstream as NVIDIA#3262).

Known limitations

• The prebuilt wheel's software dequantize flushes the 2^-127 grid point until TE is rebuilt with the exp2f fix (raw encoding and hardware GEMM are exact; tests auto-detect the build state).
• For amax > 6*2^125 the fake-dequant saturates at the cap while packed deployment moves to scale 2^126 (unmaterializable in bf16/fp32; pinned by test, unreachable for real weights).
• Distributed (FSDP/TP) + real-optimizer master-step end-to-end verification and a fused projection+quantize kernel are future work.

[xiuhu17]

Re-targeting to NVIDIA/TransformerEngine main.