[nvfp4_training] Optimize CuteDSL kernels for NVFP4 training - #4798
Open
rdspring1 wants to merge 6 commits into
Open
[nvfp4_training] Optimize CuteDSL kernels for NVFP4 training#4798rdspring1 wants to merge 6 commits into
rdspring1 wants to merge 6 commits into
Conversation
Seven rounds of SASS-driven optimization over the CuTeDSL NVFP4 quantize and amax
kernels, the fast-math path finished across both backends, and three correctness
fixes the optimization work surfaced.
The method that produced the optimization rounds: dump the compiled SASS before
spending a benchmark, find the dominant instruction family, and fix that. Three
separate register/occupancy hypotheses were killed for free that way, before any
code was written.
--------------------------------------------------------------------------------
Optimization
--------------------------------------------------------------------------------
1. BF16-to-FP4 epilogue fusion. A shared exact-mode primitive scales eight BF16-origin
values with four packed mul.f32x2 and converts them straight to four packed FP4
bytes, removing the intermediate 16-element scaled FP32 tensor.
2. Epilogue SMEM reads and packing. The 2D weight kernel spent 53% of its 3432 static
instructions on address arithmetic feeding 256 scalar LDS.U16. Vectorizing the
rowwise reads (128 LDS.U16 -> 16 LDS.128), pairing the columnwise reads so a warp
moves 128 B instead of 64 B and a row pair shares one amax and reciprocal, and
fusing the RHT-accumulator epilogue took it to 2712 instructions.
3. Grouped stochastic rounding. SR ran 2360 more static instructions than RTNE; IMAD
plus LOP3 were 76% of that and FMUL plus FMNMX 16%. Fusing the SR multiply and clamp
into the convert (two cvt.rs.satfinite.e2m1x4.f32 replacing four
cvt.rn.satfinite.e2m1x2.f32) removed the arithmetic excess entirely, and philox4_all
-- one Philox counter per 16-element block with all four output words consumed --
cut the RNG from 124 multiplies per block to 34. Excess over RTNE fell +2360 -> +736.
4. Linear stochastic rounding. The linear kernels share _quant16_from_amax, so they
inherited rounds 1-3's arithmetic fusion for free. Their SASS then showed FMUL and
FMNMX excess already at zero and Philox at 91% of the remaining 3872-instruction
excess, so the same one-draw-per-block change was applied there: +3872 -> +1112.
5. Linear RHT amax and columnwise scale factors. The amax row epilogue still issued 128
scalar LDS.U16 where the grouped amax issued 13 LDS.128 -- the round-2 fix, applied
everywhere except here. Vectorizing it took 2088 static instructions to 1704. The
columnwise scale-factor store wrote one byte at a time down a contiguous mode;
staging four in a register tile and committing one word took 25 STS.U8 to 9.
6. BF16 widening by shift. The row epilogues read bf16 out of SMEM and converted element
by element through the DSL's BFloat16 type, which ptxas lowers as a PRMT extract plus
a widen -- two instructions per value. bfloat16 is the truncation of float32, so
reading the pair as one u32 and shifting is exact and costs one, which is what
cutlass::bfloat16_t::operator float() and therefore TransformerEngine's epilogue
already get. Grouped row block 455 -> 423 static instructions, kernel dynamic 23.05M
-> 21.77M. Linear exact -3.05%, linear fast -1.12%, grouped fast -0.65%, grouped
exact neutral.
7. 64-bit row code stores. A warp's row epilogue covers 4 tokens x ROW_HB hidden blocks,
so consecutive lanes differ in hb. Writing w0 and w1 as separate u32 stores put 4 B
per lane at an 8 B lane stride: 8 lanes spanned 64 B to fill 32 B, and every sector
came back half wasted -- 8.00 sectors per instruction against an ideal of 4.00, where
TE's STG.E.64 already measured 8.00 against 8.00. row_fp4 now reaches the kernel as a
uint64 view, the same reinterpret-before-passing the wrapper does for col_sf, and the
pair goes out through a (2,) u32 rmem tensor recast to Uint64. Eight STG.E become
four STG.E.64. Global store sectors 4,587,520 -> 3,670,016, the full row-code excess;
requests 430,080 -> 315,392; bank conflicts 3.89M -> 3.44M. Grouped fast -3.05%,
grouped exact -1.13%.
--------------------------------------------------------------------------------
Fast math
--------------------------------------------------------------------------------
The CuteDSL grouped quantize had taken a use_fast_math flag since the fast-math kernels
landed, but nothing passed it, so AUTO always compiled the exact variant. nvfp4_linear
never took the flag at all, and Triton had no fast variant, which would have cost
AUTO/TRITON bitwise equivalence. All three are fixed here, and the flag defaults on.
Triton gains fast math on both families. _nvfp4_quantize is shared by the linear and
grouped kernels, so the approximate reciprocal is one edit point; _rcp_approx_ftz emits
rcp.approx.ftz.f32 via inline asm, since no tl builtin or libdevice entry reaches it
(rcp_rn/rd/ru/rz are correctly rounded and fast_dividef is div.approx, a different
instruction). The bfloat16 round-through of the tcgen05 accumulator is gated per call
site: it applies to the columnwise path only, the rowwise one having no accumulator.
The flag is threaded through the public grouped op, its autograd Function, both quantize
call sites, nvfp4_linear, NVFP4TrainingConfig and both tensor-parallel autograd
Functions. Both backends stay bitwise identical to each other and to TE in either mode,
so AUTO and TRITON agree.
Weight quantization is deliberately untouched: TE has no 2D fast path either
(quantize_transpose_nvfp4.cuh has zero use_fast_math references), and without an RHT
there is no accumulator round-through to skip. The amax kernels are likewise always
exact, in TE and in CuteDSL both.
The PyTorch oracle cannot host this. torch.reciprocal, 1/x and x.pow(-1) are all bitwise
equal to the correctly rounded reciprocal, on CPU and GPU alike, and nothing in ATen or
inductor lowers to MUFU.RCP -- verified over 4M values. nvfp4_reference therefore stays
exact-only and keeps its CPU-pure contract, and Triton takes the fast-path oracle role
instead. That works because triton's tl.dot fp32 accumulator is bitwise identical to
TE's tcgen05 UMMA accumulator: test_fast_math_matches_transformer_engine passes on both
backends across all 25 shapes, which is what the bf16 cast had been hiding. Fast math is
tied back to the reference transitively, by bounding it against the exact path the
reference does pin bitwise.
The fast-vs-exact SQNR floor is 25 dB, set from measurement: columnwise measures
30.2-32.6 dB linear and 30.4-31.9 grouped, with ~1.2% of code bytes differing. That is
the shape of the error rather than a defect -- skipping the bfloat16 round-through moves
the value by ~2**-9, but an element near an E2M1 midpoint then flips a whole FP4 step, so
the error is large and rare. It sits above NVFP4's own 20 dB quantization noise. Rowwise
shares the loose floor deliberately: block scales are identical between modes, so a
rowwise flip needs rcp.approx and div_rn to disagree on pvscale_fp8 * global_decode_scale,
and that product takes only a handful of distinct values per tensor -- either none
disagrees and the output is bitwise identical, or one does and every element sharing that
scale becomes eligible at once. For the seed224 fixture all 185 flips land in group 1 and
none in group 0. A tight rowwise bound would be flaky by construction.
Both RHT benchmark scripts take --math {exact,fast,all}, mirroring the existing
--rounding, so the README's fast rows are reproducible from checked-in scripts rather than
measured out of band.
--------------------------------------------------------------------------------
CUDA-graph warm-up
--------------------------------------------------------------------------------
cutedsl_prepare_for_cuda_graph exists so no lazy cute.compile fires inside a captured
region. It never did that. An lru_cache key is the literal (args, kwargs) shape, so
_compile_fused_kernel(idx, True, sr, apply_rht=True, ...) and the call site's
_compile_fused_kernel(dev.index, swizzle, sr, bool(apply_rht), ...) are two different
entries compiling the same kernel. The warm-up compiled twelve kernels and hit none of
them; measured on a GB200, prepare() followed by a default CUTEDSL nvfp4_linear reported
hits=0 with two fresh compiles. The grouped twin missed for a second reason: it omitted
fast_math entirely, so its 3-tuple key never matched the runtime's 4-tuple.
Hand-matching call shapes is what broke, so the fix removes the ability to diverge:
_compile_fused_kernel and _compile_group_fused_kernel take every parameter as a required
positional, and all four call sites pass them that way with inline names. A future keyword
call is now a TypeError instead of a silent second cache entry. fast_math joins both
warm-up loops -- it is a compile-time specialization that defaults on, so the variant
production actually runs was the one variant prepare could not reach even after the keys
lined up. The grouped caches also drop their bounds: maxsize=16 and maxsize=4 admit
eviction, and an evicted entry is a lazy recompile mid-capture, the exact hazard the
sibling cache documents when it pins maxsize=None.
Nothing observed any of this. test_torch_compile_cuda_graph passes either way, because
cudagraph-trees warms up eagerly before it captures, so a no-op warm-up looks identical to
a working one. The new test asserts the property instead: snapshot cache misses after
prepare, run forward and backward over both math modes (backward reaches the SR variant,
another key), and require zero new compiles. It covers all four caches -- linear fused,
grouped fused, and both amax -- driven by a linear and a grouped MoE forward+backward.
Verified by stubbing out only the grouped half of the warm-up: fails with group_fused=4,
group_amax=1 and the linear caches still at zero.
--------------------------------------------------------------------------------
Grouped allocation tail
--------------------------------------------------------------------------------
The grouped kernels zero-filled every row of the output allocation, including rows past
logical_packed_length == offsets[-1]. Those rows are storage capacity that no grouped
consumer addresses, so the zero-fill was pure work. The kernels now stop at the last
logical row and the tail is documented as untouched and unspecified across the op
docstrings, with a device-side assert that SAME_BOTH_DIMS carries no spare capacity. The
padded-capacity test correspondingly poisons the tail and asserts it cannot reach any
group-addressable output, rather than asserting the tail is zero.
--------------------------------------------------------------------------------
Config surface and AUTO
--------------------------------------------------------------------------------
nvfp4_linear's kernel_preference defaults AUTO, but _KERNEL_PREFS only parametrized TRITON
and CUTEDSL, so the default's fallback path was untested: on any SM100 box with the
CuteDSL runtime -- every box that runs the rest of the file -- AUTO always resolves to
CuteDSL and the Triton-fallback branch is unreachable. _resolve_use_cutedsl is extracted
out of nvfp4_linear, the same split the grouped path already made with _resolve_backends,
so the decision has a test target that does not need a box without the CuteDSL runtime to
reach it. AUTO joins _KERNEL_PREFS, with a monkeypatched fallback test mirroring the
grouped one and an end-to-end check that AUTO under a forced fallback matches an explicit
TRITON call.
use_fast_math and the AUTO default are both documented in NVFP4TrainingConfig with the
exact config that recovers prior numerics, so a loss-curve diff across this branch does not
carry two silently-moved variables. That default flip also exposed an asymmetry AUTO's own
docstring did not carry: the grouped CuteDSL and Triton kernels are bitwise identical under
RTNE but draw different stochastic-rounding streams, so AUTO's silent fallback changes
backward-pass numerics on a node without the CuteDSL runtime, not just which kernel runs.
The linear kernels retain SR parity; only grouped does not.
deepseek_v3_shapes.py declared the 671B config at expert_parallel_degree=2, 128 local
experts -- not this model's training layout, which is EP=64, 4 local experts, as every
caller here already assumes by passing factorized_experts=4 directly. Harmless in practice
because local_experts was consequently never called, but it is the only place in the repo
documenting this layout.
MAX_GROUPS=64's comment cited TE's kMaxTensorsPerKernel as the reason for the number. It
reproduces the value but not the cause: that TE cap is a 4 KB kernel-argument budget on
TE's pointer-list kernels, which do not apply here -- this design passes a packed tensor
plus device-side offsets, the same shape as TE's own graph-safe kernels, which run an
unbounded search and hit no such limit. The real constraint is that the group lookup
unrolls to a fixed depth for branch-free epilogues, so it resolves exactly
2**GROUP_SEARCH_STEPS groups; past that the search exits with hi - lo > 1 and returns a
group index off by one -- a silently wrong per-group amax, not a failure. An assert now
couples the two constants so they cannot drift apart silently.
Dead code retired along the way: philox_c0 has had no caller since philox4_all superseded
it, and _Tcgen05GroupRowColFused.swizzle_sf was stored and never read (the columnwise store
is hardwired to the swizzled layout), so the swizzle argument on _compile_group_fused_kernel
goes with it and the cache key narrows to (device_idx, sr, fast_math).
--------------------------------------------------------------------------------
Performance
--------------------------------------------------------------------------------
GB200, CUDA 13.4, PyTorch 2.15.0a0+git0f3e7e2, Triton 3.8.0, nvidia-cutlass-dsl 4.5.2.
Device kernel self-time in microseconds, median of three full script passes at 15 warmups
/ 50 timed profiler iterations.
CuteDSL against Triton, DeepSeek-V3 671B at E=4, the grouped quantize stage. This is the
full (backend x math x rounding) cross product, all of it reproducible from
bench_group_rht_quantize_row_col --experts 4 --math all:
projection rounding math CuteDSL Triton speedup CuteDSL fast/exact
gate/up RTNE exact 53.95 87.49 1.62x -
gate/up RTNE fast 39.73 64.16 1.62x 1.36x
gate/up SR exact 74.75 159.96 2.14x -
gate/up SR fast 54.93 141.28 2.57x 1.36x
down RTNE exact 55.23 86.91 1.57x -
down RTNE fast 41.14 63.71 1.55x 1.34x
down SR exact 76.51 159.30 2.08x -
down SR fast 56.51 140.86 2.49x 1.35x
The grouped amax has no fast-math variant (exact on both backends and in TE), and the
grouped 2D weight kernel has none either -- without an RHT there is no accumulator
round-through to skip. Both are complete at two cells rather than four: amax 22.95 vs
39.91 (1.74x) gate/up and 22.42 vs 40.54 (1.81x) down; 2D weights 59.76 vs 109.31 (1.83x)
and 60.02 vs 107.99 (1.80x).
Against TransformerEngine, DeepSeek-V3 671B FFN at E=4, complete pipeline (amax plus
quantize) versus TE's split_quantize. TE is 2.19.0.dev0 built from
NVIDIA/TransformerEngine@172bd93, an untagged
main-branch commit. No checked-in script reproduces this comparison -- the benchmark
modules cover CuteDSL-vs-Triton only, and these are the (2048, 7168) / (7168, 2048)
shapes --shape-set does not offer. The pair below is retained from a single run in which
both sides were measured together, and predates the bf16-shift and 64-bit-store rounds,
so the CuteDSL column is a lower bound; the TE-relative ratio is still the best available,
because re-measuring one side alone would be worse than leaving the pair intact:
projection math rounding CuteDSL TE TE speedup
gate/up standard RTNE 86.22 64.03 1.35x
gate/up standard SR 106.44 87.90 1.21x
gate/up fast RTNE 70.44 55.66 1.27x
gate/up fast SR 86.24 77.74 1.11x
down standard RTNE 88.35 62.59 1.41x
down standard SR 108.59 86.67 1.25x
down fast RTNE 71.54 54.73 1.31x
down fast SR 87.43 76.47 1.14x
Linear (2048, 7168), same construction versus a single-tensor NVFP4Quantizer:
math rounding CuteDSL TE TE speedup
standard RTNE 29.82 18.61 1.60x
standard SR 36.69 31.33 1.17x
fast RTNE 21.75 15.95 1.36x
fast SR 28.24 24.63 1.15x
2D weight quantize. Both sides consume a precomputed amax, so the comparison is against
TE's quantize kernel alone, not its full three-kernel call (which adds amax 5.40 and
zero_amax 1.34):
kernel projection CuteDSL TE quantize TE speedup
2D linear gate/up 16.33 13.80 1.18x
2D linear down 16.36 13.76 1.19x
2D grouped, E=4 gate/up 60.26 55.20 (x4) 1.09x
2D grouped, E=4 down 60.57 55.04 (x4) 1.10x
At the op level, where each side computes its own amax, CuteDSL wins the 2D comparison.
Neither side has a 2D fast path: NVTE_USE_FAST_MATH moves TE's kernels by at most 0.02 us.
Fast math is worth 1.57-1.91x to CuteDSL on the linear path and 1.14-1.22x to Triton;
grouped is flatter at 1.27-1.38x for CuteDSL, and splits for Triton -- 1.28-1.36x under
RTNE but only 1.13-1.14x under SR, where the Philox work it still carries dominates what
fast math removes. CuteDSL at N=28672 goes from 33.4% to 63.6% of peak bandwidth.
Against the Triton backend, CuteDSL leads 1.62x on grouped RTNE, 2.14x on grouped SR,
1.74x on the grouped amax, 1.83x on grouped 2D weights, 1.60-1.91x on the linear amax,
1.63-1.82x on the linear quantize (2.43-2.65x with fast math on both sides), and
2.71-2.86x on 2D weights.
Full tables, per-stage splits, methodology and the rejected experiments are in
benchmarks/prototype/nvfp4_training/README.md.
--------------------------------------------------------------------------------
Correctness
--------------------------------------------------------------------------------
Round-to-nearest-even output is bitwise identical to both oracles on every path, and no
round here changes that: cases against the Triton backend and against the TE-derived
PyTorch reference in
test/prototype/moe_training/nvfp4_training/nvfp4_reference.py.
Stochastic rounding is deliberately NOT bitwise-comparable to Triton, on either the linear
or the grouped path. Both CuteDSL kernels draw one Philox counter per 16-element block and
consume all four output words rather than reproducing Triton's per-packed-byte counter
stride; the stream is a different, equally valid one. It stays a pure function of tile
coordinates and the caller's rng_state, so results remain reproducible under the persistent
CLC scheduler, whose visit order is not fixed.
SR is covered instead by, in decreasing strength:
- test_rht_quantize_rs_at_most_one_fp4_step_from_rtne: every SR code sits within one
FP4 magnitude step of the RTNE code with matching sign, on both backends over one
tile, several tiles and a short trailing column group. The RTNE code it pins against
is itself bitwise-checked against Triton and TE, so a wrong nibble order, scale or
block index still fails loudly without any SR oracle.
- test_group_rht_sr_reconstructs: SQNR through the same per-group reference the RTNE
correctness test uses. The bar is 15 dB rather than RTNE's 20; SR measures 17.2 dB on
both backends and all four fixtures within 0.2 dB, which is the expected ~3 dB
variance cost of unbiased rounding, not a defect.
- test_cutedsl_rht_quantize_sr_unbiased and test_group_rht_sr_unbiased: averaging SR
draws of an exactly-halfway value converges to it with a ~50/50 grid split, which
catches a degenerate or position-correlated stream that SQNR cannot see.
- test_group_rht_rng_state_controls_stochastic_rounding: identical rng_state yields
identical codes, advanced state differs. Backend-agnostic, and the real guard on the
coordinate-derived counter.
761 passed / 36 skipped single GPU. The tensor-parallel and FSDP2+TP suites need
torchrun and were not run in this pass. Ruff lint and format clean.
--------------------------------------------------------------------------------
Refuted, recorded so they are not re-derived
--------------------------------------------------------------------------------
- warpgroup_reg_alloc/dealloc does not help the linear amax. setmaxnreg redistributes
registers within a CTA's existing allocation; it does not lower the compiled base, and
the base is what sets blocks per SM. Measured at +-0.5% with REG:98 unchanged.
- Raising occupancy does not help the linear amax either. The 128-row supertile does
compile to REG:42, genuinely fitting four blocks per SM against the 256-row config's
one, but pushing GRID past NUM_SMS is monotonically worse at every shape, co-resident
or not. ncu's 3.32 active warps per scheduler and 60.6% no-eligible-warp cycles
describe the kernel accurately without identifying its constraint.
- SR does not spill against the shared REG_COL/REG_ROW budget: both variants report
REG:128 with no MOV.SPILL, which killed a planned register-budget round outright.
- The same scale-factor packing applied to the rowwise store takes 9 STS.U8 to 1 but
moves every case within +-0.5%, so it was discarded.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
rdspring1
requested review from
andrewor14,
jerryzh168 and
vkuzo
as code owners
August 20, 2026 04:08
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4798
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
Contributor
Author
|
@vkuzo After talking with MLPerf folks, I've changed the recipe to quantize all the FFN layers with NVFP4 while the Attention layers follow the MXFP8 recipe in TorchTitan. I'm currently running DSV3 16B 1500 step run from Experiment 2: DeepSeek-V3 16B MoE Loss Convergence Validation. Can you advise on how you want users to pick between (CuteDSL, Triton) and (Standard, Fast-Math)? At the moment, this PR has |
Contributor
How about using |
Contributor
Author
|
The ideal cost for quantizing 16 fp32 to fp4 is 1 Philox-4x32-10 invocation for 4
Quantizing 16 fp32 to fp4 yields 8 packed uint8. A There are two separate areas that waste random numbers.
|
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
3 participants
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
Brings the CuteDSL NVFP4 quantization kernels to parity-or-better with the Triton backend across the DeepSeek-V3 shapes, adds a fast-math path to both backends, and promotes CuteDSL + fast math to the default via
KernelPreference.AUTO.CuteDSL now leads Triton by 1.6x on grouped RTNE, 2.1x on grouped SR, 1.7-1.8x on the grouped amax, and 1.8x on grouped 2D weights at 671B, and fast math is worth a further 1.3-1.4x on top. Round-to-nearest-even output stays bitwise identical to both the Triton backend and the TransformerEngine-derived reference on every path.
The latest configuration is NVFP4 for all dense linears, shared and routed experts and MXFP8 attention with fast math enabled. In E2E run with 64 GB300 GPUs for 1 hour, NVFP4 is 37.3% faster than BF16 and +12.8% faster than MXFP8 (1004.6 vs 731.6 and 890.8 TFLOP/s) while using less memory than either (217.08 GiB, −8.38 GiB vs bf16 and −5.41 GiB vs mxfp8). The training loss curve is inline with BF16 and MXFP8.
DSV3 671B E2E Run, GBS 512, 64 GB300 GPUs, Training Loss, Throughput, and Memory
Throughput (steps >= 50):
Memory (peak reserved per GPU, 276 GiB device):
Key Changes
Details
Fast Math Path
Fast math matches TransformerEngine under
NVTE_USE_FAST_MATH=1and does two things:div.rnencode reciprocal withrcp.approx.ftz.f32(oneMUFU.RCP).Triton previously had no fast variant at all. Adding it keeps the two backends bitwise identical to each other and to TE in either mode, so
AUTOandTRITONagree.Weight quantization is deliberately untouched: without an RHT there is no accumulator to skip, and TE has no 2D fast path either (
quantize_transpose_nvfp4.cuhhas zerouse_fast_mathreferences). The amax kernels are always exact, in TE and CuteDSL both.KernelPreference.CuteDSL and Fast Math as the default
NVFP4TrainingConfig.kernel_preferencemovesTRITON->AUTOanduse_fast_mathdefaultsTrue.AUTOresolves to CuteDSL on SM100 with the CuteDSL runtime present and falls back to Triton otherwise.Both defaults are documented in the config with the exact settings that recover prior numerics (
kernel_preference=KernelPreference.TRITON, use_fast_math=False), so a loss-curve diff across this branch does not silently carry two moved variables.Fast-vs-exact measures 30.2-32.6 dB SQNR on the linear path and 30.4-31.9 grouped — roughly 10 dB quieter than NVFP4's own ~20 dB quantization noise. The error is large and rare rather than small and uniform (~1.2% of code bytes differ): skipping the bf16 round-through moves a value by ~2**-9, but an element near an E2M1 midpoint then flips a whole FP4 step.
Philox divergence between CuteDSL and Triton
Triton and CuteDSL are bitwise identical for RTNE but not for SR
Triton's SR path calls
tl.randint(seed, offset)once per packed byte. Two things waste work there:tl.randintcomputes a full Philox-4x32-10 draw (4 words) and returns one of them.cvt.rsconversion goes throughtl.inline_asm_elementwisewithpack=4, andpackapplies uniformly to every argument. So the asm block receives 4 random words while its twocvt.rs.satfinite.e2m1x4.f32instructions consume only 2 ($9,$10) — the other two are computed and discarded. There is no way to request a different pack factor for the random-bit argument than for the data arguments.CuteDSL is not bound by that shape.
philox4_alldraws one counter per 16-element block and consumes all four output words, computing all four in the last round: 34 multiplies per block instead of 124. The counter is derived from tile coordinates rather than a running per-thread value, because the persistent CLC scheduler's visit order is not fixed — so results stay reproducible for a givenrng_state.The consequence: CuteDSL's SR stream is a different, equally valid stream, not a matching one. It is validated on properties rather than bit-equality — every SR code within one FP4 magnitude step of the bitwise-checked RTNE code, reconstruction SQNR, unbiasedness at exact-halfway values, and
rng_statereproducibility.This is the one behavioural asymmetry in
AUTO: on a node without the CuteDSL runtime, the silent Triton fallback changes backward-pass numerics under SR, not just which kernel runs. Pinkernel_preferenceexplicitly for runs that must reproduce bitwise across machines.Performance: CuteDSL vs Triton
GB200, CUDA 13.4, PyTorch 2.15.0a0+git0f3e7e2, Triton 3.8.0, nvidia-cutlass-dsl 4.5.2.
Device kernel self-time in microseconds, median of three full script passes at 15 warmups and 50 timed profiler iterations.
E = 4local experts, the EP=64 training layout.Reproducible with:
DeepSeek-V3 671B with 4 local experts
Grouped RHT quantize — round-to-nearest-even
Grouped RHT quantize — stochastic rounding
The SR lead is larger than the RTNE lead precisely because of the Philox change: SR costs
CuteDSL ~1.4x its own RTNE time against Triton's ~1.8x.
Grouped amax and 2D weights (no fast-math variant — see above)
Fast-math value by backend
Under RTNE both backends gain about the same. The split appears under SR, where Triton
gains only 1.13-1.14x: the Philox work it still carries dominates what fast math removes.
AI tools used