Stick-breaking attention kernels

What I built

Kernels

I built Triton forward and backward kernels for stick-breaking attention, merged in #1319. On H100, Nsight Compute points backward cost to fp32 scans at about 12% occupancy. That attribution is my reading, because the tool reports stall classes, not scans. On H100, 64×32 tiles fit 57-89 KiB of shared memory where 64×64 needed 129-193 KiB. The kernels still take 1.04-1.42× the original authors' time on H100. The 1.29-1.48× gain on H100 is against the PR's own earlier head, mostly from the backward rewrite.

Stick-breaking gives the nearest key a share beta of a stick, the next key beta of what is left, and so on. #1319 continues the unmerged #599. In review, the maintainer found a zero gradient at saturated logits, in bf16 as well as fp16. The fix rewrites it so nothing cancels when beta rounds to one.

Inference

I added KV-cache decoding, merged in #1412 and tested against the uncached forward. On H200, at batch 1 and 64K keys, 32 programs on 132 SMs take 4.3 ms. Reading that K/V once would take about 0.06 ms, computed from the H200's rated bandwidth, not measured. The H200 model-level step sits at about 9 ms, which I did not profile.

The decoding kernel is forward-only and runs one program per query and head, each walking that query's keys nearest first.

What I measured

On H100, on forward plus backward, the backward rewrite alone gave 1.27-1.34×, and the retile and launch order add 1.02-1.12×. On the forward alone, those two give 1.11-1.82×. The retile and launch order cost L2 reuse: in Nsight Compute on H100, the forward's hit rate fell from 92% to 78%.

On H100, the tests pass 99/99 in two runs, and 56 of the 99 fail on the earlier head. The maintainer's GB200 run also passed. My #1319 timing is H100 only. Saturated-logit coverage reaches 15 in fp16 and 30 in bf16.

On H200, 90.1× is the best of 18 op-level rows, and any cache gains like that against recomputing. Tests compare a 67-token prompt and later steps against one uncached forward over 74 tokens. They passed on H200.

At model level on H200, a cached step stays at 8.98-9.07 ms per token at every prompt and batch size I ran. Against recomputing, that ranges from 0.9×, at B=1 with 512- and 2048-token prompts, to 45.5×. The model has random weights, and each time is a median of 3 calls.

What limited it

On H100, no pipe saturates in the backward: ALU runs 22-32% and XU 14-18%. Snapshot traffic isn't the limit: dq's DRAM writes are about 3% of HBM bandwidth, and dkv reads at a 96% L2 hit rate. Registers cap occupancy: at 255 registers per thread, 2 CTAs of 4 warps fit on an SM, which is 8 of 64 warps, or 12.5%. The retile did not raise it: CTAs per SM went from 1 to 2, and each CTA from 8 warps to 4. I read the remaining cost as stalls on dependent fp32 scan and exp2/log2 results. That is inference, not a line-level measurement.

What it doesn't show

• That these kernels are faster than the authors'. On H100, tensor-core scans made the forward about 3× faster than the earlier head, but 20 of the 28 bf16 saturated-logit tests failed, so the scans stay in fp32. The backward also leaves out the authors' spin locks. I did not isolate either cost.
• Hot-loop latency. do_bench clears L2 before each call, and clocks weren't locked.
• A measured decode limit. On H200, 32 programs on 132 SMs is grid arithmetic, not a counter reading. The read time of about 0.06 ms is a floor no kernel reaches.
• A cause for the H200 step at about 9 ms. It stays flat while the op time grows with keys, and I can't reconcile the two.
• Production decode. It is one H200, forward only, no split-KV. The H200 numbers predate the rebase, which left the kernels byte-identical but changed model code around them. The non-Hopper path ran only in the CPU interpreter.

Links

• flash-linear-attention #1319 and #1412
• Upstream kernels: parallel.py and decoding.py
• Operator note with the formulation
• Full H100 results and Nsight Compute analysis on #1319
• Original authors' kernel, compared at 7619a3c