[CuteDSL] Add CuTe DSL fused-linear-cross-entropy (FLCE) - #1287
Closed
justinhh4 wants to merge 4 commits into
Closed
[CuteDSL] Add CuTe DSL fused-linear-cross-entropy (FLCE)#1287justinhh4 wants to merge 4 commits into
justinhh4 wants to merge 4 commits into
Conversation
Add an optional CuTe DSL (nvidia-cutlass-dsl) backend implementing Liger cross-entropy forward+backward as a single fused kernel, signature-compatible with the Triton LigerCrossEntropyFunction. Supports z-loss, softcap, label smoothing, class weights, ignore_index, token-accuracy and predicted-tokens, across bf16/fp16/fp32. B200 optimizations: online-softmax pass with a cp.async multi-stage pipeline, FMA-folded exp2 arguments, power-of-2 ring-buffer stage advance, and reduced per-call host overhead (cached CUstream, batched validation D2H sync, and reused DLPack handles for unused optional outputs). Adds the `cutedsl` optional dependency (nvidia-cutlass-dsl) and parity tests against the Triton kernel (158 cases). The fused-linear-CE (FLCE) variant will follow in a later PR. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Follow-up to linkedin#1279 (the CE PR promised the FLCE variant "in a later PR"). Adds the CuTe DSL FLCE (forward + backward), signature-compatible with the Triton LigerFusedLinearCrossEntropyFunction. It reuses the CuTe DSL CE kernel per token-chunk and mirrors the upstream Triton FLCE control flow byte-for-byte — memory-minimal chunk sizing, fp32 loss/grad accumulators, torch.addmm(out_dtype=fp32) grad_weight accumulation (PR linkedin#1239), and the token-scaling transform — so the ONLY difference between the two FLCE backends is the CE kernel. Full feature parity: ce_weight, softcap, label_smoothing, z-loss, ignore_index, token-accuracy, predicted-tokens, use_token_scaling, accum_dtype, reductions mean/sum, plus forward-only reduction='none' (none+grad is refused by design, since the fused pass accumulates grad_weight over tokens and cannot honor a per-token upstream grad). - src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py: the kernel wrapper. - src/liger_kernel/ops/cutedsl/ops/__init__.py: export the FLCE Function + fwd/bwd. - test/transformers/test_cutedsl_fused_linear_cross_entropy.py: cutedsl-vs-Triton (and vs torch) parity across dtype x reduction x every feature x grad, AMP, functional API + structured output, dispatch wiring, production vocab (V=32000 tail predication), and the host-side validation / by-design-refusal contracts. 186 passed, 1 skipped on B200. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The bf16 atol=1e-8 (rtol-only) bar failed test_ce_token_accuracy_matches_triton[mean-bf16] on B200: the logit=50 peaked-softmax rows produce near-zero grad entries where bf16's coarse mantissa lands ~8.7e-6 from Triton and rtol*|ref| collapses toward zero, so atol must carry them. Set bf16 -> (5e-5, 5e-2): covers the measured ~8.7e-6 with ~5.7x headroom, still 100x tighter than the original 5e-3. Same near-zero-grad phenomenon already documented for fp16/fp32. Full suite: 158 passed, 1 skipped on B200. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
Author
|
Superseded by #1291 |
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, whose body notes: "The fused-linear-CE
(FLCE) variant will follow in a later PR."). This is that PR, stacked on
justinhh4/cutedsl-cross-entropy.Adds the CuTe DSL FLCE (forward + backward), signature-compatible with the Triton
LigerFusedLinearCrossEntropyFunction. It reuses the CuTe DSL CE kernel from #1279 pertoken-chunk and mirrors the upstream Triton FLCE control flow byte-for-byte — memory-minimal
chunk sizing, fp32 loss/grad accumulators,
torch.addmm(out_dtype=fp32)grad_weight accumulation(PR #1239), and the token-scaling transform — so the only difference between the two FLCE
backends is the CE kernel.
Feature parity:
ce_weight,softcap,label_smoothing, z-loss,ignore_index,token-accuracy, predicted-tokens,
use_token_scaling,accum_dtype, reductionsmean/sum,plus forward-only
reduction='none'.reduction='none'with grad is refused by design (thefused pass accumulates
grad_weightover tokens in the forward and cannot honor a per-tokenupstream grad — refusing loudly beats silently mis-scaling), matching the documented FLCE
contract.
Files
src/liger_kernel/ops/cutedsl/ops/fused_linear_cross_entropy.py— the FLCE Function + fwd/bwd.src/liger_kernel/ops/cutedsl/ops/__init__.py— export the FLCE symbols (soLIGER_KERNEL_IMPL=cutedslrewires the public FLCE too).test/transformers/test_cutedsl_fused_linear_cross_entropy.py— parity suite (below).Correctness
test/transformers/test_cutedsl_fused_linear_cross_entropy.py— 186 passed, 1 skipped (theskip is the
LIGER_KERNEL_IMPL=cutedsldispatch test, which passes when that env varis set).Full parity vs the Triton FLCE (and vs
torchwhere applicable) across dtype × reduction × everyfeature × grad, plus: AMP/autocast, the functional API +
CrossEntropyOutputstructured output,LIGER_KERNEL_IMPL=cutedsldispatch wiring, production vocab (V=32000, the partial-tiletail-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, H=4096.speedup = Triton ÷ CuteDSL (>1.00 → CuteDSL faster). Reference config: BT=8192, V=128256.
FLCE is matmul-dominated (the
x @ Wᵀlogits + the two backward GEMMs dwarf the CE kernel), andthose matmuls are the same cuBLAS calls on both backends — so CuteDSL FLCE runs at parity
with Triton FLCE, edging ahead where the CE kernel's slightly-faster forward shows through (large
BT bf16).
bf16
full (fwd+bwd) — vocab sweep (BT=8192, H=4096)
full — BT sweep (V=128256, H=4096)
fp32
full (fwd+bwd) — vocab sweep (BT=8192, H=4096)
full — BT sweep (V=128256, H=4096)
Peak memory (the FLCE headline win) —
full, BT=8192, V=128256, H=4096CuteDSL FLCE matches Triton FLCE's peak memory exactly (both avoid materializing the BT×V logits),
~2.2× under the stock
Linear → CrossEntropypath.