Repository navigation
[moe_training] Fused dual-axis MXFP8 quantization with correct metadata - #4995
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>
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4995
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 86f7e2a with merge base 9d0c494 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
## Agent details Prepared with AI assistance. Reuse TorchAO's fused rowwise/columnwise activation cast and existing single-axis CuTeDSL casts, consuming blocked E8M0 scales without separate rearrangement kernels. Keep CUDA as the default and expose CuTeDSL as an explicit activation_quantization_backend option through both the linear and its converter. Input-save format, weight quantization, output ownership, FSDP lifetime and gradient accumulation policy are unchanged. CuTeDSL requires 128-aligned BF16 matrices and the required runtime/API at construction. Unsupported shapes fail explicitly rather than padding activations to a larger GEMM or silently selecting another implementation. Represent dense matrices as one group, view the valid column-scale prefix, and retain immutable real endpoints for CUDA-graph pointer lifetime. Fake endpoints never enter the real cache. Explicitly request RCEIL for every activation cast, matching the weight-scale policy independently of provider defaults. The existing operand tests check this forwarding while still executing each real provider kernel. Depends on pytorch/ao#4995 as the single AO dependency. That PR preserves Protonu's kernel from #4895 and includes current nonfinite handling, the 32x1 FakeTensor fix (also addressed by Aleksandar in #4702), and byte-oracle regressions. The original PRs are provenance, not additional dependencies. ## Test Plan Real CUDA comparisons validate all cast directions and two backward contributions, including in-place accumulation. Real two-rank FSDP checks cover compiled RAF=true, retained PP weights and CUDA-graph replay. Schema/FakeTensor opchecks validate provider layouts. Matched CUDA-graph replay checks defined output bytes, timing and actual target kernels. ```bash python -m pytest tests/unit_tests/cpu/quantization/test_quantization.py tests/unit_tests/cpu/quantization/test_quantized_fsdp_tensor.py -q python -m pytest tests/unit_tests/gpu/quantization/test_mxfp8_linear.py -q python -m pytest tests/unit_tests/gpu/quantization/test_mxfp8_fsdp.py -q -k 'cutedsl or (tensor_lifecycle and not simple)' pre-commit run --all-files ```
## Agent details Prepared with AI assistance. Reuse TorchAO's fused rowwise/columnwise activation cast and existing single-axis CuTeDSL casts, consuming blocked E8M0 scales without separate rearrangement kernels. Keep CUDA as the default and expose CuTeDSL as an explicit activation_quantization_backend option through both the linear and its converter. Input-save format, weight quantization, output ownership, FSDP lifetime and gradient accumulation policy are unchanged. CuTeDSL requires 128-aligned BF16 matrices and the required runtime/API at construction. Unsupported shapes fail explicitly rather than padding activations to a larger GEMM or silently selecting another implementation. Represent dense matrices as one group, view the valid column-scale prefix, and retain immutable real endpoints for CUDA-graph pointer lifetime. Fake endpoints never enter the real cache. Explicitly request RCEIL for every activation cast, matching the weight-scale policy independently of provider defaults. The existing operand tests check this forwarding while still executing each real provider kernel. Depends on pytorch/ao#4995 as the single AO dependency. That PR preserves Protonu's kernel from #4895 and includes current nonfinite handling, the 32x1 FakeTensor fix (also addressed by Aleksandar in #4702), and byte-oracle regressions. The original PRs are provenance, not additional dependencies. ## Test Plan Real CUDA comparisons validate all cast directions and two backward contributions, including in-place accumulation. Real two-rank FSDP checks cover compiled RAF=true, retained PP weights and CUDA-graph replay. Schema/FakeTensor opchecks validate provider layouts. Matched CUDA-graph replay checks defined output bytes, timing and actual target kernels. ```bash python -m pytest tests/unit_tests/cpu/quantization/test_quantization.py tests/unit_tests/cpu/quantization/test_quantized_fsdp_tensor.py -q python -m pytest tests/unit_tests/gpu/quantization/test_mxfp8_linear.py -q python -m pytest tests/unit_tests/gpu/quantization/test_mxfp8_fsdp.py -q -k 'cutedsl or (tensor_lifecycle and not simple)' pre-commit run --all-files ```
| @pytest.mark.skipif( | ||
| not _mxfp8_cutedsl_kernels_available, reason="MXFP8 cutedsl kernels not available" | ||
| ) | ||
| def test_mxfp8_quantize_2d_1x32_32x1_cutedsl_input_validation(): |
There was a problem hiding this comment.
nit: i'd remove this, low value test
|
@claude review, pay special attention to how the new kernel is tested vs reference |
|
Claude finished @vkuzo's task in 2m 59s —— View job Review of PR #4995: fused dual-axis MXFP8 quantization
I reviewed this by reading the code only; I didn't run anything on GPU (the runner has no SM10x). I traced the kernel's indexing by hand: the rowwise and colwise scale swizzles ( Testing vs reference (the main focus)1. Nothing compares the kernel to
On current main I don't think that's true any more:
Suggestion: add a direct 2. Special values are only checked against the shared contract for the rowwise half. 3. Empty groups are never tested.
4. The flat-prefix 5. 6. Small test-hygiene points
Kernel / wrapper7. Dead code: the 8. 9. The 10. Offsets that aren't ascending get through. 11. Minor
SummaryThe kernel looks correct everywhere I traced it, and the destination-buffer, uncovered-tail and NaN/Inf tests are good additions. My main request, given the question asked: anchor the tests to Fix the dead branch and |
| "torchao::mxfp8_quantize_2d_1x32_32x1_cutedsl", mutates_args=() | ||
| ) | ||
| def _mxfp8_quantize_2d_1x32_32x1_cutedsl_custom_op( | ||
| x: torch.Tensor, |
There was a problem hiding this comment.
can we match the argument order of surrounding APIs like _mxfp8_quantize_2d_1x32_cutedsl_custom_op
x: torch.Tensor,
block_size: int = 32,
scaling_mode: str = "rceil",
offs: torch.Tensor = None,
then assert for any unsupported values
There was a problem hiding this comment.
Matched the public, custom-op and fake signatures to (x, block_size=32, scaling_mode="rceil", offs=None) and added consistent unsupported-value checks. offs=None selects dense placement in the same kernel without offset processing or scale slack; tensor offsets retain grouped, empty-expert and uncovered-tail behavior. Updated all AO callers and TT#5334's consumer, removing its synthetic endpoint/cache. Real/fake opchecks, actual Inductor compilation and the two-rank consumer checks passed.
|
lgtm, can we just update the public arg order (see comment inline) and remove the unnecessary test |
| q_row, s_row = mxfp8_quantize_2d_1x32_cutedsl( | ||
| x, block_size=32, scaling_mode="rceil", offs=offs | ||
| ) | ||
| mx = _to_mxfp8_dim1_kernel_wrapper( |
There was a problem hiding this comment.
optional nit: the authorative numerics reference is
ao/torchao/prototype/mx_formats/mx_tensor.py
Line 227 in 9d0c494
There was a problem hiding this comment.
Switched the primary numerical checks to direct to_mx(RCEIL) for both axes, with per-group unblocking. This exposed a real scale0/1 cutoff bug: the fused kernel used 0x04E00000 instead of 0x04600000 (448*2^-127). Corrected it and added exact neighboring values on both CTA paths, mixed magnitudes, independent columnwise semantic cases and explicit empty groups. Finite bytes are checked exactly; FP8 NaN signs follow the shared semantic contract. The current 63-test suite and strict API documentation build passed.
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
5fb606c to
86f7e2a
Compare
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: > Reuse TorchAO's fused dense rowwise/columnwise activation cast and existing > single-axis CuTeDSL casts, consuming directly blocked E8M0 scales. Keep CUDA > as the default and CuTeDSL as explicit activation_quantization_backend opt-in > through the Linear and converter. Input-save format, weight quantization, > output ownership, FSDP lifetime and gradient accumulation policy are unchanged. > > Use the provider's offs=None dense specialization, with compact column scales > and no synthetic endpoint, cache, FakeTensor endpoint branch or padding trim. > Construction rejects an older required-offset provider as well as missing > symbols/runtime. Both real and fake paths use the same dense contract. > > CuTeDSL requires 128-aligned BF16 matrices. Unsupported shapes fail explicitly, > without padding to a larger GEMM or silently changing implementations. Pin > RCEIL for every activation cast to match the 32x32 weight-scale policy. > > Depends only on pytorch/ao#4995. Preserve credit to Protonu's original #4895 > kernel and Aleksandar's #4702 single-axis metadata work; those PRs remain > provenance, not additional dependencies. > > Test plan: current CPU contracts, real CUDA cast directions/backward formats > and accumulation, two-rank compiled/retained/CUDA-graph FSDP lifecycle, > strict configuration/README rendering, complete repository lint and matched > replay timings with actual target-kernel traces.
|
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:
|
|
@claude review |
|
Claude finished @vkuzo's task in 3m 1s —— View job Re-review of PR #4995 (after 86f7e2a)
I reviewed by reading the code only. I couldn't run anything: the runner has no SM10x, and ad-hoc scripts weren't permitted here. @vkuzo's requests: both done
Prior Claude review items
Switching to Checking
|
vkuzo
left a comment
vkuzo
left a comment
There was a problem hiding this comment.
amazing! thank you @sanketpurandare for making this ready, and also thank you @protonu for the great work in #4895!
4 participants
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: