refactor: one KV storage abstraction with dense and paged backends #2171

Part of #2166. Phase 4a. Depends on #2170 (Phase 3: one prefill plan and one trim hook). Phase 4b builds on this.

Problem / Background

The CLI and the server run different attention kernels for the same model and the same request. The CLI's dense KVCache goes through MLX fused SDPA. On the server, pool-backed paged decode goes through mlxcel's paged kernels. The choice is made in each model forward, by the call site. So two families can disagree on when paged attention applies, a storage policy change means editing model files, and #2128 and #2162 (cuDNN SDPA nondeterminism) could affect only the CLI. Until storage is one abstraction, the epic's "same input, same output" invariant cannot hold for B=1 dense vs paged.

Current Behavior

• src/models/qwen3.rs:354 and src/models/llama3.rs:734 branch on cache.is_paged_backed() (src/lib/mlxcel-core/src/cache.rs:699) to call paged_batch_decode_attention (src/lib/mlxcel-core/src/cache/paged_batch_decode.rs:307) for a single pooled sequence.
• src/models/qwen3.rs:498 and src/models/llama3.rs:877 branch again in the batched forward.
• Batched forwards take a DecodeBatchContext (src/lib/mlxcel-core/src/generate.rs:1091-1123, which has dense(), paged() and is_paged_decode()) through forward_batched_with_context (generate.rs:1011). Models that consume it: qwen3, qwen3_5, qwen3_moe, llama3, llama4, gemma3, helium, muse_glimmer and model_owned.
• The server picks storage with --decode-storage-backend (DecodeStorageBackend, src/server/config.rs:72, field at :871). auto resolves to dense at --max-batch-size 1 (docs/CONTINUOUS_BATCHING.md:65-79).

Proposed Solution

Make attention dispatch a property of the cache, not of the model.

• Add one method on the KV cache layer, KVCache::attend(&mut self, q, k, v, scale, mask) -> UniquePtr<MlxArray>, plus a batched counterpart over &mut [&mut KVCache]. It updates storage and runs attention through whichever kernel suits the storage behind it: dense uses MLX fused SDPA (including the Turbo4 and KV-quant variants already selected inside the cache), and paged uses paged_batch_decode_attention or the gather path from ADR 0001.
• Move the l == 1 && mask.is_none() && is_paged_backed() decision, and the batched equivalent, out of the model files into the cache.
• DecodeBatchContext becomes a storage policy that the cache pool sets per sequence, so model forwards no longer read it to choose a kernel.
• Dense stays the single-sequence default wherever it is faster. The engine picks a storage policy, not a separate code path.
• Record the kernel-selection rule in ADR 0007, or in a short follow-up ADR if 0007 has already merged.

Rejected alternative: a trait KvStorage with dynamic dispatch per layer per token. It adds a virtual call in the hot loop, and the session seam in ADR 0004 already rejected per-op dispatch. Use an enum inside KVCache, which is what the cache already does for its quantization modes.

Scope

In scope: the KV cache layer in src/lib/mlxcel-core/src/cache.rs and layers.rs; the nine families above (start with qwen3 and llama3, which branch at the call site today); the cache pool's per-sequence policy; and dense vs paged parity tests.

Out of scope: new kernels, kernel tuning, and the engine step API (Phase 4b).

Implementation Notes

• Reuse: paged_batch_decode_attention, the existing dense SDPA dispatch inside KVCache (update_and_*_attention methods), and the paged-detach machinery in src/lib/mlxcel-core/src/cache/paged_detach_tests.rs. Prompt-cache detach and adopt, and the disaggregated handoff (src/distributed/disaggregated/handoff_impl.rs, which reads is_paged_backed()), must keep working unchanged.
• Constraints: no extra host sync or allocation per decode step. MLA caches (src/lib/mlxcel-core/src/mla/cache.rs) keep their own attention but implement the same attend entry. RotatingKVCache now carries an opt-in decode undo log (src/lib/mlxcel-core/src/cache/decode_undo.rs, perf(scheduler): restore Gemma 3 decode lookahead via ring rewind #2182) that Gemma 3 enables on its scheduler sequences so the decode lookahead can rewind a wrapped sliding window exactly. Single-token FP16 writes are logged in the write path (record_decode_write, with row copies only inside a DecodeLookaheadAppendScope), and every other mutation clears the log: multi-token appends, buffered and Turbo4Asym writes, trim, snapshot restore, gather_within_tail, and detach and adopt (cache/detach.rs). attend must route decode writes through that same update path and keep each clear, or rewind_decode_writes will refuse or restore the wrong rows. scheduler_model_owned_lookahead_tests and cache::decode_undo_tests cover this.
• Edge cases: a family without paged support (supports_paged_decode_backend() == false) always gets dense. A masked step (prefill or verify) never takes the paged single-token kernel, matching today's mask.is_none() guard.

Acceptance Criteria

• rg -n 'is_paged_backed\(\)|is_paged_decode\(\)' src/models returns nothing for the families touched (PR refactor: move KV attention dispatch into the cache (#2171 Phase 4a) #2210: qwen3, llama3, helium; the remaining hits are gemma3 and model_owned, the model-owned families ADR 0008 leaves for Phase 4b), so model forwards no longer branch on storage type
• (PR refactor: move KV attention dispatch into the cache (#2171 Phase 4a) #2210: before and after runs identical row for row on Qwen3-1.7B and Llama-3.2-1B; paged arm ran 1820 and 1764 paged kernel launches on Qwen3 in both) Under MLXCEL_SDPA_DETERMINISTIC=1, the Phase 0 harness shows B=1 dense and B=1 paged giving identical greedy output on Qwen3-1.7B and Llama-3.2-1B, or the difference is measured (first divergent token, logit max-abs-diff) and documented in docs/CONTINUOUS_BATCHING.md
• (left to the end-of-run measurement) Phase 0 B=1 and bench_decode numbers stay within the threshold, and the continuous-batching benches show no batched regression beyond noise
• Prompt-cache, paged-detach, disaggregated-handoff, cache::decode_undo_tests and scheduler_model_owned_lookahead_tests are green (PR refactor: move KV attention dispatch into the cache (#2171 Phase 4a) #2210, plus tests/paged_scheduler_parity.rs on the real qwen3-0.6b checkpoint)
• The new dispatch is what the CLI and the server actually run, with no parallel unused path (PR refactor: move KV attention dispatch into the cache (#2171 Phase 4a) #2210: KVCache::attend in the single-sequence forward, attend_batched in the batched forward; the dense-pointer paged route is removed from these forwards)

Verification

make verify                                   # macOS gate
make verify-fmt verify-test-cuda              # Linux/CUDA gate
cargo test --profile test-fast --features cuda -p mlxcel-core cache
MLXCEL_SDPA_DETERMINISTIC=1 make <Phase 0 parity target> MODEL=<Qwen3-1.7B-4bit path>

References verified against main at fc21784 (2026-10-07).