[PyTorch] MXFP4 weight QAT on MXFP8 and FP8 block-scaling recipes

[xiuhu17]

Description

Weight-only MXFP4 quantization-aware training: 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.

The only new numerics is the fused projection bf16/fp32 -> MXFP4 grid -> bf16/fp32 (tex.mxfp4_fake_quantize: per 1x32 block, scale 2^clamp(ceil(log2(amax/6)), -126, 125) derived from the amax bit pattern, RTNE onto E2M1, straight-through-estimator gradient). All MXFP8 and 128x128 blockwise FP8 encoding/decoding of the projected weight runs through the existing quantizer pipeline unchanged. Since every projected value is a grid point p * 2^e, the existing encodings are decoded-value exact: MXFP8 rowwise unconditionally, 128x128 blockwise for tile scale spreads up to 2^14.

Recipes: MXFP4QATMXFP8BlockScaling(MXFP8BlockScaling) and MXFP4QATFloat8BlockScaling(Float8BlockScaling); QAT can also be enabled on the stock host recipes via the mxfp4_qat_weights field (NVTE_MXFP4_QAT=1). backward_override keeps its base-recipe semantics. Also fixes ptx::exp2f(e8m0) to decode UE8M0 code 0 as 2^-127 and code 255 as NaN.

Test results (B200, tests/pytorch/test_mxfp4_qat.py)

Test	Compared quantities	Corpus	Result
CUDA kernel vs composite-torch reference	raw bits of the projected weight	all 65536 bf16 patterns; fp32 bit fuzz; edges: +-0, bf16 subnormal ladder (2^-127 ... 2^-133), 2^-128 tie, every RTNE midpoint, 1.52^k +-1ulp triples, saturation 62^125, non-finite block poisoning, misaligned storage-offset views	bitwise equal
CUDA kernel vs independent fp64 RTNE oracle	decoded values	same corpora (finite blocks)	bitwise equal
--use_fast_math rebuild of the same kernel source	raw bits vs normal build	same edge corpora incl. the 2^-127 FTZ-critical block	identical
CUDA kernel vs TileKernels (independent deployment MXFP4 implementation, power-of-two scales)	projected values, bitwise	bf16-exhaustive (2036/2048 finite below-saturation blocks); fp32 edges + fuzz; full edge corpus	bitwise equal on the shared domain; divergence only for amax > 6*2^125 and the non-finite policy
MXFP8 rowwise raw-encode losslessness	payload + scale bytes decoded independently	incl. the 2^-127 grid point (UE8M0 code 0)	exact
UE8M0 extreme scale codes through real dequantize	planted codes 0 / 255	code 0 -> 2^-127, code 255 -> NaN	exact
128x128 blockwise exactness boundary	per-element E4M3-lattice prediction vs actual quantizer	spread d = 0..16 x full signed payload grid	exact through 2^14, first loss at 2^15
End to end	fwd/dgrad/wgrad with dgrad-lineage discriminators	Linear + GroupedLinear, 2 recipes x backward_override {None, dequantized, high_precision} x fused/unfused wgrad; recipe-switch cache invalidation; loud rejection of unsupported surfaces (ops API, fused Userbuffers/grouped-MLP, fp16, primary FP8 weights)	pass
Perf, 8192x8192 bf16	kernel vs torch reference	--	67 us (~4.0 TB/s) vs 3.1 ms

Known limitations

• 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).
• fp16 weights and fp16 activation/dequantize dtypes are rejected: the MXFP4 grid exceeds fp16 range.
• Distributed (FSDP/TP) + real-optimizer master-step end-to-end verification and a fused projection+quantize kernel are future work.

Type of change

• New feature (non-breaking change which adds functionality)

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

[greptile-apps]

Greptile Summary

The PR adds weight-only MXFP4 quantization-aware training on MXFP8 and FP8 block-scaling recipes.

• Adds the fused MXFP4 fake-quantization kernel and PyTorch binding.
• Adds QAT recipe configuration and integrates projected weights with module and fusible-op paths.
• Adds recipe/cache invalidation, distributed-weight handling, and unsupported-configuration guards.
• Corrects UE8M0 decoding at the extreme scale codes.
• Adds numerical, end-to-end, FSDP2, and generalized tensor-parallel coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously reported cached-weight issue is resolved because the skip-update flag now reaches the in-place quantization operation, which preserves the existing workspace on non-first microbatches.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/base.py	Integrates MXFP4 projection with weight-workspace caching and correctly forwards the CUDA-graph skip-update flag to in-place quantization.
transformer_engine/pytorch/mxfp4_qat.py	Exposes the MXFP4 fake-quantization operation through a straight-through autograd wrapper.
transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh	Implements fused per-block MXFP4 projection and E2M1 rounding.
transformer_engine/common/recipe/init.py	Adds MXFP4-QAT recipe options and specialized host-recipe variants.
tests/pytorch/test_mxfp4_qat.py	Adds extensive numerical, cache, backward-override, and end-to-end QAT coverage.
tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_modules.py	Verifies module-level QAT behavior and gradient routing under FSDP2.
tests/pytorch/distributed/fsdp2_tests/run_gtp_mxfp4_qat.py	Verifies supported distributed-weight paths and explicit rejection of unsupported combinations.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  W["Master BF16/FP32 weight"] --> QAT["MXFP4 fake quantization<br/>1x32 E2M1 grid"]
  QAT --> P["Projected high-precision weight"]
  P --> Host["Existing MXFP8 or FP8 block quantizer"]
  Host --> GEMM["Existing quantized GEMM"]
  GEMM --> FWD["Forward output"]
  GEMM --> BWD["Backward using configured override"]
  BWD --> STE["Straight-through gradient"]
  STE --> W

Loading

_{Reviews (7): Last reviewed commit: "update" | Re-trigger Greptile}

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.