[nvfp4 training][rl] Add a CuTe DSL fast-path kernel for four-over-six quantize - #13
Closed
wolfcomos wants to merge 1 commit into
Closed
[nvfp4 training][rl] Add a CuTe DSL fast-path kernel for four-over-six quantize#13wolfcomos wants to merge 1 commit into
wolfcomos wants to merge 1 commit into
Conversation
This was referenced Aug 30, 2026
wolfcomos
force-pushed
the
4over6/ao1-dense
branch
from
August 31, 2026 02:00
7bcd854 to
f114077
Compare
wolfcomos
force-pushed
the
4over6/ao2-cutedsl
branch
from
August 31, 2026 02:01
e712a0e to
cd50f68
Compare
wolfcomos
force-pushed
the
4over6/ao1-dense
branch
from
August 31, 2026 03:09
f114077 to
3a10fef
Compare
wolfcomos
force-pushed
the
4over6/ao2-cutedsl
branch
3 times, most recently
from
August 31, 2026 04:37
5113820 to
7f1e7d4
Compare
wolfcomos
force-pushed
the
4over6/ao1-dense
branch
from
August 31, 2026 04:44
f8a0a5a to
e47975f
Compare
wolfcomos
force-pushed
the
4over6/ao2-cutedsl
branch
from
August 31, 2026 04:44
7f1e7d4 to
ff415d7
Compare
wolfcomos
force-pushed
the
4over6/ao1-dense
branch
from
August 31, 2026 06:35
e47975f to
8ac6288
Compare
wolfcomos
force-pushed
the
4over6/ao2-cutedsl
branch
from
August 31, 2026 06:35
ff415d7 to
1e847a8
Compare
wolfcomos
changed the title
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
force-pushed
the
4over6/ao1-dense
branch
from
August 31, 2026 07:01
8ac6288 to
f2f7647
Compare
wolfcomos
force-pushed
the
4over6/ao2-cutedsl
branch
from
August 31, 2026 07:02
1e847a8 to
2c69dbc
Compare
Owner
Author
|
Opened upstream as pytorch#4852 (same head). This fork draft stays as the staging copy. |
Owner
Author
|
Live upstream as pytorch#4852 — closing the staging draft. The branch stays: it is the upstream PR's head. |
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 #12. 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.
Stack: #12 ← this PR ← #14.