[inureyes]

Summary

Sub-step (a) of #121 (epic #116): wire the live batch scheduler's paged decode backend onto the PagedBlockPool that #118–120 built and #152 proved through a real model. Path A — the pool rides inside the KVCache, so the LanguageModel trait and forward signature stay unchanged.

This is one of three sub-steps and does not close #121 — the radix adopt/donate + try_adopt_cached_prefix guard removal lands in sub-step (b), and the multi-model RMS sweep in sub-step (c).

What changed

Phase 1 — CachePool Rc refactor (behaviour-preserving)

• CachePool.paged_pool → Option<Rc<RefCell<PagedBlockPool>>> and SequenceCacheSet.paged → Option<Rc<RefCell<PagedSequenceState>>>, so a pool-backed sequence shares ONE authoritative block table across its per-layer caches (each cache's PagedBacking holds a cheap Rc clone of the same handles).
• Updated every access site — allocate*, detach_paged/adopt_paged/release_detached_paged/park_* (paged_detach.rs), sync_paged_state_with_*, paged_stats/stats_for_sequences, restore_paged_state, ensure_paged_pool*, release, and the paged_pool_ref/paged_pool_mut/get_paged_state/paged_state accessors (now hand out Ref/RefMut guards) — with tightly-scoped borrows. Request-boundary pool/state borrows never nest with the forward's update_and_fetch borrows; adopt_paged drops every pool borrow before the prepare_sequence_state callback.
• No behaviour change: the Phase 1: Global block-pool tensor storage #118–120 detach/adopt/round-trip suites pass unchanged (adapted only for the Ref/RefMut type change).

Phase 2 — live paged wiring

• CachePool::allocate builds pool-backed caches (KVCache::new_paged) for paged sequences whose model's natural backend is the dense external KVCache slice (qwen3/llama3) under an Fp16 layout. Model-owned families (gemma3/llama4/qwen3_5, which drive their own ModelOwnedSequenceState and ignore the caller's caches) and Turbo4 layouts keep their existing dense placeholders — byte-identical to before.
• Added KVCache::is_paged_backed(). The qwen3/llama3 batched-decode dispatch returns None from its native-paged-kernel selection when any cache is pool-backed, falling through to the per-sequence update_and_fetch loop — the feat: transparent pool-backed KVCache + real-model paged parity test #152-validated pool path (write_prefill + gather_visible). The gemma3/llama4 shared dispatch helpers are untouched because those models never receive pool-backed caches.
• KVCache::update gained the same pool intercept as update_and_fetch (factored into a shared write_paged), so llama3's single-sequence fused causal prefill path — which calls update without a fetch — also writes to the pool. qwen3's single-sequence prefill already routes through update_and_fetch.
• sync_paged_state_with_dense skips pool-backed sequences (the pool is authoritative — write_prefill advances state.len in lockstep with the cache offset). trim/trim_front are dense-side no-ops for pool-backed caches (they would otherwise desync from / unwrap the absent dense buffers; paged trimming is a pool-API concern).
• Hybrid-SSM carve-out: confirmed already in place — effective_decode_storage_backend gates on supports_paged_decode_backend() (trait default false) and supports_batching() (which mamba/mamba2/jamba/nemotron_h/nemotron_nas/recurrent_gemma/qwen3_next/rwkv7 override to false); either flag forces the dense backend. Added a test pinning the combined SSM profile across batch sizes.

Phase 3 — verification

• New tests/paged_scheduler_parity.rs (#[ignore], skip-if-absent) drives the scheduler's paged allocation path via CachePool::allocate_with_layout (the exact API BatchScheduler::allocate_sequence_state calls for the paged backend) and asserts identical greedy tokens to the dense backend, for both the single-sequence path and a batched B=2 decode (which exercises the is_paged_backed() decode guard). BatchScheduler::new is pub(crate), so this CachePool-driven harness is the faithful constructible driver from an integration test.

Drive-by

• Fixed a pre-existing clippy::identity_op in src/lib/mlxcel-core/src/utils.rs (0 + 2 → 2) that cold clippy fails on (from perf: support variable-length prompts in B>1 batched MTP bursts #162).

Test plan

• cargo test --lib -p mlxcel-core --features metal,accelerate cache:: — 370 pass (incl. 2 new sync tests)
• cache::paged_detach / cache::detach regression suites — unchanged, pass
• cargo test --lib --features metal,accelerate hybrid_ssm_workers_never_use_paged_decode_backend — pass
• cargo clippy -p mlxcel-core --tests --features metal,accelerate -- -D warnings (cold) — clean
• cargo clippy --lib --features metal,accelerate -- -D warnings — clean
• cargo check --test paged_scheduler_parity --features metal,accelerate — compiles
• cargo fmt --check — clean

Real-model parity test (loads qwen3-0.6b-4bit on the GPU; run by the orchestrator):

cargo test --test paged_scheduler_parity --release --features metal,accelerate -- --ignored --nocapture

Part of #121

Follow-up fixes (after the real-model parity run)

The orchestrator's paged_scheduler_parity real-model run surfaced two issues, now fixed on this branch (both crates cold-clippy clean; core regression 832 pass; scheduler/batch 173 pass; parity now passes for BOTH single-sequence and batched B=2 — qwen3-0.6b-4bit, 16 greedy tokens identical to dense):

fast_rope_batched stride aliasing — pre-existing, affects dense batched decode too

Batched B=2 parity initially diverged (row 1 looped). Root cause: fast_rope_batched's uniform-batch fast path collapsed to a single fast_rope over the full [B, n_heads, T, D] tensor. In T==1 batched decode the Q/K/V arrive transposed from [B, T, n_heads, D], where swapping the size-1 sequence dim leaves stride[batch] == stride[sequence]; MLX's fast::rope kernel then rotates batch element b to position offset + b instead of offset. This silently corrupted RoPE for batch rows past the first even with identical offsets — a latent bug in dense batched decode (the function's existing callers) that the paged same-prompt test exposed. Fix: drop the uniform fast path, always go per-row (B==1 cannot alias) + regression test test_fast_rope_batched_non_contiguous_t1_symmetric.

Int8/Turbo3 paged → dense (pool-backing is Fp16-only)

The Fp16 pool-backing gate keyed on the layout's cache_mode, but the scheduler built Int8/Turbo3 paged layouts with cache_mode = Fp16, so those sequences were pool-backed and stored as Fp16 (correct output, lost Int8 memory saving). Fix: the scheduler now builds the non-Turbo4 paged layout with the real effective_mode, so the gate pool-backs only genuine Fp16 and Int8/Turbo3 keep their dense path. Pool-native quantized storage is a later phase.