[CuteDSL] FLCE: 2-SM online-softmax fast path (Blackwell), ~5x over Triton#1291
Open
justinhh4 wants to merge 1 commit into
Open
[CuteDSL] FLCE: 2-SM online-softmax fast path (Blackwell), ~5x over Triton#1291justinhh4 wants to merge 1 commit into
justinhh4 wants to merge 1 commit into
Conversation
…riton Adds the CuTe DSL fused-linear-cross-entropy (forward + backward), signature-compatible with the Triton LigerFusedLinearCrossEntropyFunction and selected via LIGER_KERNEL_IMPL=cutedsl. The forward is store-free: built on quack's 2-SM (CtaGroup.TWO) tcgen05 GEMM, it computes only the per-row LSE + the target logit in a fused epilogue and never materializes the (BT, V) logits (not even per chunk). The backward recomputes the masked softmax on the fly, tiled over the vocabulary so the grad_weight GEMM contracts over K=BT and writes each slice once. Net on B200 (bf16): up to ~5x faster than the Triton FLCE at LLM vocab sizes and ~1.5-1.9x less peak memory. fp32 / non-Blackwell / softcap / argmax-metrics delegate to the upstream Triton FLCE (the 2-SM GEMM is 16-bit-input only), so those paths are byte-identical to and exactly as fast/lean as Triton. Full feature parity with the Triton FLCE (ce_weight, bias, softcap, label_smoothing, z-loss, ignore_index, token-accuracy, predicted-tokens, use_token_scaling, accum_dtype, mean/sum, forward-only reduction='none'). reduction='none' with grad is refused by design (matches the Triton contract). Tests: 277 passed (LIGER_KERNEL_IMPL=cutedsl). The AMP+bias=False parity tests exercise a Triton addmm-under-autocast path fixed by linkedin#1283; stack this on linkedin#1283. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.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
1 participant
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 CuTe DSL cross-entropy PR). This adds the CuTe DSL fused-linear-cross-entropy (FLCE) — forward + backward, signature-compatible with the Triton
LigerFusedLinearCrossEntropyFunction, selected viaLIGER_KERNEL_IMPL=cutedsl.Unlike a straightforward port, this kernel is built on quack's 2-SM (
CtaGroup.TWO)tcgen05GEMM for Blackwell and a store-free online-softmax forward: the(BT, V)logits are never materialized (not even per chunk). The forward computes only the per-row LSE + the target logit in a fused GEMM epilogue; the backward recomputes the masked softmax on the fly, tiled over the vocabulary so thegrad_weightGEMM contracts overK = BTand writes each slice once. The result is both faster and leaner than the Triton FLCE on B200 — up to ~5× faster at LLM vocab sizes and ~2× less peak memory.Feature parity with the Triton FLCE:
ce_weight,bias,softcap,label_smoothing, z-loss (lse_square_scale),ignore_index, token-accuracy, predicted-tokens,use_token_scaling,accum_dtype, reductionsmean/sum, plus forward-onlyreduction='none'.reduction='none'with grad is refused by design (the fused pass accumulatesgrad_weightover tokens and cannot honor a per-token upstream grad — matching the documented FLCE contract, and the Triton FLCE's own limitation).Files
src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py— the FLCE Function + fwd/bwd (2-SM online-softmax fast path + Triton-delegated fallback).src/liger_kernel/ops/cutedsl/ops/__init__.py— export the FLCE symbols soLIGER_KERNEL_IMPL=cutedslrewires the public FLCE too.src/liger_kernel/ops/cutedsl/ops/utils.py—infer_device_arch(the Blackwell fast-path gate).test/transformers/test_cutedsl_fused_linear_cross_entropy.py— parity suite (below).setup.py— addquack-kernels~=0.5.0to thecutedsloptional-deps group.Correctness
test/transformers/test_cutedsl_fused_linear_cross_entropy.py— 277 passed (LIGER_KERNEL_IMPL=cutedsl pytest ...). Full parity vs the Triton FLCE (and vstorchwhere applicable) across dtype × reduction × every feature × grad, plus: AMP/autocast, not-last-layer scalar grad, non-contiguous inputs, the functional API +CrossEntropyOutputstructured output,LIGER_KERNEL_IMPL=cutedsldispatch wiring, production vocab (V=32000, the partial-tile tail-predication path), and the host-side validation / by-design-refusal contracts.Benchmarks — CuTe DSL vs Triton FLCE (NVIDIA B200)
Interleaved A/B, median p50 of the
full(fwd+bwd) mode,grad_output=1.0, H=4096, bf16. speedup = Triton ÷ CuteDSL (>1.00 → CuteDSL faster). Peak memory ismax_memory_allocatedfor the samefullpass. Reference config: BT=8192, V=128256.bf16 — full (fwd+bwd), vocab sweep (BT=8192, H=4096)
bf16 — full (fwd+bwd), BT sweep (V=128256, H=4096)
Takeaways. CuteDSL FLCE is faster at every shape and ~1.5–1.9× leaner in peak memory (the store-free online-softmax forward never materializes chunk×V logits, unlike the Triton chunked path). The speedup is largest at the memory-bound regime (large V, small–moderate BT) where Triton's chunked cuBLAS path is dominated by the logit materialization + separate softmax passes; at very large BT both approach the GEMM roofline and the lead narrows to ~1.2×.
fp32 (delegates to Triton) — full, vocab sweep (BT=8192, H=4096)
fp32 has no Blackwell tensor-core GEMM path (the 2-SM
tcgen05GEMM is 16-bit-input only), so it delegates to the Triton FLCE and is tied on speed and memory by construction:Comparison with other memory-efficient linear-CE kernels (NVIDIA B200, bf16)
For context, here is the CuteDSL FLCE swept against the two well-known store-free / chunked peer libraries: Apple Cut-Cross-Entropy (
cut-cross-entropy25.9.3, default preset = gradient filtering + bf16 grad accum) and quack (Dao-AILabchunked_linear_cross_entropy0.5.0, chunk_size=4096). (The Triton FLCE baseline is in the tables above.) All bf16,full(fwd+bwd) p50 latency (ms) + peakmax_memory_allocated(MB). The(N×)next to each peer's ms and MB is its value ÷ cutedsl — for both,>1×means cutedsl is better (faster / leaner) and<1×means the peer is. By design class: cutedsl and CCE are store-free (no chunk×V logits); quack is chunked-store.vocab sweep (BT=8192, H=4096)
BT sweep (V=128256, H=4096)
Reading the frontier. cutedsl and CCE are both store-free and sit at ~half the peak memory of the chunked-store quack across the sweep. Within that lean class, cutedsl is dramatically faster than CCE — ~4.8× at V=128k, and the gap widens with BT (CCE's backward scales poorly: 25.7 vs 123 ms at BT=8k → 194 vs 983 ms at BT=65k, a ~5× lead). CCE's own peak is a hair below cutedsl (it stores nothing but a running LSE; 3422 vs 3714 MB at V=128k), so the honest framing is: cutedsl and CCE tie on memory, and cutedsl is 4–5× faster.
vs quack — quack is the fastest kernel measured (~1.3× faster than cutedsl — the
0.6–0.8×column — roughly flat across the sweep, since its 3-GEMM chunked-store does less total work), but it pays ~1.5–2× the peak memory (chunked-store class). A clean speed-for-memory trade: quack wins throughput, cutedsl wins memory at near-matched throughput and is the only kernel that is both store-free and faster than Triton/CCE.