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

[xiuhu17]

Description

Weight-only MXFP4 quantization-aware training on the MXFP8 and FP8 block-scaling (128x128) recipes: training sees exactly the deployment weight (E2M1 payloads, 1x32 power-of-two UE8M0 scales) while activations and gradients keep the base recipe untouched. Supersedes #3263 (closed; its commits are included here).

Motivation: the DeepSeek-V4 technical report observes that FP4-to-FP8 dequantization is lossless -- E4M3's wider exponent range fully absorbs the fine-grained 1x32 scales whenever the scale spread inside an FP8 quantization block stays under a threshold -- letting MXFP4 QAT reuse the existing FP8 training pipeline.

Two weight-conversion modes:

1. Projection (default): fused 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, those encodings are decoded-value exact (MXFP8 rowwise unconditionally; 128x128 blockwise up to tile scale spread 2^14).
2. Direct (mxfp4_qat_direct field / NVTE_MXFP4_QAT_DIRECT=1): converts the MXFP4 decomposition straight into the host FP8 representation, skipping the bf16 re-encode. MXFP8 rowwise = fixed-shift-6 lift (deployment-canonical, invertible bytes; scale clamps at UE8M0 code 0 with exact payload absorption); MXFP8 columnwise = fresh 32x1 quantization with the host amax rule (dgrad bytes match mode 1 bit for bit); blockwise = per-tile exponent folding (tile scale 2^(max fp4 exp - 6) clamped at 2^-127, RTNE onto E4M3 incl. subnormals, no device-side assert).

API: MXFP4QATMXFP8BlockScaling / MXFP4QATFloat8BlockScaling recipe subclasses, or the mxfp4_qat_weights field on the stock hosts (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.

Weight dataflows

MXFP8 host:

backward_override	forward	stored for backward	backward (dgrad)
dequantized	bf16 -> mxfp4(row) -> mxfp8(row)	mxfp8(row)	dequantize(stored row) -> bf16/fp32 -> GEMM
high_precision	bf16 -> mxfp4(row) -> mxfp8(row), feeds the fwd GEMM only	original bf16 weight	stored bf16 weight -> GEMM
None	bf16 -> mxfp4(row) -> mxfp8(row) for fwd; mxfp8(col) built from the same MXFP4 decomposition	mxfp8(col)	stored mxfp8(col) -> GEMM

FP8 blockwise (128x128) host -- one square encoding, no row/col distinction:

backward_override	forward	stored for backward	backward (dgrad)
dequantized	bf16 -> mxfp4(row) -> fp8(128x128)	fp8(128x128)	dequantize(stored) -> bf16/fp32 -> GEMM
high_precision	bf16 -> mxfp4(row) -> fp8(128x128), feeds the fwd GEMM only	original bf16 weight	stored bf16 weight -> GEMM
None	bf16 -> mxfp4(row) -> fp8(128x128)	fp8(128x128)	its exact transpose -> GEMM

STE: identity gradient into the weight parameter (.grad and fused main_grad verified end to end), so updates propagate directly to the FP32 master weights. fp16 is rejected (the MXFP4 grid exceeds fp16 range).

Test results (B200)

tests/pytorch/test_mxfp4_qat.py -- the projection and the host encodings:

Test	Result
Projection kernel vs composite-torch reference and an independent fp64 RTNE oracle, raw bits: all 65536 bf16 patterns, fp32 bit fuzz, edges (+-0, 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 views)	bitwise equal
Same kernel source rebuilt with --use_fast_math vs normal build	identical bits
MXFP8 rowwise raw payload/scale bytes decode back to the projected weight (incl. 2^-127 at UE8M0 code 0); planted extreme scale codes through real dequantize (0 -> 2^-127, 255 -> NaN)	exact
128x128 exactness boundary, per-element E4M3-lattice prediction vs actual quantizer, spread d = 0..16 x full payload grid	exact through 2^14, first loss at 2^15
End to end: Linear + GroupedLinear, 2 recipes x 3 backward_override modes x fused/unfused wgrad, dgrad-lineage discriminators; recipe-switch cache invalidation; loud rejection of unsupported surfaces (ops API, fused Userbuffers/grouped-MLP, fp16, primary FP8 weights)	pass
Perf 8192x8192 bf16	67 us (~4.0 TB/s) vs 3.1 ms torch reference

tests/pytorch/test_mxfp4_qat_direct.py -- the direct converters:

Test	Result
direct vs projection mode, end to end: fwd outputs, dgrad, wgrad across 2 recipes x 3 backward_override modes x cache miss/hit/update paths	bitwise identical
Direct row + blockwise kernels vs composite-torch reference, payload and scale bytes, on the edge corpora above plus the code-0 absorption domain and the spread sweep; normal and --use_fast_math builds	bitwise equal
Direct columnwise bytes vs the host quantizer's	bitwise equal
2^-127 (UE8M0 scale code 0) through the real GEMM paths (MXFP8 and the Blackwell blockwise-to-MXFP8 route)	bit-exact
Perf 8192x8192 bf16	row 87 us, blockwise 418 us vs ~3 ms torch reference

Cross-validation against TileKernels (the independent deployment MXFP4 implementation, power-of-two scale mode): the projection values match bit-for-bit on the bf16-exhaustive corpus (2036/2048 finite below-saturation blocks) and fp32 edge/fuzz corpora; direct-row payload bytes and scale codes, and direct-blockwise payload bytes and tile scale codes, match TileKernels' per_token_cast / per_block_cast_lossless device kernels on the same corpora plus spread-banded and zero/2^-127/negative-zero tiles. Divergence only outside the shared domain (amax > 6*2^125, non-finite policy, spreads > 2^11 where TileKernels device-asserts while this implementation stays exact to 2^14). The tests carrying these comparisons depend on the external TileKernels package and were dropped from this PR in 2f40ca8 and 7a4e47e.

Known limitations

• For amax > 6*2^125 the projection 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.
• 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)