Fused paged-attention decode v2: CSR page table, cross-CTA split-KV, and variable-length merge kernels

[inureyes]

Background

mlxcel already ships a fused split-K flash-decoding paged-attention kernel for both backends in src/lib/mlx-cpp/turbo/paged_attention.cpp: the Metal body PAGED_ATTENTION_DECODE_SOURCE (around line 79) and its CUDA port PAGED_ATTENTION_DECODE_CUDA_SOURCE (around line 232), launched through mlx::core::fast::metal_kernel / cuda_kernel JIT strings. The Rust entry points are paged_decode_attention_pooled and the dispatch selectors in src/lib/mlxcel-core/src/layers.rs (select_pooled_paged_dispatch, resolve_pooled_paged_dispatch).

This kernel has a structural scalability limit: its KV split happens inside a single threadgroup. The launch shape is grid (32, NumSplits, B*Hq) with threadgroup (32, NumSplits, 1), where the 32 lanes partition the head dimension and NumSplits SIMD groups stripe KV tokens with online softmax. NumSplits is bounded by the 28 KB threadgroup memory budget (see paged_attention.cpp, the launch-shape computation around lines 463-505). Consequence: one CTA/threadgroup serves one (batch, q_head) pair regardless of context length, so a small batch with a long context leaves most of the GPU idle, and adding context length does not add parallelism.

docs/adr/0001-paged-attention-gather-vs-fused-kernel.md quantifies what the production fallback (gather-then-SDPA) costs instead: gather overhead vs the contiguous-SDPA lower bound is under ~15% below ~4096 tokens single-sequence, but ~56% at 16384 and ~67% at 32768; batched decode (batch 4) is already ~48% at 1024 tokens and 2x-3x past 4096. The ADR's stated trigger for building a proper fused path ("single-sequence context past ~16384, or any sustained batched decode") is met by the current server defaults (--parallel 4 batched decode since v0.4).

The scalable design, standard in the flash-decoding lineage used by modern GPU serving engines, combines four composable pieces:

1. CSR page table: flatten the whole batch's page tables into three arrays: indices (flat physical page ids), indptr ([batch+1] prefix sums delimiting each request's pages), last_page_len ([batch] valid entries in each request's final page). Inside the kernel a logical token position is converted with a fast-division by page_size into (page_iter, entry) and indices[page_iter] resolves the physical page. No gather pre-pass, one kernel launch for the whole batch.
2. Cross-CTA split-KV: split each request's KV range into chunks assigned to different CTAs. Each CTA computes an online-softmax partial result and writes (partial_V, LSE) to a workspace; the grid becomes (padded_batch_size_with_chunks, num_kv_heads). Handle GQA by making one CTA process all query heads of one KV head (block-y dimension = group size) so KV is loaded once per group.
3. Variable-length merge-state kernel: partial attention states (V_sum, lse) merge with the closed-form online-softmax rescale: m = max(lse_a, lse_b); w_a = exp2(lse_a - m); w_b = exp2(lse_b - m); out = (w_a*v_a + w_b*v_b)/(w_a+w_b); lse_out = log2(w_a+w_b)+m. A merge indptr describes how many partials each output row has, so one persistent kernel merges the whole batch. The algebra is associative and commutative, which later enables shared-prefix (cascade) composition with the same kernel.
4. Host-side plan: pick the pages-per-chunk by binary search so total CTA count matches an occupancy-derived target for the device, and emit flat index arrays (request_indices, kv_tile_indices, o_indptr) plus workspace offsets as a POD plan object that the run step consumes without recomputation.

Task

Build "paged decode v2" as a library-level capability (production wiring is the follow-up issue #899):

1. CSR view builder on PagedBlockPool (src/lib/mlxcel-core/src/cache/paged.rs): given the active batch's sequence states, produce indices: Vec<i32>, indptr: Vec<i32> (len B+1), last_page_len: Vec<i32>, and per-sequence RoPE position offsets. The existing pool layout [num_blocks, block_size, n_kv_heads, head_dim] (PagedKvLayout, paged.rs:39) already addresses any token as (block, entry) and must not change. Unit-test the builder against synthetic pools including shared (refcounted) blocks and partial last pages.
2. Kernel v2 in src/lib/mlx-cpp/turbo/paged_attention.cpp, following the existing dual JIT-string pattern (Metal + CUDA in one file, runtime backend selection): add a KV-chunk grid axis; each CTA computes an online-softmax partial over its chunk with GQA-grouped q heads (KV read once per group); when a request has one chunk, write O directly; otherwise write fp32 (partial_V, LSE) into a caller-provided workspace buffer.
3. Merge kernel (Metal + CUDA): variable-length merge over the workspace driven by a merge indptr, implementing the fp32 exp2/log2 rescale algebra above. This kernel is deliberately generic (it takes (V, LSE) arrays and an indptr) because the cascade-attention issue Cascade attention: compute shared prompt prefixes once per decode batch #903 will reuse it unchanged.
4. Plan in Rust: a PagedDecodePlan produced from (seq_lens, num_kv_heads, device parallelism): choose kv_chunk_size via binary search so total CTA count reaches a target derived from SM count (CUDA; a small C++ occupancy helper is acceptable) or GPU core count (Metal); emit the flat index arrays and workspace size. The plan must be a plain-old-data struct that can be cached and reused across decode steps while the batch composition is unchanged.
5. Entry point + flag: expose the v2 path through the existing library entry (paged_attention_decode in src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp / paged_decode_fused), selected by env MLXCEL_PAGED_ATTENTION_V2=1. Default off in this issue. Keep kernel v1 intact for comparison.
6. Correctness harness: extend examples/profile_paged_decode_kernel (or add a sibling example) to compare v2 output elementwise against the dense reference (paged_decode_attention_dense_fallback in layers.rs) across a matrix of head_dim {64, 128}, GQA ratio {1, 4, 8}, block_size {16, 32, 64}, context {512, 4096, 16384, 32768}, batch {1, 4, 8}, with randomized partial last pages.

Out of scope: server/scheduler integration (issue #899), sliding-window and softcap variants (documented fallback), speculative multi-token verify (stays on its current path).

Performance validation (mandatory)

• Extend examples/page_gather_microbench.rs to a three-way comparison: gather-then-SDPA vs fused v1 vs fused v2, over the ADR-0001 sweep (context 1K-32K, batch 1/4/8, block 16/32/64).
• Required outcomes: v2 >= v1 across the entire sweep; v2 beats gather-then-SDPA at the ADR trigger points (batch 4 at >=1024 ctx; single-sequence at >=16384 ctx).
• Record results in docs/benchmark_results/paged-decode-v2-<hw>-<date>.md following the existing report format, on at least one Apple Silicon machine, and on GB10 when a CUDA machine is available.

Regression guard (mandatory)

• Full cargo test -p mlxcel-core passes (document any pre-existing environment failures separately; none may be new).
• With MLXCEL_PAGED_ATTENTION_V2 unset, behavior is byte-identical: no new code on the default path, existing paged-cache tests and the v1 kernel tests unchanged.
• Correctness matrix vs the dense reference passes with a documented f16-KV tolerance (suggested: relative error <= 2e-2 on outputs, exact LSE algebra in fp32).
• CUDA build compiles and the harness passes on sm_121 (GB10) when hardware is available; note in the report if validated on Metal only.

References

• docs/adr/0001-paged-attention-gather-vs-fused-kernel.md
• docs/benchmark_results/paged-attention-cuda-port-gb10-2026-07-10.md
• src/lib/mlx-cpp/turbo/paged_attention.cpp
• src/lib/mlxcel-core/src/cache/paged.rs
• src/lib/mlxcel-core/src/layers.rs

Line numbers are indicative as of current main; search for the named symbols.

---

Part of epic #909.