perf(cutedsl): fused dX+dW backward + host-enqueue caching for RMSNorm - #1336
Open
Charlesrizzly wants to merge 1 commit into
Open
perf(cutedsl): fused dX+dW backward + host-enqueue caching for RMSNorm#1336Charlesrizzly wants to merge 1 commit into
Charlesrizzly wants to merge 1 commit into
Conversation
The CuTe-DSL RMSNorm op (added in linkedin#1299) documents "fused dx/dw pass, vectorized/pipelined loads" as future work, and its device kernels are already fast enough that per-call Python launch overhead dominates at benchmark shapes. This delivers both, with no change to the public API or backend registry. Backward: - Add a fused single-pass dX+dW strip-reduction backward and make it the default, with a split fallback for shapes the fast vector path can't serve and an env escape hatch (LIGER_RMS_FORCE_SPLIT_BWD). Numerically equivalent to the existing split path within fp reduction-order noise: bf16 dW and Y are bit-identical, bf16 dX differs by <=4e-3 (well inside the parity suite's 5e-2 bf16 tolerance), fp32 matches to ~1e-5. Host enqueue (no device-code math changed): - Hoist all LIGER_RMS_* env knobs to module constants read once at import. - Add a bounded FIFO cache of marshaled CuTe tensor handles keyed on data_ptr (cap 8, override via LIGER_RMS_TENSOR_CACHE_CAP). The bound is load-bearing: an unbounded cache pins fresh output buffers and triggers a cudaMalloc storm; cap 8 keeps the speed win while holding peak memory to ~1.5x Triton. - Cache multi_processor_count per device. Perf (repo harness, median us, eager, bf16, hidden=4096): - B200 (cc 10.0): seq 8192 forward 0.93x / backward 0.83x vs Triton (wins); ~parity at 4096. Wins universally under CUDA-graph replay -- the device kernels were never the bottleneck, host enqueue was. - H100 (cc 9.0): seq 8192 ~parity (fwd 1.02x, bwd 1.10x); slower at small shapes. Peak memory ~1.5x Triton, architecture-independent. The cutedsl backend stays opt-in via LIGER_KERNEL_IMPL=cutedsl. Adds a CPU-only unit test for the fast-path eligibility helpers; the existing test/transformers/test_cutedsl_rms_norm.py parity suite is untouched. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
kolehma8
reviewed
Aug 4, 2026
| tid, _, _ = cute.arch.thread_idx() | ||
| lane = tid % 32 | ||
| warp = tid // 32 | ||
| row, _, _ = cute.arch.block_idx() |
Collaborator
There was a problem hiding this comment.
persistent kernels (e.g. have each CTA to process multiple rows) will likely produce a faster kernel than launching one CTA per row
| out_frag = cute.make_rmem_tensor((VEC,), mY.element_type) | ||
| partial = Float32(0.0) | ||
| n_vec = N_COLS // VEC | ||
| for ct in cutlass.range_constexpr(NUM_VEC_TILES): |
Collaborator
There was a problem hiding this comment.
If this is sum over the hidden states, I don't think NUM_VEC_TILES can be constant expression as that would depend on the input. Not 100% sure of this as maybe in cuteDSL that is acceptable
| for ct in cutlass.range_constexpr(NUM_VEC_TILES): | ||
| vec_idx = ct * _FAST_THREADS + tid | ||
| if vec_idx < n_vec: | ||
| cute.autovec_copy(gXv[None, vec_idx], x_frags[None, ct]) |
Collaborator
There was a problem hiding this comment.
doing the copy using int4 might be better if you are using raw CUDA threads. That will fully utilize the 4Kib page that one warp can read at the time.
| partial = Float32(0.0) | ||
| n_vec = N_COLS // VEC | ||
| for ct in cutlass.range_constexpr(NUM_VEC_TILES): | ||
| vec_idx = ct * _FAST_THREADS + tid |
Collaborator
There was a problem hiding this comment.
might be worthwhile to experiment with a design where each warp handles separate tokens. We usually have at least ~2000 tokens for training (typically more than that). With 128 CTAs and 8-16 warps for each, you can process that in single instance. That should also saturate HBM bandwidth which is likely the bottle neck here.
kolehma8
reviewed
Aug 4, 2026
| @@ -0,0 +1,42 @@ | |||
| """CPU-only checks for CuTe DSL RMSNorm fast-path launch policy.""" | |||
Collaborator
There was a problem hiding this comment.
is there any numerical tests here?
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 optimization to the CuTe-DSL RMSNorm op added in #1299. That kernel's
own docstring lists "fused dx/dw pass, vectorized/pipelined loads" as future
work, and profiling shows its device kernels are already fast enough that
per-call host (Python) launch overhead dominates at benchmark shapes. This PR
delivers the fused backward and removes that launch overhead.
No public-API or backend-registry change — the cutedsl backend stays opt-in
via
LIGER_KERNEL_IMPL=cutedsl.What changed (3 files)
ops/cutedsl/ops/rms_norm.py— the optimization:fallback for shapes the fast vector path can't serve and an env escape
hatch (
LIGER_RMS_FORCE_SPLIT_BWD).LIGER_RMS_*env knobs hoisted to moduleconstants (read once at import rather than ~12
os.environlookups perfwd+bwd pass); a bounded FIFO cache of marshaled CuTe tensor handles
keyed on
data_ptr(LIGER_RMS_TENSOR_CACHE_CAP, default 8); a per-devicemulti_processor_countcache.ops/cutedsl/ops/rms_norm_fastpath.py(new) — two pure-Python fast-patheligibility helpers (vector width + Triton-parity warp count), split out so
they are unit-testable without CUDA.
test/transformers/test_cutedsl_rms_norm_fastpath.py(new) — CPU-onlyunit tests for those helpers.
Numerics — equivalent to the merged kernel
Verified head-to-head against the current
mainkernel on identical inputs, thefused backward matches the existing split path to within floating-point
reduction-order noise:
Every delta is far inside the parity suite's bf16 tolerance (atol 5e-2). The
existing correctness oracle
test/transformers/test_cutedsl_rms_norm.pyisleft untouched.
Performance (repo harness, median µs, eager, bf16, hidden = 4096)
ratio = cutedsl / Triton, so < 1.0 is faster:
A Blackwell-favorable play: on B200 cutedsl ties Triton at 4096 and wins at
8192, and wins universally under CUDA-graph replay (the device kernels were
never the bottleneck — host enqueue was). On H100 the device-kernel margin is
smaller, so eager stays parity-or-slightly-worse. Peak memory is ~1.5× Triton at
cache cap 8, architecture-independent.
Testing
test/transformers/test_cutedsl_rms_norm.py(existing parity oracle) +test_cutedsl_rms_norm_fastpath.py(new) — run withLIGER_RMS_*/LIGER_KERNEL_IMPLscrubbed from the env.ruff check/ruff format --checkclean on all three files.Tuning knobs (all opt-in, safe defaults)
LIGER_RMS_FORCE_SPLIT_BWD,LIGER_RMS_BACKWARD_WARPS,LIGER_RMS_FUSED_STRIP_MULT,LIGER_RMS_TENSOR_CACHE_CAP,LIGER_RMS_COMPILE_BUCKET,LIGER_RMS_FORCE_NO_FAST,LIGER_RMS_RELOAD_POLICY,LIGER_RMS_AUTOTUNE_FILE,LIGER_RMS_DEBUG— eachdefaults to the fast behavior and can be trimmed if a narrower surface is
preferred.
Authored with GitHub Copilot.