[CuteDSL] Add benchmarking, CE memory optimizations, and match Triton num_warps - #1290
Merged
vaibhavjindal merged 8 commits intoAug 3, 2026
Merged
Conversation
Follow-up to linkedin#1279 addressing the review follow-ups (memory + benchmark scripts). Memory parity fix ----------------- The CuTe DSL CE backward, for the scalar-grad (reduction mean/sum) path, computed `_input * grad_output`, allocating a *second* BT×V buffer. That doubled peak memory vs the Triton CE, which scales the saved gradient IN PLACE via a raw element-wise kernel (no autograd version bump). Measured on B200 (full fwd+bwd, BT=8192, V=128256): cutedsl 4008 MB vs triton 2004 MB in bf16. Fix: mirror Triton exactly — scale in place with `element_mul_kernel`. A raw Triton kernel (not `_input *= grad_output`) is used so the in-place write doesn't bump the autograd version counter and trip backward-through-backward. Peak memory now equals Triton across the whole vocab/BT sweep. CE parity suite: 158 passed. Benchmark tooling (reviewer ask: "add benchmarking scripts for cutedsl, like cutile") ------------------------------------------------------------------------------------- - benchmark/scripts/run_cutedsl_compare.py: Triton-vs-CuTe-DSL compare driver (mirrors run_cutile_compare.py), runs benchmark_<kernel>.py twice under LIGER_KERNEL_IMPL and tags the providers (liger_triton / liger_cutedsl) into benchmark/data/all_benchmark_data_cutedsl.csv. Supports cross_entropy. - benchmark/README.md: document the compare drivers (CuTile + CuTe-DSL). - benchmark/data/all_benchmark_data_cutedsl.csv: sample CE speed+memory dataset (liger_triton / liger_cutedsl / torch), generated by the new driver on B200. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The CuTe DSL CE forward baked 8 warps/CTA for any 2-byte dtype (a Blackwell-only tuning). On Hopper (sm_90) that underfills the SMs and loses ~0.90x to the 32-warp Triton forward. Mirror the Triton CE convention exactly via infer_device_arch()/is_hip(): - Blackwell (sm_100+): bf16/fp16 -> 8, fp32 -> 32 - Hopper (sm_90) and earlier: 32 for all dtypes - AMD (ROCm): 16 num_warps is baked into the kernel and is part of the compile-cache key, so this is a pure launch-config change with no numerics impact. On H100 bf16 goes from ~0.89-0.91x to parity-to-a-win on the vocab sweep (0.91-1.11x); memory is unchanged (exact Triton parity). Update the white-box test_num_warps_matches_dtype_convention to assert the arch-aware value instead of the fixed 8/8/32. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
| # without that bookkeeping — exactly how the Triton CE backward dodges the same issue. | ||
| BT, V = _input.shape | ||
| BLOCK_SIZE = min(_MAX_FUSED_SIZE, triton.next_power_of_2(V)) | ||
| element_mul_kernel[(BT,)]( |
Collaborator
There was a problem hiding this comment.
Let's not use triton kernel in cutedsl code. This is a bit confusing and could lead to issues when triton is not installed.
Let's revert to the original code, or maybe write an equivalent kernel in cutedsl.
Contributor
Author
There was a problem hiding this comment.
@vaibhavjindal makes sense. I have implemented a cutedsl version and got roughly the same speed and memory numbers. Let me know if there is anything else to change!
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
kolehma8
reviewed
Jul 28, 2026
| x_frag = cute.make_rmem_tensor((VEC,), gX.element_type) | ||
|
|
||
| for i in cutlass.range(0, cute.ceil_div(num_vec, 1024)): | ||
| vec_idx = tid + i * 1024 |
Collaborator
There was a problem hiding this comment.
where is the 1024 coming? If that is to max out the 4kB page size, can we make this dependent on the numerical precision.
Contributor
Author
There was a problem hiding this comment.
1024 is threads per block, so vec_idx is obtained through thread index tid and loop index i
| if is_hip(): | ||
| expected_warps = 16 | ||
| else: | ||
| is_blackwell = infer_device_arch().startswith("blackwell") |
Collaborator
There was a problem hiding this comment.
nit, should these be moved to a utility function?
Contributor
Author
There was a problem hiding this comment.
infer_device_arch() is already a utility function
pull Bot
pushed a commit
to dumpmemory/Liger-Kernel
that referenced
this pull request
Jul 29, 2026
## Summary This PR isolates one optimization on top of the already-merged CE improvements in linkedin#1266 (`exp2`) and linkedin#1267 (dtype-aware `num_warps`): replace the per-element true-class gradient selection with a single post-loop `dx_y` correction. At the reference shape (`BT=8192`, `V=128256`, full forward + backward), the isolated impact is: | GPU | dtype | baseline (ms) | this PR (ms) | speedup | |:----|:------|--------------:|-------------:|--------:| | B200 | bf16 | 1.4229 | 1.1787 | **1.207x** | | B200 | fp32 | 1.8820 | 1.8022 | **1.044x** | | H100 | bf16 | 2.3162 | 2.2734 | **1.019x** | | H100 | fp32 | 4.5362 | 4.5674 | **0.993x** | ## What changed The CE gradient subtracts 1 only at the true-class index `y`. Previously, both pass-2 paths evaluated a per-element `tl.where(X_offsets != y, ...)` across the entire vocabulary. This PR stores the plain softmax gradient in the loop, then applies one scalar correction at `y` after the loop. This removes an `O(V)` compare/select used to update one element and replaces it with an `O(1)` store. A `tl.debug_barrier()` orders the loop's in-place writes before the correction reads and updates `dx_y`. The correction retains the same weighting, reduction, label-smoothing, and softcap factors. ## Speed benchmarks `baseline` is `main` with linkedin#1266 and linkedin#1267 but without this PR's single-shot `dx_y` change. `speedup = baseline / this PR`, so values above `1.00x` mean this PR is faster. Each latency is the average of two independent medians. B200 runs were collected on swapped NVIDIA B200 GPUs; H100 runs were collected on one NVIDIA H100 80GB HBM3. Each median uses 40 timed iterations after 20 warmup iterations with CUDA events. `full` is forward plus `loss.backward()` with `grad_output=1.0`, matching the standard last-layer CE benchmark used in linkedin#1290. ### bf16 **Speed - full (forward + backward) - vocab sweep (`BT=8192`)** | vocab | B200 baseline (ms) | B200 this PR (ms) | B200 speedup | H100 baseline (ms) | H100 this PR (ms) | H100 speedup | |------:|-------------------:|------------------:|-------------:|-------------------:|------------------:|-------------:| | 32000 | 0.5179 | 0.4455 | **1.163x** | 0.8484 | 0.7890 | **1.075x** | | 50304 | 0.7967 | 0.6663 | **1.196x** | 1.2223 | 1.1375 | **1.075x** | | 102400 | 1.3693 | 1.1129 | **1.230x** | 2.0434 | 1.8907 | **1.081x** | | 128256 | 1.4240 | 1.1802 | **1.207x** | 2.3179 | 2.2790 | **1.017x** | | 152064 | 1.7677 | 1.4415 | **1.226x** | 2.9131 | 2.7976 | **1.041x** | | 201088 | 2.4596 | 1.9703 | **1.248x** | 3.8283 | 3.7046 | **1.033x** | | 262144 | 2.8544 | 2.3144 | **1.233x** | 4.7663 | 4.6687 | **1.021x** | **Speed - full - `BT` sweep (`V=128256`)** | BT | B200 baseline (ms) | B200 this PR (ms) | B200 speedup | H100 baseline (ms) | H100 this PR (ms) | H100 speedup | |---:|-------------------:|------------------:|-------------:|-------------------:|------------------:|-------------:| | 1024 | 0.3744 | 0.3356 | **1.116x** | 0.5599 | 0.5552 | **1.008x** | | 2048 | 0.5264 | 0.4595 | **1.146x** | 0.8055 | 0.7913 | **1.018x** | | 4096 | 0.8250 | 0.6987 | **1.181x** | 1.3031 | 1.2690 | **1.027x** | | 8192 | 1.4229 | 1.1787 | **1.207x** | 2.3162 | 2.2734 | **1.019x** | | 16384 | 2.6574 | 2.1671 | **1.226x** | 4.2325 | 4.1427 | **1.022x** | | 32768 | 5.0804 | 4.1005 | **1.239x** | 8.0368 | 7.8853 | **1.019x** | | 65536 | 9.9058 | 7.9347 | **1.248x** | 15.6748 | 15.3520 | **1.021x** | ### fp32 **Speed - full (forward + backward) - vocab sweep (`BT=8192`)** | vocab | B200 baseline (ms) | B200 this PR (ms) | B200 speedup | H100 baseline (ms) | H100 this PR (ms) | H100 speedup | |------:|-------------------:|------------------:|-------------:|-------------------:|------------------:|-------------:| | 32000 | 0.6876 | 0.6391 | **1.076x** | 1.0472 | 1.0402 | **1.007x** | | 50304 | 1.0058 | 0.9423 | **1.067x** | 1.5734 | 1.6267 | **0.967x** | | 102400 | 1.7260 | 1.6429 | **1.051x** | 3.7195 | 3.7151 | **1.001x** | | 128256 | 1.8806 | 1.8026 | **1.043x** | 4.5767 | 4.5742 | **1.001x** | | 152064 | 2.3784 | 2.2591 | **1.053x** | 5.3767 | 5.3722 | **1.001x** | | 201088 | 3.2668 | 3.1830 | **1.026x** | 7.0200 | 7.0008 | **1.003x** | | 262144 | 4.0601 | 3.9960 | **1.016x** | 9.0001 | 9.0043 | **1.000x** | **Speed - full - `BT` sweep (`V=128256`)** | BT | B200 baseline (ms) | B200 this PR (ms) | B200 speedup | H100 baseline (ms) | H100 this PR (ms) | H100 speedup | |---:|-------------------:|------------------:|-------------:|-------------------:|------------------:|-------------:| | 1024 | 0.4218 | 0.4141 | **1.019x** | 0.8335 | 0.8518 | **0.979x** | | 2048 | 0.6379 | 0.6167 | **1.034x** | 1.3555 | 1.3693 | **0.990x** | | 4096 | 1.0574 | 1.0149 | **1.042x** | 2.4686 | 2.4647 | **1.002x** | | 8192 | 1.8820 | 1.8022 | **1.044x** | 4.5362 | 4.5674 | **0.993x** | | 16384 | 3.5883 | 3.4314 | **1.046x** | 8.6929 | 8.7241 | **0.996x** | | 32768 | 6.8904 | 6.5826 | **1.047x** | 17.0411 | 17.0401 | **1.000x** | | 65536 | 13.6248 | 13.0628 | **1.043x** | 33.7193 | 33.6499 | **1.002x** | ## Correctness The full CE suite passes across bf16 and fp32 configurations, including class weights, softcap, label smoothing, z-loss, and `ignore_index`. --------- Co-authored-by: Justin Hu <181588904+justinhh4@users.noreply.github.com> Co-authored-by: Vaibhav Jindal <vaibhav.jndl@gmail.com>
kolehma8
approved these changes
Aug 3, 2026
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
Follow-up to #1279 (the merged CuTe DSL cross-entropy PR), addressing @vaibhavjindal's review
follow-ups. The headline is a memory optimization to the CuTe DSL CE backward that brings its peak
memory to exact Triton parity (it was ~2× before). Also adds the CuTe DSL benchmarking scripts
(like the cuTile ones), and — in this description — B200 speed + memory numbers vs the fully
optimized Triton CE, an H100 section (memory parity confirmed; Triton slightly ahead on speed), and
an explanation of the source of the gains.
TL;DR on speed: against Triton with every CE optimization applied (through #1268) on the standard
last-layer step, the two backends are at parity — CuTe DSL edges bf16 by ~1.0–1.18× and is within
noise on fp32. This isn't a speed PR; it's a memory-parity PR (plus benchmark tooling).
Reviewer asks addressed:
B200). Memory: exact Triton parity at every point. Speed: after a Hopper warp-count fix (arch-aware
num_warps, mirroring Triton), bf16 is parity-to-a-win on vocab / ~0.90× on the BT sweep and fp32 is
parity — see the H100 section.
run_cutedsl_compare.py, mirroringrun_cutile_compare.py).The optimization: CE backward memory parity
As merged in #1279, the CuTe DSL CE backward, on the scalar-grad path (
reduction='mean'|'sum', i.e.the normal training case), computed
_input * grad_output— which allocates a secondBT×Vbuffer. The Triton CE instead scales the saved gradient in place over the logits via a raw
element-wise kernel, so it only ever holds 1× the logits tensor. Net effect: the CuTe DSL CE used
~2× the peak memory of Triton in the fwd+bwd step (e.g. bf16, BT=8192, V=128256: 4008 → 2004 MB).
Fix (
src/liger_kernel/ops/cutedsl/ops/cross_entropy.py): scale the gradient in place, reusing thesame
element_mul_kernelthe Triton CE uses. A raw Triton kernel is used (rather than_input *= grad_output) on purpose: an in-place torch mul on the tensor returned fromforwardbumpsits autograd version counter and trips backward-through-backward; the raw kernel writes through the
pointer without that bookkeeping — exactly how the Triton CE backward dodges the same issue. Peak
memory now equals Triton across the entire sweep (tables below).
Benchmarks — CuTe DSL vs Triton (NVIDIA B200)
Speed: median of 40 iters, warmup 20 (CUDA events).
full= fwd +loss.backward()Peak memory: median of 5,
torch.cuda.max_memory_allocated. The Triton baseline is the fullyoptimized CE — every merged CE perf change plus #1268 (exp2 #1266, dtype-aware
num_warps#1267, single post-loopdx_ycorrection #1268) — so this is CuTe DSL vs Triton at its best.speedup = Triton ÷ CuteDSL (>1.00 → CuteDSL faster). Reference config: BT=8192, V=128256.
bf16
Speed — full (fwd+bwd) — vocab sweep (BT=8192)
Speed — full — BT sweep (V=128256)
Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)
Peak memory — full — BT sweep (V=128256)
fp32
Speed — full (fwd+bwd) — vocab sweep (BT=8192)
Speed — full — BT sweep (V=128256)
Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)
Peak memory — full — BT sweep (V=128256)
Notes
parity on the BT sweep (0.99–1.05×). These numbers line up with [CuteDSL] Add CuTe DSL cross-entropy kernel and CuteDSL integration scaffolding (B200) #1279's original table (which
compared against the same optimized Triton), now on the memory-fixed kernel.
the 128-bit vectorization width and doubles the bytes streamed, so the kernel is fully DRAM-bound and
both backends saturate the same bandwidth.
backward now scales the gradient in place instead of allocating a second
BT×Vbuffer) — this isthe headline of the PR.
Benchmarks — CuTe DSL vs Triton (NVIDIA H100 80GB HBM3)
bf16 (H100)
Speed — full (fwd+bwd) — vocab sweep (BT=8192)
Speed — full — BT sweep (V=128256)
Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)
Peak memory — full — BT sweep (V=128256)
fp32 (H100)
Speed — full (fwd+bwd) — vocab sweep (BT=8192)
Speed — full — BT sweep (V=128256)
Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)
Peak memory — full — BT sweep (V=128256)
Source of the gains vs Triton
First, the honest framing: with the Triton CE fully optimized (through #1268) and the standard
last-layer step (grad_output = 1.0, where the backward is a no-op), the two backends are close to
parity — CuTe DSL wins bf16 by a modest ~1.0–1.18× and is within noise on fp32. The value of this PR
is the memory parity fix; the speed is a wash-to-slight-win. Where CuTe DSL does pull ahead (bf16),
it comes from the forward kernel:
cp.asyncmulti-stage smem pipeline — 128-bit vectorized, L1-bypassing loads with apower-of-2 ring buffer, prologue prefetch, and
cp_async_wait_groupstaging. Global loads areprefetched ahead of the online-softmax compute, hiding HBM latency somewhat better than
compiler-scheduled loads on Blackwell (sm_100). This is the main bf16 edge.
(
torch.stack((count, max, min)).tolist()), theCUstreamis cached, and DLPack handles forunused optional outputs are reused. That constant ≈25–40 µs helps the small-vocab corners most.
Both backends run the same base-2 online-softmax with hardware
ex2.approx(Triton via #1266), sothe math and numerics are identical — there is no algorithmic advantage, only scheduling/overhead. In
fp32 the kernel is fully DRAM-bound and that overhead is amortized away, so Triton (which does the same
in-place-gradient bandwidth work) matches or slightly beats CuTe DSL.
Benchmark scripts
benchmark/scripts/run_cutedsl_compare.py— Triton-vs-CuTe-DSL compare driver, mirroringrun_cutile_compare.py. Runsbenchmark_<kernel>.pytwice underLIGER_KERNEL_IMPLand tags eachseries (
liger_triton/liger_cutedsl/torch) intobenchmark/data/all_benchmark_data_cutedsl.csv. Supportscross_entropy.benchmark/README.md— document both compare drivers (CuTile + CuTe-DSL).benchmark/data/all_benchmark_data_cutedsl.csv— sample CE dataset generated by the new driveron B200.
Correctness
test/transformers/test_cutedsl_cross_entropy.py: 158 passed, 1 skipped on both B200 and H100 (the not-last-layer / scalar-grad tests exercise exactly the fixed backward path;test_num_warps_matches_dtype_conventionnow asserts the arch-aware warp count on each GPU).