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

[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

This PR adds MXFP4 weight-only quantization-aware training on MXFP8 and block-scaled FP8 hosts.

• Adds the fused BF16/FP32-to-MXFP4 fake-quantization kernel, C API, PyTorch binding, and straight-through autograd wrapper.
• Adds MXFP4-QAT recipe classes and recipe-field configuration.
• Integrates projected-weight handling into fused PyTorch modules while rejecting unsupported operator and fused paths.
• Corrects UE8M0 extreme-code decoding and adds extensive numerical and end-to-end tests.

Confidence Score: 3/5

The PR is not yet safe to merge because Linear can reconstruct its backward weight under a different recipe than the one used for forward.

Linear does not retain the forward MXFP4-QAT state and instead queries mutable global recipe state when rebuilding FSDP2 or distributed weight workspaces during backward, so leaving the forward autocast context can cause dgrad to use an incorrectly projected weight.

Files Needing Attention: transformer_engine/pytorch/module/linear.py

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/linear.py	Integrates QAT projection into Linear weight handling, but backward reconstruction still derives QAT state from the ambient recipe rather than forward state.
transformer_engine/pytorch/module/layernorm_linear.py	Adds MXFP4 projection to LayerNormLinear weight preparation and backward workspace reconstruction.
transformer_engine/pytorch/module/layernorm_mlp.py	Adds projected-weight handling for both MLP weights while preserving the existing fused backward structure.
transformer_engine/pytorch/module/base.py	Extends shared weight-workspace validation and quantization to account for MXFP4-QAT recipe state.
transformer_engine/pytorch/mxfp4_qat.py	Introduces the validated PyTorch fake-quantization wrapper and identity-gradient estimator.
transformer_engine/common/cast/mxfp4/fake_quantize_mxfp4.cuh	Implements the fused blockwise MXFP4 projection CUDA kernel.
transformer_engine/common/recipe/init.py	Adds MXFP4-QAT configuration and host-recipe subclasses.
tests/pytorch/test_mxfp4_qat.py	Adds broad numerical, recipe, integration, and unsupported-path coverage.

Sequence Diagram

sequenceDiagram
  participant User
  participant Module as PyTorch fused module
  participant QAT as MXFP4 fake quantizer
  participant Host as MXFP8 / FP8 block quantizer
  participant GEMM
  User->>Module: Forward under MXFP4-QAT recipe
  Module->>QAT: Project master weight to MXFP4 grid
  QAT-->>Module: BF16/FP32 projected weight
  Module->>Host: Encode using host recipe
  Host-->>GEMM: Quantized weight workspace
  GEMM-->>User: Forward output
  User->>Module: Backward
  Module->>Host: Reconstruct weight workspace if needed
  Host->>GEMM: Dgrad weight operand

Loading

_{Reviews (7): Last reviewed commit: "Drop the TileKernels cross-check tests" | Re-trigger Greptile}