[nvfp4 training][rl] Add a CuTe DSL fast-path kernel for four-over-six quantize - #4852
Open
wolfcomos wants to merge 2 commits into
Open
[nvfp4 training][rl] Add a CuTe DSL fast-path kernel for four-over-six quantize#4852wolfcomos wants to merge 2 commits into
wolfcomos wants to merge 2 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>
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/4852
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. |
meta-cla
Bot
added
the
CLA Signed
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 — this PR's diff includes its commit; the change to review here is the top commit. Adds this series' one minimal kernel prototype: the quantizer is the recipe's hot op — it runs on every activation and weight of every forward — so it gets a CuTe DSL fast-path kernel, dispatched automatically for eligible inputs (CUDA bf16/fp32, SM100+, contiguous, C % 64 == 0). Ineligible inputs silently fall through to the pure-PyTorch reference; the CuTe DSL runtime is assumed installed wherever the gate passes (no package probing in the dispatch).
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, and the width-16 error-reduction order is pinned (clamped shuffle-down tree), so no compiler lowering choice can change a rounding and the kernel is insulated from cutlass-dsl wrapper churn. One documented divergence: NaN inputs take NaN-dropping block amaxes and encode to +6 codes, while
torch.amaxpropagates NaN into the reference's scales. Compile caching keys on every shape/dtype/knob that changes codegen, and atorch.librarycustom op with a fake impl keepstorch.compiletracing through the dispatch.Numerics / perf: 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; 26–30x over the reference body at training shapes (bench file included). Further kernel work (fused dequantize, grouped variants) lands in future PRs.
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 ← this PR ← #4853.