Add cutedsl implementation of Liger Rope Kernel

[Celaena24]

Summary

Adds a new NVIDIA CUTLASS Python-DSL (CuteDSL) implementation of the RoPE (Rotary Position Embedding) kernel, alongside the existing Triton kernel. It applies the rotary transform to q and k in place using the B200 Tensor Memory Accelerator (TMA) for bulk HBM movement, exposed via a drop-in autograd Function (LigerRopeCuteDSLFunction).

RoPE is a purely memory-bound op — it reads q/k once and writes them once with a little fp32 math — so the only thing that matters is achievable HBM bandwidth. This kernel runs the forward at the HBM wall and matches Triton on every axis.

Result on B200 — three headline outcomes:

• Forward at 93–100 % of HBM peak (8.2–8.9 TB/s, bf16) — mha fully saturates the bus; matches Triton's hand-tuned forward (~1.0×).
• Full-pass + backward parity: end-to-end fwd+bwd is ≈1.0× Triton across both sweeps, and ≈2× faster than HuggingFace. No regression.
• Memory: peak memory at Triton parity (within a couple MB), and ~25–30 % below HuggingFace on the full pass.

File	Change
src/liger_kernel/ops/rope_cutedsl.py	New — CuteDSL TMA fwd/bwd kernels + autograd Function
test/transformers/test_rope_cutedsl.py	New — 28 correctness tests (vs Triton and HF)
benchmark/scripts/benchmark_rope.py	adds a liger_cutedsl provider; forward times the raw op; opt-in interleaved provider measurement
benchmark/scripts/utils.py	adds opt-in interleave_providers to run_benchmarks (default off)
benchmark/scripts/benchmark_rope_cutedsl.py	New — isolated head-to-head bandwidth benchmark (interleaved-median timer)
benchmark/scripts/plot_rope_* , run_rope_*	New — B200 artifact/roofline plotters + H100 comparison runner

Details

RoPE rotation couples element j with j + hd/2 (the two halves of each head): the kernel bulk-loads a (TILE_M, head_dim) tile global→shared with one TMA instruction, does the cos/sin math in fp32 on shared memory, and bulk-stores the tile back in place. Three things make it match Triton:

• TMA bulk copies (forward). A single async instruction per tile keeps the memory pipeline saturated — the mechanism Blackwell kernels use to reach HBM peak — versus a synchronous per-thread copy that stalls on the load→use dependency (~0.45× of peak).
• Backward grad-layout discipline (the part that keeps the full pass at parity). In real attention, q/k are view(b,s,h,hd).transpose(1,2) views, so the gradient flowing into RoPE backward is non-contiguous and the leaf .grad buffer is (b,s,h,hd)-contiguous. The backward therefore operates on that native (b,s,h,hd) storage and returns a transpose(1,2) view — exactly like Triton — so PyTorch's AccumulateGrad accumulates two layout-matched tensors (coalesced). Returning a (b,h,s,hd)-contiguous grad instead makes AccumulateGrad's add_ run uncoalesced (~2.9× slower) and forces an extra .contiguous() materialization (≈+37 % peak memory). Profiling (torch.profiler + a standalone add_ microbench: 32.8 µs matched vs 92.2 µs mismatched) confirmed this was the whole gap; mirroring Triton's layout brings backward, full-pass, and memory to parity. Bonus: on the native layout the backward reuses the forward's cos/sin dedup fast path.
• Host launch-overhead elimination. The CuteDSL TMA handle (from_dlpack + memref build, ~4.4 µs/tensor) is cached and only its data_ptr field is swapped per call (~0.1 µs); config-constant Int32 scalars and the flat reshape are likewise cached. This lifts the small/medium configs off a host-bound floor to parity.

Forward — isolated kernel bandwidth (B200, interleaved-median, realistic non-contiguous HF layout)

Speedup = CuteDSL / Triton:

Config	dtype	Triton GB/s	CuteDSL GB/s	speedup
llama-8B (8k, 32×8 h, hd128)	bf16	8672	8224	0.95×
llama-8B (4k)	bf16	4669	4669	1.00×
gqa-32x8	bf16	8907	8280	0.93×
mha-32x32	bf16	7466	7473	1.00×
fp16	fp16	8827	8357	0.95×
fp32	fp32	6098	6153	1.01×

The CuteDSL forward measures 8.2–8.9 TB/s = 93–100 % of the B200 HBM peak (≈8.9 TB/s); mha already saturates. ncu confirms the kernel is HBM-bound, not occupancy- or tail-bound: raising occupancy 58→79 % (TILE_M=32) left DRAM throughput and duration flat — no kernel-side headroom.

End-to-end (forward + backward) — parity, no regression

Ratios are liger / liger_cutedsl latency (>1 ⇒ CuteDSL faster), B200, bf16, both sweeps in agreement:

mode	model_config sweep	token_length sweep
forward	1.00× (deepseek_v3 0.54×)	1.00×
backward	0.97–1.01× (dsv3 0.92×)	0.99–1.06×
full	0.99–1.01× (dsv3 0.89×)	1.00–1.06×

Forward is at parity because the op is already at the HBM wall; backward/full hold parity thanks to the grad-layout discipline above. Both Liger backends are ~2× faster than HuggingFace on the full pass. deepseek_v3 is the lone config where CuteDSL trails — its head_dim=56 (hd_half=28) isn't a multiple of the 128-bit vector width, so it falls back to the portable (latency-bound) token-per-CTA kernel instead of the TMA path. Reported honestly rather than hidden; widening TMA coverage to awkward head_dims is logged as future work.

Memory — Triton parity

The forward writes in place (forward peak bit-identical to Triton). On the full fwd+bwd pass, the backward grad-layout fix removes the extra .contiguous() copy, so CuteDSL's peak matches Triton and sits well below HuggingFace. B200, bf16, full peak (MB):

Config	HF	Triton (liger)	CuteDSL	vs HF
llama_3_8b	144.4	112.5	113.5	−21 %
qwen2.5_7b	120.4	92.5	94.5	−22 %
qwen2.5_72b	256.4	192.5	196.5	−23 %
deepseek_v3	347.9	291.9	297.9	−14 %
T=4096	281.6	217.7	219.2	−22 %
T=8192	547.6	419.7	423.2	−23 %

(The head-to-head benchmark_rope_cutedsl.py reports exact parity, e.g. llama-8B 180/180 MB; the few-MB deltas above are CuteDSL compile-pool noise, not a per-call allocation.)

Figures

Model-config sweep — real model widths (llama / qwen / deepseek), three providers (huggingface / liger / liger_cutedsl), bf16, T=2048:

Forward speed — cutedsl ≈ liger
Backward speed — cutedsl ≈ liger
Full speed — all at parity, ~2× HF
Full memory — cutedsl == liger, both below HF

Token-length sweep — llama-3-8B width, bf16, T = 1024→8192:

Forward speed vs T
Backward speed vs T
Full speed vs T
Full memory vs T — exact parity

Isolated kernel + cross-GPU + high-T (B200):

Isolated kernel speedup vs T
Isolated kernel across model widths
Isolated forward GB/s, high-T
Isolated backward GB/s, high-T
Full ms, high-T
Full speed, B200 vs H100
Full memory, B200 vs H100

Contiguous Gradients

Additionally, when the incoming gradient arrives already contiguous (eager attention, or any post-RoPE  .contiguous() /reshape), the backward detects it and rotates in place — skipping the copy Triton always pays on that layout — for up to ~1.9× faster backward and ~19% lower peak memory in that regime

Testing Done

• test/transformers/test_rope_cutedsl.py — 28/28 pass: forward + backward vs both Triton and the HuggingFace reference. Coverage: 7 shapes including standard (1,128,32,32,64) / (2,128,32,32,64), GQA (…,32,8,…), non-tile-aligned (3,423,73,213,92) / (3,423,73,155,92) exercising the token-fallback path, and the perf-relevant long sequence (1,8192,32,8,128); fp32 (atol 1e-5) and bf16 (atol 1e-1); both expand_position_ids ∈ {True, False}.

• Hardware Type: NVIDIA B200 (sm_100)

• run make test to ensure correctness (test_rope_cutedsl.py 28/28)

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

[kolehma8]

would warps be better idea here since if the number of threads exceeds the number of columns this would idle them?

[Celaena24]

According to my understanding, the threads wouldn't actually idle here right? hg = tid // HD_HALF (line 88) would reuse them for the next head so they stay busy. Threads only idle in the last partial warp which is unavoidable.

[kolehma8]

right, I was not thinking this one fully

[kolehma8]

this makes a copy, would it be possible to write this without having to create copy of the dq and instead do the index math within the kernel?

[Celaena24]

@kolehma8 I changed the code so that when the grad comes in already contiguous, it rotates in place with the seq-index math done inside the kernel, so there's no copy. This works with both the TMA kernel and the token kernel, whichever the head_dim supports, and both stay fully coalesced, so it comes out ~5× faster for TMA and ~2.5× faster for token kernel than copying first.

When the grad arrives as a  transpose-view of  (bsz, seq, n_head, hd)  storage, the  transpose(1,2)  exposes the already-contiguous storage and  .contiguous()  is a no-op. The kernel rotates the original storage in place and there is no copy

But for grad layouts where the head dimension is not the innermost, we need  .contiguous() so that we can use TMA. If we don't do .contiguous(), we can do index-math-in-place but we would have to take the non-TMA token kernel with uncoalesced access.

I benchmarked copy-then-fast-kernel vs. index-math-in-place, and the copy wins 2–8×.
Because RoPE's cost is purely memory movement, and  .contiguous() moves data in the efficient (coalesced, full-bandwidth) way, that "extra" copy is cheap enough that copy-then-fast-kernel beats skipping the copy and running slow uncoalesced access.

[kolehma8]

TMA should allow you to get away without the transpose as it performs very well on strided data access. Coalescence is important when warps reads the data as you want to max out the 4kB page.

[Celaena24]

@kolehma8 Yeah that's true, TMA's descriptor handles strided outer dimensions fine, but its innermost box dimension must be stride-1 (contiguous).

In the case where we are doing .transpose, it works because head_dim is already innermost-contiguous. So I could technically drop the transpose and feed the (b,h,s,hd) view straight to TMA, but then the rows have a non-uniform stride, so I'd need a multi-dim strided descriptor that issues the exact same loads but with more complexity and no speedup, and I'd still have to transpose the returned grad for AccumulateGrad to stay coalesced.

The transpose here isn't actually costing anything.  transpose(1,2)  is a metadata-only view op (it just swaps strides, shares storage) so there is no copy in this case.

[kolehma8]

should you pipeline the TMA load, compute and TMA store? That is a big benefit that comes with TMA.

[Celaena24]

@kolehma8
I built and benchmarked the pipeline version by sweeping over blocks/SM {1,2,4,6} × NSTAGE {2,3,4} on B200. It loses in every config — best (b6s2) is 0.91x–0.94× the current kernel.

The reason is that pipelining's main benefit is that it does not stall while waiting on a load. But prefetching the next tile while computing the current one is already provided here by occupancy, so it's redundant:

• The load-wait is already hidden, across blocks instead of within one. A current-kernel block does stall on its own load, but during that wait the SM just runs one of the other ~12 resident blocks. So pipelining doesn't hide new latency, it hides the same latency in a more expensive way. To get 3 stages, each block needs 48 KB, so only 4 blocks/SM fit instead of 13, and occupancy drops from 57.6% → 24.3% active warps, and DRAM falls with it (61.5% → 47.6%).

• There's almost nothing to overlap the load with. Pipelining overlaps a load with compute, but RoPE's per-tile compute is a few fp32 FMAs (~tens of cycles) vs a ~hundreds-of-cycles load. So a shallow prefetch hides almost none of the load latency. To hide it fully you'd need many loads in flight (a deep pipeline) — but "many loads in flight" is just memory-level parallelism, which is SMEM-bounded (~13 buffers/SM) either way and costs occupancy to buy the stages.

• In-place rotation also adds a bubble that pipeline can't avoid. RoPE overwrites its input, so each tile's SMEM buffer is both the store source and the next load's destination. We have to wait for the store to drain before reloading it. The 1-tile-per-CTA design never pays this. Each CTA stores once and exits, and the store drains in the background while other CTAs run.

And there's no idle-TMA gap to reclaim. The current kernel is already ~62% DRAM, matching the hand-tuned Triton kernel which is the HBM wall for this op. So the best pipelined variant just re-reaches, at extra cost and complexity, the latency-hiding the simple high-occupancy design already gets for free.

[kolehma8]

Thanks for this analysis. This is quite surprising as I would think that the TMA would get choked by 12 warps issuing instructions per SM.

[kolehma8]

Approved the PR, I guess what you said makes sense that the 12 CTAs enqueue the TMA instructions and the TMA is running hot even without pipelining as the queueing comes from the separate CTAs. I did not think of this angle before.

[kolehma8]

Why 12 blocks per SM? There is only single TMA unit per SM so this seems unnecessary. Also, running more blocks will increase the pipeline bubble (if you do pipelining).

[Celaena24]

@kolehma8 Around 12 blocks is just how many CTAs fit on an SM once each one uses 16 KB of shared memory. Also Rope is memory bound so is it not better to have more transfers in flight?

I tested using fewer blocks directly on B200 and it is consistently slower. On llama-8B, 1 → 2 → 4 blocks/SM gives 0.45× → 0.67× → 0.89× of the current kernel.

[kolehma8]

Might be interesting experiment to see what happens with 1 CTA per SM, and only issuing the TMA load instructions. That should already max out the BW.

[vaibhavjindal]

@Celaena24 let's rebase this one on top of the newly merged changes for cutedsl: #1279