[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

This PR adds direct MXFP4-to-FP8 weight conversion for MXFP4 QAT.

• Adds CUDA kernels and PyTorch bindings for rowwise MXFP8 and 128x128 blockwise conversion.
• Adds recipe flags and module weight-cache integration for selecting direct conversion.
• Adds MXFP4 fake quantization, direct-conversion helpers, and extensive numerical and end-to-end tests.
• Explicitly rejects unsupported operations-API and fused Userbuffers paths.

Confidence Score: 3/5

This PR should not merge until direct cached-weight updates honor CUDA-graph requests to preserve the existing workspace.

The new direct cache-hit branch bypasses the noop flag used by every existing quantized update path, allowing non-first CUDA-graph microbatches to overwrite weights that the replay contract requires them to reuse.

Files Needing Attention: transformer_engine/pytorch/module/base.py; transformer_engine/pytorch/mxfp4_qat_direct.py

Important Files Changed

Filename	Overview
transformer_engine/common/cast/cast.cu	Adds validated C entry points that dispatch MXFP4 fake quantization and direct FP8 conversion kernels.
transformer_engine/common/cast/mxfp4/direct_convert_mxfp4.cuh	Implements rowwise MXFP8 and 128x128 blockwise direct-conversion CUDA kernels.
transformer_engine/pytorch/mxfp4_qat_direct.py	Builds and updates direct-conversion tensor workspaces, but its updater cannot honor conditional cache-update suppression.
transformer_engine/pytorch/module/base.py	Integrates MXFP4 QAT into shared weight caching, but the direct cache-hit branch drops the CUDA-graph skip-update contract.
transformer_engine/common/recipe/init.py	Adds recipe fields and predicates for MXFP4 QAT and direct conversion.
tests/pytorch/test_mxfp4_qat_direct.py	Adds broad converter and end-to-end parity coverage, without covering skipped cached updates during CUDA-graph replay.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  W[BF16/FP32 master weight] --> Q[MXFP4 decomposition]
  Q --> R[Direct MXFP8 rowwise conversion]
  Q --> C[Composite MXFP8 columnwise conversion]
  Q --> B[Direct FP8 128x128 blockwise conversion]
  R --> WS[Cached quantized workspace]
  C --> WS
  B --> WS
  WS --> G[Forward/backward GEMM]
  S[CUDA-graph skip-update flag] -. must gate updates .-> WS

Loading

_{Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | 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)