[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

Adds MXFP4 weight-only quantization-aware training while retaining the host MXFP8 or FP8 block-scaling pipeline.

• Introduces fused MXFP4 fake quantization with a straight-through gradient.
• Adds recipe configuration, PyTorch bindings, module and fusible-op integration, workspace invalidation, and unsupported-path validation.
• Corrects UE8M0 extreme-code decoding and adds extensive numerical, distributed, cache, backward-policy, and integration coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh	Implements the fused CUDA projection from BF16/FP32 weights onto the block-scaled MXFP4 grid.
transformer_engine/pytorch/mxfp4_qat.py	Exposes fake quantization through a validated PyTorch straight-through autograd operation.
transformer_engine/pytorch/module/base.py	Integrates QAT projection into cached weight quantization and preserves conditional cache-update handling.
transformer_engine/common/recipe/init.py	Adds MXFP4-QAT recipe configuration for MXFP8 and FP8 block-scaling hosts.
transformer_engine/common/util/ptx.cuh	Corrects UE8M0 decoding behavior for the minimum exponent and NaN sentinel codes.
tests/pytorch/test_mxfp4_qat.py	Adds broad numerical, recipe, cache, backward-policy, and end-to-end QAT coverage.
tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mxfp4_qat_ops.py	Adds FSDP2 parity coverage for QAT-enabled linear and grouped operations.

Sequence Diagram

sequenceDiagram
  participant Train as Training module
  participant QAT as MXFP4 fake quantizer
  participant Host as Host FP8/MXFP8 quantizer
  participant GEMM as GEMM
  Train->>QAT: High-precision master weight
  QAT-->>Train: MXFP4-grid projected weight
  Train->>Host: Quantize projected weight
  Host-->>GEMM: Existing FP8/MXFP8 representation
  GEMM-->>Train: Forward and backward results
  Train-->>Train: STE gradient updates master weight

Loading

_{Reviews (6): Last reviewed commit: "[PyTorch] MXFP4 weight QAT on MXFP8 and ..." | Re-trigger Greptile}

[greptile-apps]

Skip-update flag is dropped

When a non-first CUDA-graph microbatch requests cached-weight reuse, this direct branch calls mxfp4_qat_direct_update_ without forwarding the skip flag, so it unconditionally rewrites the workspace and the replayed GEMM can consume different weights than the cached first-microbatch weights.

Knowledge Base Used: PyTorch Fused Modules (transformer_engine/pytorch/module)

[greptile-apps]

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