[nvfp4 training][rl] Add the four-over-six grouped GEMM and dispatcher wiring - #4853
Open
wolfcomos wants to merge 3 commits into
Open
[nvfp4 training][rl] Add the four-over-six grouped GEMM and dispatcher wiring#4853wolfcomos wants to merge 3 commits into
wolfcomos wants to merge 3 commits into
Conversation
Four-over-six is an adaptive NVFP4 block-scaling recipe: every quantization block is encoded twice — the standard map-to-6 encoding and a 1.5x-scale map-to-4 candidate whose denser FP4 grid lowers error for blocks with mass below the amax — and the lower-error candidate is stored. It is part of the NVFP4 RL training recipe roadmap (miles). This PR adds the recipe surface and a correctness-first reference implementation; optimized kernels land separately: - four_over_six_quantize: pure-PyTorch quantizer with per-tensor or row-scaled global scales and 1x16/16x16 blocks, every rounding step pinned so faster implementations can be validated bitwise against it. - nvfp4_dequantize: the standard NVFP4 decode (recipe-agnostic), a pure-PyTorch correctness helper — an optimized kernel is future work. - four_over_six_mm / four_over_six_linear: differentiable matmul with quantized, high_precision, and dequantized backward modes; the dequantized mode differentiates the quantized-forward function itself (the RL train/inference-consistency mode) while saving only 4-bit codes and scales. - NVFP4FourOverSixLinear: a stateless leaf module for the recipe, and recipe="default"|"four_over_six" on NVFP4TrainingConfig so quantize_() installs the right module. The existing NVFP4Linear is untouched: it owns RHT/SR state (sign-vector and seed buffers, TP machinery) that four-over-six deliberately has none of, so the recipes stay separate leaf modules behind one config. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The quantizer is the recipe's one hot op — it runs on every activation and weight of every forward — so this PR adds a minimal CuTe DSL kernel prototype for it, dispatched automatically for eligible inputs (CUDA bf16/fp32, SM100+, contiguous, C % 64 == 0); ineligible inputs silently fall through to the pure-PyTorch reference body. The CuTe DSL runtime packages are assumed present wherever the gate passes — the dispatch does no package probing. The kernel is bitwise-identical to the reference: the rounding-critical steps (correctly-rounded FP32 division, the FP4/E4M3 casts and their exact decodes) are emitted as raw inline PTX so no compiler lowering choice can change a rounding, and the width-16 error-reduction order is pinned (clamped shuffle-down tree), insulating the kernel from cutlass-dsl wrapper churn. Compile caching keys on every shape/dtype/knob that changes codegen, and compilation runs under the input's device context. One documented divergence: NaN inputs take NaN-dropping block amaxes and encode to +6 codes, where torch.amax propagates NaN. The kernel is wrapped in a CUDA-only torch.library custom op with a fake impl so torch.compile traces through the dispatch; direct op callers get the dispatch gate's and wrapper's checks mirrored as errors. The bench file measures the kernel against the pure-torch reference (26-30x at training shapes). Further kernel work (fused dequantize, grouped variants) lands in future PRs. Numerics: all 440 swept shape/knob configurations (block x err_mode x bound x scale granularity x dtype x data construction x shape) bitwise identical to the pure-torch reference on GB200. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…atcher Integrate the four-over-six recipe with MoE training: routed-expert layers get four_over_six_grouped_mm, built by reusing the existing pieces rather than adding kernels — the quantizer from the dense PRs (the CuTe DSL fast path dispatches automatically) and the existing F.scaled_grouped_mm for the GEMMs. - Per-tensor activations: each token group's amax expands to a per-row vector so the whole packed tensor quantizes in one call, bitwise identical to quantizing each group separately; one scaled_grouped_mm with per-group second-level scales runs the forward. - Expert weights quantize in one flattened call for both block shapes (the quantizer accepts a per-row amax with 16x16 blocks under a tile-uniform contract), one kernel launch regardless of expert count. - Row-scaled activations: one scaled_grouped_mm carrying the constant per-tensor factor, its bf16 output upcast and scaled by the raw per-row amaxes (the grouped GEMM emits bf16 only, one output rounding vs a dense per-group loop; tests pin it against a rounding-emulated loop oracle). - Backward is high-precision or dequantized grouped GEMMs (no quantized backward); dequantized differentiates the quantized forward itself from the saved 4-bit codes and scales. Wiring mirrors the Float8/MXFP8 recipes exactly: NVFP4FourOverSixTrainingOpConfig(TrainingOpBaseConfig) — a pytree-constant dataclass with explicit __eq__/__hash__ — plus an isinstance branch in _quantize_then_scaled_grouped_mm. The dispatcher branch owns the offs[-1] tail slice / zero-extend for padded token dispatchers that over-allocate A past the logical rows. Also adds an NVFP4 four-over-six section to the moe_training README. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
wolfcomos
requested review from
andrewor14,
jerryzh168 and
vkuzo
as code owners
August 31, 2026 07:24
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4853
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Aug 31, 2026
meta-cla
Bot
added
the
CLA Signed
This was referenced Aug 31, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #4851 / #4852 — this PR's diff includes their commits; the change to review here is the top commit. Integrates the four-over-six recipe with MoE training for routed-expert layers, reusing existing pieces rather than adding kernels: the #4852 quantizer (its CuTe DSL fast path dispatches automatically) and the existing
F.scaled_grouped_mmfor the GEMMs.scaled_grouped_mmwith per-group second-level scales runs the forward.scaled_grouped_mmcarrying the constant per-tensor factor in every group's slot, its bf16 output upcast and scaled by the raw per-row amaxes. The grouped GEMM emits bf16 only, which costs one output rounding vs per-group dense GEMMs; tests emulate that rounding in a loop oracle and compare bitwise (with a reduction-order SQNR fallback), and separately bound the rounding cost.high_precisionordequantizedgrouped GEMMs (four-over-six has no quantized backward);dequantizeddifferentiates the quantized forward itself from the saved 4-bit codes and scales.Wiring mirrors the Float8/MXFP8 recipes exactly:
NVFP4FourOverSixTrainingOpConfig(TrainingOpBaseConfig)— a pytree-constant dataclass with explicit__eq__/__hash__— plus anisinstancebranch in_quantize_then_scaled_grouped_mm. The dispatcher branch owns theoffs[-1]tail slice / zero-extend for padded token dispatchers that over-allocateApast the logical rows.quantize_()model conversion for the grouped config is future work; framework integrations drive the grouped GEMM dispatcher (see the torchtitan converters). Adds an NVFP4 four-over-six section to the moe_training README.Tests: 40 grouped tests (grouped-vs-dense forward parity, row-scaled fused-vs-oracle, backward bitwise references including empty groups and ragged padding, dispatcher round-trip with over-allocated tails, fullgraph compile in both scale granularities) plus the dense suite with the relaxed quantizer combinations, all green on GB200.
Reference
NVFP4 RL training recipe roadmap references from Ziang Li from humans&: https://humansand.ai/blog/nvfp4-rl, radixark/miles#615, NVIDIA/TransformerEngine#2972
Stack: #4851 ← #4852 ← this PR.