perf(ce): single-shot dx_y gradient correction - #1268
Merged
Conversation
justinhh4
force-pushed
the
ce-opt-3-dxy
branch
2 times, most recently
from
June 26, 2026 17:48
bd42e40 to
5b13f51
Compare
Collaborator
|
@justinhh4 since #1266 and #1267 are merged, we can try merging this now. Can you update the PR description reporting just the benefits from single shot dx_y improvement? This will help us evaluate the individual impact of this change. |
Contributor
Author
|
@vaibhavjindal thanks for the suggestions. I have updated the PR body, please take a look! |
vaibhavjindal
left a comment
vaibhavjindal
left a comment
Collaborator
There was a problem hiding this comment.
LGTM, great catch. Just one minor nit about removing some comments.
| X_block += -eps | ||
| # special handle dx_y | ||
| X_block = tl.where(X_offsets != y, X_block, X_block - (1 - label_smoothing)) | ||
| # dx_y corrected once after the loop (removed the per-element tl.where) |
Collaborator
There was a problem hiding this comment.
nit: Remove the "(removed the per-element tl.where)" comment
| dloss_ori = (1 - label_smoothing) * softmax_X | ||
| # specially handle dx_y | ||
| dloss_ori = tl.where(X_offsets != y, dloss_ori, dloss_ori - (1 - label_smoothing)) | ||
| # dx_y corrected once after the loop (removed the per-element tl.where) |
vaibhavjindal
approved these changes
Jul 28, 2026
3 tasks
pull Bot
pushed a commit
to dumpmemory/Liger-Kernel
that referenced
this pull request
Aug 3, 2026
…rs (linkedin#1329) linkedin#1268 replaced the per-element `tl.where(X_offsets != y, ...)` selection with a post-loop read-modify-write of the true-class gradient: tl.store(X_ptr + y, tl.load(X_ptr + y) + dxy) X_ptr is the logits buffer, which is bf16/fp16 in mixed-precision training (FLCE computes the gradient in place). The round-trip therefore rounds softmax(x_y) / N to the buffer dtype *before* the -(1 - label_smoothing) / N term is subtracted. Since dx_y = (softmax(x_y) - 1) / N cancels catastrophically as softmax(x_y) -> 1, the pre-rounding error is amplified by roughly 1 / (1 - softmax(x_y)). Measured max relative error of dx_y against an fp64 reference (bf16, V=32000, H100): softmax(x_y) before linkedin#1268 at linkedin#1268 torch 0.90 0.24% 1.95% 0.43% 0.99 0.30% 19.99% 0.46% 0.999 0.38% 100.00% 0.42% Confident predictions are the common case once a model has trained for a while, so this silently degrades the dominant gradient signal. It surfaced as a convergence failure in test_mini_model[mini_gemma3_text-32-1e-05-dtype25-0.01-0.01-0.1-0.01-0.01-0.01], where 32 steps of accumulated drift push a top-k logprob past tolerance. Recompute dx_y once in fp32 after the loop, from the already-live ori_X_y, m and d, and write it with a single store. This keeps linkedin#1268's O(1) correction (no O(V) compare/select is reintroduced) while folding in the -(1 - label_smoothing) term before the result is rounded to the buffer dtype. Accuracy returns to the pre-linkedin#1268 values exactly (0.02% / 0.36% / 0.24% / 0.30% / 0.38%), and the speedup is retained: 3.585 ms vs 3.571 ms (this PR vs linkedin#1268) at BT=8192, V=128256 bf16 forward+backward on H100, against 3.619 ms before linkedin#1268. The existing CE suite passed with the regression because its random logits leave softmax(x_y) near zero, so add a regression test that drives softmax(x_y) -> 1 and requires the true-class gradient to be no less accurate than torch's own low-precision backward. It fails on the linkedin#1268 kernel and passes here with ~6x headroom. ## Summary <!--- This is a required section; please describe the main purpose of this proposed code change. ---> <!--- ## Details This is an optional section; is there anything specific that reviewers should be aware of? ---> ## Testing Done <!--- This is a required section; please describe how this change was tested. ---> <!-- Replace BLANK with your device type. For example, A100-80G-PCIe Complete the following tasks before sending your PR, and replace `[ ]` with `[x]` to indicate you have done them. --> - Hardware Type: <BLANK> - [x] run `make test` to ensure correctness - [x] run `make checkstyle` to ensure code style - [x] run `make test-convergence` to ensure convergence Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
pull Bot
pushed a commit
to dumpmemory/Liger-Kernel
that referenced
this pull request
Aug 3, 2026
… num_warps (linkedin#1290) ## Summary Follow-up to linkedin#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 linkedin#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: 1. ✅ **Memory** — measured; the backward is now Triton-parity (details + fix below, tables below). 2. ✅ **H100 speed/memory** — filled below from an H100 run on this branch (full sweep, same tables as 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. 3. ✅ **Source of the gains** — explained below. 4. ✅ **CuTe DSL benchmark scripts** — added (`run_cutedsl_compare.py`, mirroring `run_cutile_compare.py`). ### The optimization: CE backward memory parity As merged in linkedin#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 second `BT×V` buffer**. 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 the same `element_mul_kernel` the 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 from `forward` bumps its 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 fully optimized CE — every merged CE perf change plus linkedin#1268** (exp2 linkedin#1266, dtype-aware `num_warps` linkedin#1267, single post-loop `dx_y` correction linkedin#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)** | vocab | cutedsl (ms) | triton (ms) | speedup | |------:|-------------:|------------:|--------:| | 32000 | 0.4110 | 0.4519 | 1.100× | | 50304 | 0.5673 | 0.6717 | 1.184× | | 102400 | 0.9719 | 1.1194 | 1.152× | | 128256 | 1.1660 | 1.1863 | 1.017× | | 152064 | 1.3395 | 1.4501 | 1.083× | | 201088 | 1.7105 | 1.9816 | 1.158× | | 262144 | 2.2184 | 2.3193 | 1.045× | **Speed — full — BT sweep (V=128256)** | BT | cutedsl (ms) | triton (ms) | speedup | |---:|-------------:|------------:|--------:| | 1024 | 0.3454 | 0.3412 | 0.988× | | 2048 | 0.4655 | 0.4660 | 1.001× | | 4096 | 0.6969 | 0.7099 | 1.019× | | 8192 | 1.1651 | 1.1884 | 1.020× | | 16384 | 2.0957 | 2.1927 | 1.046× | | 32768 | 4.0186 | 4.1108 | 1.023× | | 65536 | 7.7521 | 7.9471 | 1.025× | **Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)** | vocab | cutedsl (MB) | triton (MB) | |------:|-------------:|------------:| | 32000 | 500 | 500 | | 50304 | 786 | 786 | | 102400 | 1600 | 1600 | | 128256 | 2004 | 2004 | | 152064 | 2376 | 2376 | | 201088 | 3142 | 3142 | | 262144 | 4096 | 4096 | **Peak memory — full — BT sweep (V=128256)** | BT | cutedsl (MB) | triton (MB) | |---:|-------------:|------------:| | 1024 | 250 | 250 | | 2048 | 502 | 502 | | 4096 | 1002 | 1002 | | 8192 | 2004 | 2004 | | 16384 | 4008 | 4008 | | 32768 | 8017 | 8017 | | 65536 | 16034 | 16033 | ### fp32 **Speed — full (fwd+bwd) — vocab sweep (BT=8192)** | vocab | cutedsl (ms) | triton (ms) | speedup | |------:|-------------:|------------:|--------:| | 32000 | 0.7377 | 0.6491 | 0.880× | | 50304 | 0.9406 | 0.9516 | 1.012× | | 102400 | 1.4447 | 1.6519 | 1.143× | | 128256 | 1.8539 | 1.8128 | 0.978× | | 152064 | 2.2744 | 2.2852 | 1.005× | | 201088 | 3.1161 | 3.1875 | 1.023× | | 262144 | 4.0071 | 3.9991 | 0.998× | **Speed — full — BT sweep (V=128256)** | BT | cutedsl (ms) | triton (ms) | speedup | |---:|-------------:|------------:|--------:| | 1024 | 0.4252 | 0.4224 | 0.993× | | 2048 | 0.6288 | 0.6250 | 0.994× | | 4096 | 1.0395 | 1.0222 | 0.983× | | 8192 | 1.8552 | 1.8100 | 0.976× | | 16384 | 3.5445 | 3.4421 | 0.971× | | 32768 | 6.8232 | 6.5943 | 0.966× | | 65536 | 13.3974 | 13.1099 | 0.979× | **Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)** | vocab | cutedsl (MB) | triton (MB) | |------:|-------------:|------------:| | 32000 | 1000 | 1000 | | 50304 | 1572 | 1572 | | 102400 | 3200 | 3200 | | 128256 | 4008 | 4008 | | 152064 | 4752 | 4752 | | 201088 | 6284 | 6284 | | 262144 | 8192 | 8192 | **Peak memory — full — BT sweep (V=128256)** | BT | cutedsl (MB) | triton (MB) | |---:|-------------:|------------:| | 1024 | 502 | 502 | | 2048 | 1002 | 1002 | | 4096 | 2004 | 2004 | | 8192 | 4008 | 4008 | | 16384 | 8016 | 8016 | | 32768 | 16033 | 16033 | | 65536 | 32066 | 32065 | ### Notes - **bf16**: CuTe DSL edges out the fully optimized Triton on the vocab sweep (**1.02–1.18×**) and is at parity on the BT sweep (**0.99–1.05×**). These numbers line up with linkedin#1279's original table (which compared against the same optimized Triton), now on the memory-fixed kernel. - **fp32**: essentially parity (**0.97–1.14×**); Triton is slightly ahead on the BT sweep. fp32 halves the 128-bit vectorization width and doubles the bytes streamed, so the kernel is fully DRAM-bound and both backends saturate the same bandwidth. - **Memory**: CuTe DSL now **matches Triton exactly** at every point (1× the logits tensor, since the backward now scales the gradient in place instead of allocating a second `BT×V` buffer) — **this is the headline of the PR.** ## Benchmarks — CuTe DSL vs Triton (NVIDIA H100 80GB HBM3) ### bf16 (H100) **Speed — full (fwd+bwd) — vocab sweep (BT=8192)** | vocab | cutedsl (ms) | triton (ms) | speedup | |------:|-------------:|------------:|--------:| | 32000 | 0.7215 | 0.7975 | 1.105× | | 50304 | 1.0494 | 1.1438 | 1.090× | | 102400 | 1.9960 | 1.8839 | 0.944× | | 128256 | 2.4283 | 2.2067 | 0.909× | | 152064 | 2.8312 | 2.7233 | 0.962× | | 201088 | 3.6500 | 3.6274 | 0.994× | | 262144 | 4.6563 | 4.6008 | 0.988× | **Speed — full — BT sweep (V=128256)** | BT | cutedsl (ms) | triton (ms) | speedup | |---:|-------------:|------------:|--------:| | 1024 | 0.6004 | 0.5601 | 0.933× | | 2048 | 0.8644 | 0.7973 | 0.922× | | 4096 | 1.3858 | 1.2731 | 0.919× | | 8192 | 2.4278 | 2.2088 | 0.910× | | 16384 | 4.5233 | 4.0762 | 0.901× | | 32768 | 8.7104 | 7.8105 | 0.897× | | 65536 | 17.0571 | 15.2874 | 0.896× | **Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)** | vocab | cutedsl (MB) | triton (MB) | |------:|-------------:|------------:| | 32000 | 500 | 500 | | 50304 | 786 | 786 | | 102400 | 1600 | 1600 | | 128256 | 2004 | 2004 | | 152064 | 2376 | 2376 | | 201088 | 3142 | 3142 | | 262144 | 4096 | 4096 | **Peak memory — full — BT sweep (V=128256)** | BT | cutedsl (MB) | triton (MB) | |---:|-------------:|------------:| | 1024 | 250 | 250 | | 2048 | 502 | 502 | | 4096 | 1002 | 1002 | | 8192 | 2004 | 2004 | | 16384 | 4008 | 4008 | | 32768 | 8017 | 8017 | | 65536 | 16034 | 16033 | ### fp32 (H100) **Speed — full (fwd+bwd) — vocab sweep (BT=8192)** | vocab | cutedsl (ms) | triton (ms) | speedup | |------:|-------------:|------------:|--------:| | 32000 | 1.3057 | 1.0447 | 0.800× | | 50304 | 1.9576 | 1.6267 | 0.831× | | 102400 | 3.7017 | 3.6407 | 0.984× | | 128256 | 4.5613 | 4.5006 | 0.987× | | 152064 | 5.3596 | 5.3013 | 0.989× | | 201088 | 6.9944 | 6.9316 | 0.991× | | 262144 | 9.0163 | 8.9290 | 0.990× | **Speed — full — BT sweep (V=128256)** | BT | cutedsl (ms) | triton (ms) | speedup | |---:|-------------:|------------:|--------:| | 1024 | 0.8693 | 0.8518 | 0.980× | | 2048 | 1.3944 | 1.3716 | 0.984× | | 4096 | 2.4547 | 2.4148 | 0.984× | | 8192 | 4.5600 | 4.5009 | 0.987× | | 16384 | 8.7760 | 8.6826 | 0.989× | | 32768 | 17.1979 | 17.0310 | 0.990× | | 65536 | 34.0251 | 33.6152 | 0.988× | **Peak memory — full (fwd+bwd) — vocab sweep (BT=8192)** | vocab | cutedsl (MB) | triton (MB) | |------:|-------------:|------------:| | 32000 | 1000 | 1000 | | 50304 | 1572 | 1572 | | 102400 | 3200 | 3200 | | 128256 | 4008 | 4008 | | 152064 | 4752 | 4752 | | 201088 | 6284 | 6284 | | 262144 | 8192 | 8192 | **Peak memory — full — BT sweep (V=128256)** | BT | cutedsl (MB) | triton (MB) | |---:|-------------:|------------:| | 1024 | 502 | 502 | | 2048 | 1002 | 1002 | | 4096 | 2004 | 2004 | | 8192 | 4008 | 4008 | | 16384 | 8016 | 8016 | | 32768 | 16033 | 16033 | | 65536 | 32066 | 32065 | ## Source of the gains vs Triton First, the honest framing: with the Triton CE fully optimized (through linkedin#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: 1. **Explicit `cp.async` multi-stage smem pipeline** — 128-bit vectorized, L1-bypassing loads with a power-of-2 ring buffer, prologue prefetch, and `cp_async_wait_group` staging. Global loads are prefetched 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. 2. **Lower fixed host overhead** — the three target-validation D2H syncs are batched into one (`torch.stack((count, max, min)).tolist()`), the `CUstream` is cached, and DLPack handles for unused 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 linkedin#1266), so the 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, mirroring `run_cutile_compare.py`. Runs `benchmark_<kernel>.py` twice under `LIGER_KERNEL_IMPL` and tags each series (`liger_triton` / `liger_cutedsl` / `torch`) into `benchmark/data/all_benchmark_data_cutedsl.csv`. Supports `cross_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 driver on 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_convention` now asserts the arch-aware warp count on each GPU). --------- Co-authored-by: Justin Hu <181588904+justinhh4@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Vaibhav Jindal <vaibhav.jndl@gmail.com>
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
2 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
This PR isolates one optimization on top of the already-merged CE improvements in #1266 (
exp2) and#1267 (dtype-aware
num_warps): replace the per-element true-class gradient selection with a singlepost-loop
dx_ycorrection.At the reference shape (
BT=8192,V=128256, full forward + backward), the isolated impact is:What changed
The CE gradient subtracts 1 only at the true-class index
y. Previously, both pass-2 paths evaluateda per-element
tl.where(X_offsets != y, ...)across the entire vocabulary. This PR stores the plainsoftmax gradient in the loop, then applies one scalar correction at
yafter the loop.This removes an
O(V)compare/select used to update one element and replaces it with anO(1)store.A
tl.debug_barrier()orders the loop's in-place writes before the correction reads and updatesdx_y. The correction retains the same weighting, reduction, label-smoothing, and softcap factors.Speed benchmarks
baselineismainwith #1266 and #1267 but without this PR's single-shotdx_ychange.speedup = baseline / this PR, so values above1.00xmean 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.
fullis forward plusloss.backward()withgrad_output=1.0, matching the standard last-layer CE benchmark used in #1290.bf16
Speed - full (forward + backward) - vocab sweep (
BT=8192)Speed - full -
BTsweep (V=128256)fp32
Speed - full (forward + backward) - vocab sweep (
BT=8192)Speed - full -
BTsweep (V=128256)Correctness
The full CE suite passes across bf16 and fp32 configurations, including class weights, softcap,
label smoothing, z-loss, and
ignore_index.