Route production batched paged decode through the fused v2 kernel, retiring the gather-then-SDPA hot path

[inureyes]

Blocked by #898 (fused paged-attention decode v2 kernels and plan).

Background

The server's paged decode path does not use the fused paged-attention kernel at all. Since #720 the fused kernel is library-only: mlxcel-server --decode-storage paged routes pool-backed layers through per-sequence update_and_fetch_paged (src/lib/mlxcel-core/src/cache.rs, around line 2744) which calls pool.gather_visible to materialize a dense per-sequence K/V copy and then runs the ordinary fused SDPA (gather-then-SDPA). docs/benchmark_results/paged-attention-cuda-port-gb10-2026-07-10.md records this state, and docs/adr/0001-paged-attention-gather-vs-fused-kernel.md measures the cost: ~48% overhead at batch 4 with only 1024 tokens of context, 2x-3x past 4096 tokens, ~56%/~67% single-sequence at 16K/32K.

Since v0.4 the server defaults to batched decode (--parallel 4), which is precisely the "sustained batched decode" trigger ADR-0001 set for building the fused path. This issue closes that loop: the v2 kernel from issue #898 becomes the production decode path for paged storage.

Relevant scheduler facts: decode storage backend selection lives in effective_decode_storage_backend (src/server/batch/scheduler.rs, around line 193; paged only when max_batch_size > 1 && supports_batching && supports_paged_decode_backend); the decode step is execute_decode_step / run_decode_tick (around lines 5863/5911). The env override MLXCEL_PAGED_ATTENTION_NATIVE and parse_native_paged_override (layers.rs, around line 3476) already exist as a dispatch gate and should be preserved as the kill switch.

A design note for the plan/run contract (relevant when MLX's CUDA graph cache is in play): keeping the launch geometry fixed across steps and passing step-varying scalars (like the KV chunk size) through device memory rather than kernel constants lets a captured graph be replayed while batch composition evolves; padded grid slots carry a validity mask and exit early. Adopt this shape for the v2 launch path so graph capture never forces a recapture per step.

Task

1. Plan lifecycle in the scheduler: build the CSR plan (from Fused paged-attention decode v2: CSR page table, cross-CTA split-KV, and variable-length merge kernels #898) once per batch-composition change and cache it on the active batch; invalidate on admission, eviction/preemption, sequence finish, and whenever a sequence crosses a page boundary (new block allocated). A cheap incremental update (bump last_page_len, occasionally append a page id) is preferred over full rebuilds; measure and pick.
2. Dispatch: for models and shapes the v2 kernel supports (f16 KV pool, head_dim/GQA in the validated matrix of Fused paged-attention decode v2: CSR page table, cross-CTA split-KV, and variable-length merge kernels #898), replace the gather_visible + dense SDPA flow in the batched decode step with a single v2 kernel call over the whole batch. Keep the gather path as the fallback for unsupported shapes and as the MLXCEL_PAGED_ATTENTION_NATIVE=0 kill-switch target.
3. Variant coverage on this path:
   • Sliding-window families (RotatingKVCache / SlidingWindowKVCache users): restrict each sequence's CSR range to the visible window (adjust indptr/first-page offset and the RoPE position offsets); attention-sink retention (trim_front_keep_sink) must keep producing correct visible ranges.
   • Logits soft-cap families: fall back to the gather path initially; note the follow-up.
   • Speculative/MTP verify steps use multi-token queries; they stay on their current path (out of scope, document in code).
4. Memory accounting: the v2 workspace (partial V/LSE) must be budgeted under the existing KV-cache budget logic so admission control stays truthful (block_budget in paged.rs, memory estimate in src/execution/memory_estimate.rs).
5. Docs: update docs/CONTINUOUS_BATCHING.md (backend notes) and add a superseding note to ADR-0001 (its trigger fired; the fused path is now production). Record the decision details in a short ADR if any layout or plan-contract decision deviates from Fused paged-attention decode v2: CSR page table, cross-CTA split-KV, and variable-length merge kernels #898.
6. Rollout: default ON for paged decode storage once validation passes; --decode-storage selection logic itself is unchanged.

Performance validation (mandatory)

• Server-level benchmark, before/after, same hardware, recorded in docs/benchmark_results/paged-decode-production-<hw>-<date>.md:
  • 4 concurrent clients (the documented v0.4 default scenario), aggregate decode throughput and inter-token latency at context lengths 1K, 4K, 16K.
  • Single-sequence long-context decode at 16K and 32K.
• Required outcomes: batched decode aggregate throughput improves at 4K+ context (ADR-0001 predicts the gather overhead removed there is 2x-3x of SDPA time); no scenario in the matrix regresses beyond noise (3%).
• Run on at least one Apple Silicon machine and on GB10 when available.

Regression guard (mandatory)

• Greedy token parity: on a pinned prompt set (>= 3 model families among Llama3/Qwen3/Gemma, prompts >= 8K tokens, batch 4 with mixed lengths), greedy decode token streams must match the gather-path baseline exactly. Any mismatch is investigated and either fixed or explicitly signed off with a numerical-difference analysis in the PR (matching the repo's pinned-oracle culture).
• MLXCEL_PAGED_ATTENTION_NATIVE=0 restores the previous gather behavior end to end.
• Full server test suite and cargo test -p mlxcel-core pass; dense (--no-batch / non-paged) storage paths show zero diff in behavior and performance (spot-check single-client decode within 3%).
• Memory: admission under --parallel 4 with a constrained KV budget behaves as before (no OOM regressions; workspace accounted).

References

• src/server/batch/scheduler.rs
• src/lib/mlxcel-core/src/cache.rs (update_and_fetch_paged)
• src/lib/mlxcel-core/src/cache/paged.rs
• docs/adr/0001-paged-attention-gather-vs-fused-kernel.md
• docs/CONTINUOUS_BATCHING.md

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

---

Part of epic #909.