Repository navigation
Conversation
Quantizing a 2D tensor along both axes currently takes three kernels and two
full reads of the same tensor: the 1x32 CuTeDSL cast along N, the 32x1 CUDA cast
along M, and a standalone Triton K-groups swizzle of the colwise scales. Add a
CuTeDSL kernel that reads the input once and emits all four outputs, with both
scale tensors already in the blocked layouts their consumers expect. Byte
accounting drops from 6.125 to 4.0625 bytes per input element.
The kernel puts one CTA on each (128, 128) input tile, stages it into SMEM with
cp.async, and runs the rowwise and colwise passes sequentially at full CTA
width. The arithmetic is packed 16-bit throughout -- max.u16x2 on sign-masked
bf16 for amax, cvt.rn.satfinite.e4m3x2.bf16x2 for the cast -- and prmt.b32 byte
permutes keep the logically transposed colwise store fully coalesced. bf16 and
rceil only; both dimensions and every group size must be multiples of 128.
Exposed as mxfp8_quantize_2d_1x32_32x1_cutedsl, registered as a custom op with a
fake impl, alongside the two kernels it composes.
Uncovered tail. offs[-1] < total_M is supported and is the norm rather than a
corner case: a dropless MoE dispatcher sizes its buffer for worst-case padding
and fills only a prefix. Rows past offs[-1] are still quantized, but their
colwise scale bytes are zeroed to match the K-groups swizzle, which iterates
groups and returns new_zeros. Falling through to the placement math instead aims
the store into group 0's slab and races every n-tile CTA on one 512-byte line.
Bug-compatible where the two kernels it replaces disagree with each other, since
the point is to be a drop-in:
* A block whose amax is subnormal enough to reach E8M0 scale byte 0 gets a
2**127 multiplier on both outputs, matching the 254 - scale_biased
reciprocal both use.
* Inf: the rowwise kernel converts NOSAT, so Inf poisons the block to scale
255 and all-NaN data; the colwise CUDA kernel converts SATFINITE, so Inf
clamps to scale 254 with finite data. Both are reproduced.
The one deliberate divergence is NaN on the colwise output. The CUDA kernel
drops NaN from its amax entirely -- an all-NaN block still gets the scale its
finite elements imply -- and reproducing that would need a per-element NaN test
in the hottest loop in order to imitate a behavior that hides NaN corruption
rather than surfacing it. The fused kernel poisons instead, and the test asserts
the difference so it fails loudly if either side changes.
Tests compare raw bytes at atol=0/rtol=0 against the three-kernel composition,
not against to_mx, because being a drop-in for that composition is the contract.
They cover shapes and group counts, uncovered tails including the zero-fill,
degenerate and mixed-magnitude blocks, Inf/NaN, dirty destination buffers so a
never-written region cannot hide behind a zeroed page, determinism across
repeated launches, and input validation. The rowwise half is additionally
checked against the shared MXFP8 semantic contract.
The benchmark follows the existing 1x32 and 32x1 scripts, but times both sides
under CUDA-graph replay rather than eagerly. torchao's triton_op dispatch costs
on the order of 100us per call regardless of problem size and the baseline pays
it three times, which timed eagerly swamps every shape below roughly 0.5 GB and
reports 15x on small shapes -- launch overhead, not a kernel win. It also checks
that the kernel does not capture as an empty graph, which is what happens when a
CuTeDSL kernel ignores the caller's stream and which would read as a 100x win.
Geomean is 2.6x across the swept shapes, 2.0x on the largest, at up to 95% of
the measured SM bandwidth roof on a GB200 at 2062 MHz.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
protonu
requested review from
andrewor14,
jerryzh168 and
vkuzo
as code owners
September 14, 2026 20:47
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4895
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 1 Unrelated Failure, 1 Unclassified FailureAs of commit 1d4e8d3 with merge base ee0d7ce ( NEW FAILURE - The following job has failed:
UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
Hi @protonu! Thank you for your pull request. We require contributors to sign our Contributor License Agreement, and yours needs attention. You currently have a record in our system, but the CLA is no longer valid, and will need to be resubmitted. ProcessIn order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA. Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks! |
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
meta-cla
Bot
added
the
CLA Signed
|
Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks! |
sanketpurandare
added a commit
that referenced
this pull request
Oct 14, 2026
Performnace of this fused cute-dsl quantizer is better than cuda, a single API for dense and grouped mxfp8 linears, no need for scale rearrangements since swizzled layouts are directly output. AI-assisted implementation details: > Keep Protonu's original authored fused kernel from #4895 as the > parent commit and consolidate current-main correctness fixes in this PR. > Credit: Protonu (@protonu), original commit > 1d4e8d3; Aleksandar Samardzic's #4702 > independently addressed the single-axis FakeTensor metadata issue. > > Match the public/custom-op/fake argument order to surrounding APIs. None > offsets select dense placement in the same kernel, with compact blocked > scales and no synthetic endpoint or offset search. Tensor offsets retain > grouped placement, empty experts and uncovered-tail behavior. > > Correct the scale0/1 RCEIL cutoff to 448*2^-127, preserve NOSAT nonfinite > semantics, and compare both axes directly to to_mx with group unblocking. > Cover exact cutoff neighbors, mixed magnitudes, special values, empty groups, > dirty destinations, actual compile/fake contracts and device selection. > > Validate aligned destination metadata and nonoverlapping byte intervals. > Enable device assertions explicitly for grouped compilation; CuTeDSL strips > them by default. Reuse the endpoint validation loop for ordering and bounds. > Remove unreachable CTA work and redundant reciprocal conventions. > > Test plan: affected fused/single-axis CUDA tests, repository Ruff checks, > strict API documentation rendering, matched replay timings and traces, and > real TorchTitan FSDP/MXFP8 eager/capture consumer tests. Pull Request resolved: #4995
vkuzo
pushed a commit
that referenced
this pull request
Oct 15, 2026
…ta (#4995) * [moe_training] Add a fused 1x32 + 32x1 MXFP8 quantization kernel Quantizing a 2D tensor along both axes currently takes three kernels and two full reads of the same tensor: the 1x32 CuTeDSL cast along N, the 32x1 CUDA cast along M, and a standalone Triton K-groups swizzle of the colwise scales. Add a CuTeDSL kernel that reads the input once and emits all four outputs, with both scale tensors already in the blocked layouts their consumers expect. Byte accounting drops from 6.125 to 4.0625 bytes per input element. The kernel puts one CTA on each (128, 128) input tile, stages it into SMEM with cp.async, and runs the rowwise and colwise passes sequentially at full CTA width. The arithmetic is packed 16-bit throughout -- max.u16x2 on sign-masked bf16 for amax, cvt.rn.satfinite.e4m3x2.bf16x2 for the cast -- and prmt.b32 byte permutes keep the logically transposed colwise store fully coalesced. bf16 and rceil only; both dimensions and every group size must be multiples of 128. Exposed as mxfp8_quantize_2d_1x32_32x1_cutedsl, registered as a custom op with a fake impl, alongside the two kernels it composes. Uncovered tail. offs[-1] < total_M is supported and is the norm rather than a corner case: a dropless MoE dispatcher sizes its buffer for worst-case padding and fills only a prefix. Rows past offs[-1] are still quantized, but their colwise scale bytes are zeroed to match the K-groups swizzle, which iterates groups and returns new_zeros. Falling through to the placement math instead aims the store into group 0's slab and races every n-tile CTA on one 512-byte line. Bug-compatible where the two kernels it replaces disagree with each other, since the point is to be a drop-in: * A block whose amax is subnormal enough to reach E8M0 scale byte 0 gets a 2**127 multiplier on both outputs, matching the 254 - scale_biased reciprocal both use. * Inf: the rowwise kernel converts NOSAT, so Inf poisons the block to scale 255 and all-NaN data; the colwise CUDA kernel converts SATFINITE, so Inf clamps to scale 254 with finite data. Both are reproduced. The one deliberate divergence is NaN on the colwise output. The CUDA kernel drops NaN from its amax entirely -- an all-NaN block still gets the scale its finite elements imply -- and reproducing that would need a per-element NaN test in the hottest loop in order to imitate a behavior that hides NaN corruption rather than surfacing it. The fused kernel poisons instead, and the test asserts the difference so it fails loudly if either side changes. Tests compare raw bytes at atol=0/rtol=0 against the three-kernel composition, not against to_mx, because being a drop-in for that composition is the contract. They cover shapes and group counts, uncovered tails including the zero-fill, degenerate and mixed-magnitude blocks, Inf/NaN, dirty destination buffers so a never-written region cannot hide behind a zeroed page, determinism across repeated launches, and input validation. The rowwise half is additionally checked against the shared MXFP8 semantic contract. The benchmark follows the existing 1x32 and 32x1 scripts, but times both sides under CUDA-graph replay rather than eagerly. torchao's triton_op dispatch costs on the order of 100us per call regardless of problem size and the baseline pays it three times, which timed eagerly swamps every shape below roughly 0.5 GB and reports 15x on small shapes -- launch overhead, not a kernel win. It also checks that the kernel does not capture as an empty graph, which is what happens when a CuTeDSL kernel ignores the caller's stream and which would read as a 100x win. Geomean is 2.6x across the swept shapes, 2.0x on the largest, at up to 95% of the measured SM bandwidth roof on a GB200 at 2062 MHz. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> * [moe_training] Match MXFP8 dense and grouped contracts Performnace of this fused cute-dsl quantizer is better than cuda, a single API for dense and grouped mxfp8 linears, no need for scale rearrangements since swizzled layouts are directly output. AI-assisted implementation details: > Keep Protonu's original authored fused kernel from #4895 as the > parent commit and consolidate current-main correctness fixes in this PR. > Credit: Protonu (@protonu), original commit > 1d4e8d3; Aleksandar Samardzic's #4702 > independently addressed the single-axis FakeTensor metadata issue. > > Match the public/custom-op/fake argument order to surrounding APIs. None > offsets select dense placement in the same kernel, with compact blocked > scales and no synthetic endpoint or offset search. Tensor offsets retain > grouped placement, empty experts and uncovered-tail behavior. > > Correct the scale0/1 RCEIL cutoff to 448*2^-127, preserve NOSAT nonfinite > semantics, and compare both axes directly to to_mx with group unblocking. > Cover exact cutoff neighbors, mixed magnitudes, special values, empty groups, > dirty destinations, actual compile/fake contracts and device selection. > > Validate aligned destination metadata and nonoverlapping byte intervals. > Enable device assertions explicitly for grouped compilation; CuTeDSL strips > them by default. Reuse the endpoint validation loop for ordering and bounds. > Remove unreachable CTA work and redundant reciprocal conventions. > > Test plan: affected fused/single-axis CUDA tests, repository Ruff checks, > strict API documentation rendering, matched replay timings and traces, and > real TorchTitan FSDP/MXFP8 eager/capture consumer tests. Pull Request resolved: #4995 --------- Co-authored-by: protonu <pbasu@nvidia.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Contributor
|
we landed this in #4995, thank you for the contribution! |
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
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.
Summary
Quantizing a 2D tensor along both axes currently takes three kernels and two
full reads of the same tensor: the 1x32 CuTeDSL cast along N, the 32x1 CUDA
cast along M, and a standalone Triton K-groups swizzle of the colwise scales.
This PR adds a CuTeDSL kernel that reads the input once and emits all four
outputs, with both scale tensors already in the blocked layouts their
consumers expect.
Exposed as
mxfp8_quantize_2d_1x32_32x1_cutedsl, registered as a custom opwith a fake impl, alongside the two kernels it composes. (No call sites are
changed in this PR — see the follow-on integration PR for the backward-pass
wiring.)
Benchmark
benchmarks/prototype/moe_training/mxfp8/bench_cutedsl_quantize_2d_1x32_32x1_fused.pytimes both sides under CUDA-graph replay rather than eagerly — torchao's
triton_opdispatch costs ~100us/call regardless of problem size and thebaseline pays it three times, which swamps every shape below ~0.5 GB when
timed eagerly and reports 15x on small shapes (launch overhead, not a kernel
win). It also asserts the fused kernel does not capture as an empty graph,
which is what happens when a CuTeDSL kernel ignores the caller's stream and
would otherwise read as a spurious 100x win.
Measured locally on a GB200, 12 configs, SM clock 2062 MHz / memory clock
3996 MHz (applications clocks; the 3996 MHz memory clock is also what the
7,928 GB/s theoretical peak below is derived from):
Geomean ~2.5x over the three-kernel baseline across these shapes, fused
throughput consistently ~6.2-6.9 TB/s. Absolute rates are clock-bound — reproduce
locally rather than trusting these as a fixed constant:
Design
One CTA per
(128, 128)input tile, staged into SMEM withcp.async, runningthe rowwise and colwise passes sequentially at full CTA width. Arithmetic is
packed 16-bit throughout —
max.u16x2on sign-masked bf16 for amax,cvt.rn.satfinite.e4m3x2.bf16x2for the cast — andprmt.b32byte permuteskeep the logically transposed colwise store fully coalesced. bf16 and
rceilonly; both dimensions and every group size must be a multiple of 128 (checked
with a device-side assert, mirroring the two kernels it replaces).
Testing
Tests compare raw bytes at
atol=0/rtol=0against the three-kernelcomposition, because being a drop-in for that
composition is the contract.