feat: add physical block-pool K/V tensor storage to PagedBlockPool

[inureyes]

Summary

Phase 1 of epic #116 (unified paged KV cache). Adds the physical main K/V tensor storage to PagedBlockPool plus the write/gather primitives, byte accounting, and unit tests, following layout A and the gather-then-SDPA strategy locked by ADR 0001 (docs/adr/0001-paged-attention-gather-vs-fused-kernel.md).

This is an additive capability that later phases consume. The live model/decode flow is intentionally not rewired here (decode-read is #119, prefill-write is #120, fused kernel is #123), per the project's integration rule that integration depending on later sub-issues is deferred to them. Net effect in live operation after this PR: pool tensors stay unallocated (no live writer wired yet), dense KVCache placeholders still carry live K/V, the existing paged decode path is byte-for-byte unchanged, and all existing paged tests pass. The new storage is exercised and proven only by the new unit tests, ready for #119/#120 to wire.

What changed

• src/lib/mlxcel-core/src/cache/paged.rs
  • New per-layer physical pool state on PagedBlockPool: pool_k / pool_v (Vec<Option<UniquePtr<MlxArray>>>, layout A [capacity_blocks, block_size, n_kv_heads, head_dim], separate K and V per layer), pool_meta (lazily-inferred n_kv_heads/head_dim/dtype/capacity_blocks), block_rows (per-layer block_id -> row), free_rows, next_row. PagedBlockPool::new(layout) signature is unchanged (tensors start None).
  • write_block(block_id, layer_idx, slot_start, k_block, v_block) — lazily allocates the layer pool on first write, assigns/looks up the layer-local row (reusing freed rows), grows in POOL_GROW_CHUNK_BLOCKS (32) chunks when a fresh row exceeds capacity, and writes via slice_update reassigning the pool tensor so MLX donates the buffer (O(block) append per ADR 0001). Accepts both [1, n_kv_heads, n_slots, head_dim] (SDPA-style, primary) and [n_slots, n_kv_heads, head_dim] (bare slab) block layouts; validates shape/dtype against captured geometry on every write.
  • gather_visible(state, layer_idx) — builds the physical-row index from the layer's block table (fragmented / out-of-order rows gather in block-table order), takes the rows, flattens, slices the visible [logical_start, len) window, and transposes into [1, n_kv_heads, visible_len, head_dim], byte-identical to the equivalent dense contiguous buffer's visible slice. Returns Ok(None) for an empty layer.
  • pool_tensor_bytes() — sums array_nbytes over allocated pool_k/pool_v.
  • release_block now frees the main-K/V row (push to free_rows, remove from block_rows) in addition to the existing turbo_sidecars.forget.
  • No astype on the K/V path: dtype is preserved through transpose_axes/reshape/slice_update/take/slice. Fp16 stays Fp16, Int8 stays Int8. Int8 quantization scales stay in the dense path and are not touched.
  • Updated the PagedBlockPool doc comment to state the pool can now own physical main K/V storage (layout A) and that dense-placeholder removal lands with Phase 2: Paged decode attention over real block tables #119/Phase 3: Paged prefill into the block pool #120.
• src/lib/mlxcel-core/src/cache.rs
  • Re-export the new GatheredKv type alias; wire the paged_pool_tests module.
  • CachePool::memory_usage_bytes adds pool.pool_tensor_bytes() additively next to turbo_sidecar_bytes(); the layout-derived reserved_bytes/used_bytes/stats_for_sequences budgets are unchanged.
• src/lib/mlxcel-core/src/cache/paged_pool_tests.rs (new) — 13 unit tests covering the acceptance criterion and all required cases.

Test plan

• cargo check --lib -p mlxcel-core --features metal,accelerate — clean
• cargo clippy --lib -p mlxcel-core --features metal,accelerate -- -D warnings — clean
• cargo clippy -p mlxcel-core --tests --features metal,accelerate -- -D warnings — clean
• cargo test --lib -p mlxcel-core cache::paged_pool --features metal,accelerate — 13 passed
• cargo test --lib -p mlxcel-core cache::paged_turbo --features metal,accelerate — 26 passed (existing nbytes_* budget assertions intact)
• cargo test --lib -p mlxcel-core cache::paged --features metal,accelerate — 57 passed (detach round-trips intact)
• cargo fmt -p mlxcel-core --check — clean

Closes #118

[inureyes]

Security & Performance Review — Phase 1 block-pool storage (#118)

Reviewed scope: the three changed files (cache/paged.rs, cache.rs, cache/paged_pool_tests.rs). In-memory single-owner (&mut self) cache code with no network/parse/injection surface, so the review focused on memory-safety, donation correctness, integer-overflow/bounds, panic-safety, and stale-data on row recycle. Live-flow wiring deferred to #119/#120 is treated as in-design, not a defect.

Verdict: no CRITICAL or HIGH findings. Nothing auto-fixed (no fix warranted). PR stays at status:review.

Warm verification (all clean, no new compile/lint/test surface introduced):

• cargo check --lib -p mlxcel-core — clean
• cargo clippy --lib and --tests -- -D warnings — clean
• cargo fmt --check — clean
• cargo test --lib cache::paged — 56 passed (12 new pool tests + 44 existing; detach round-trips and nbytes_* budgets intact)

Findings

1. Unbounded growth / retention — LOW (acceptable, not a leak). The pool grows in 32-row chunks via grow_pool and never shrinks; next_row is monotonic. This is standard pool high-water retention, identical to the dense KVCache pre-allocate-and-keep discipline. Freed rows are recycled correctly: release_block pushes to free_rows, and assign_row pops free_rows before advancing next_row, so capacity is bounded by peak concurrent blocks (itself bounded by admission control, which now sees pool_tensor_bytes()). No UniquePtr<MlxArray> is leaked — every .take() is immediately followed by a reassign that consumes the old array into a slice_update graph node (freed at eval). Acceptable.

2. Donation correctness / no hidden O(pool) copy — CORRECT. write_block and grow_pool both follow the ADR 0001 donation pattern: let old = pool[i].take(); pool[i] = Some(slice_update(&old, ...)). After .take(), old is the sole Rust owner; the FFI passes const MlxArray& (no mlx::core::array copy), the lazy Slice primitive captures src.inner, and old drops before eval, so the array's refcount is 1 at eval and MLX donates the buffer in place (O(block) append). In grow_pool the freshly-zeros'd larger buffer is the donation target and old is the read-only update source — the one genuine data copy here is unavoidable on a real resize. Matches the proven dense-cache slice_update path. No lingering reference forces a full copy on the steady-state append.

3. Integer overflow / bounds — LOW (defense-in-depth only; no memory-safety risk). All hot-path arithmetic is sound: next_row += 1 and row + 1 need 2^64 blocks to wrap; grow_pool uses div_ceil + saturating_mul + .max(); growth trigger row >= capacity ⇒ grow(row+1) has correct 0-indexed boundary (no off-by-one). The only externally-arbitrary value not pre-bounded before an as i32 cast is write_block's slot_start (slot_start as i32 + n_slots): a pathological slot_start ≥ ~2^31 would wrap (debug: panic; release: wrap) and could bypass the > block_size check. Likewise n_blocks * block_size in gather_visible is a theoretical i32 product overflow at absurd context lengths. Neither can cause out-of-bounds memory access: there is no unsafe and no raw indexing — every slice/index goes through MLX, which bounds-checks at the C++ layer and throws rather than corrupting memory. Realistic callers (#119/#120 prefill/decode) always pass slot_start < block_size (default 32). Optional future hardening: validate slot_start as usize against block_size before the cast. Not blocking.

4. New unsafe — CLEAN. Zero unsafe blocks in the new code (confirmed by scan). The K/V path uses only the safe FFI wrappers take/slice_update/slice/reshape/transpose_axes/from_slice_i32/zeros/array_nbytes — never the unsafe gather/concatenate/stack variants.

5. Panic safety (.expect() in non-test code) — CORRECT. The five new .expect() calls (write_block pool_k/pool_v takes, gather_visible pool_meta, grow_pool pool_k/pool_v takes) are each guarded by an invariant established immediately prior in the same &mut self/&self method: pool_k/pool_v/pool_meta are always allocated together (the match pool_meta { None => set all three } arm), so "meta present ⇒ tensors present" and "tensors present ⇒ meta present" both hold. No interior mutability and no await points, so no scheduler/attacker-influenced interleaving can violate them. Acceptable.

6. Stale-data leak on row recycle — CORRECT (no leak path). A recycled row is not zeroed, but it cannot serve stale data: release_block removes the old block_id from block_rows, and gather_visible returns an error for any block with no row entry, so the window between re-acquire and re-write cannot gather. After re-write, the row holds fresh data. Trailing padding slots of a partially-rewritten recycled row are never exposed because gather_visible slices exactly [logical_start, len) (verified by partial_final_block_gathers_exactly_visible_len). The released_row_is_recycled_and_serves_fresh_data test confirms the full free→recycle→rewrite→gather path serves fresh bytes only.

Notes (informational, not findings)

• The write-before-read ordering for visible slots in [logical_start, len) is a contract for the Phase 2: Paged decode attention over real block tables #119/Phase 3: Paged prefill into the block pool #120 writer (a row whose entry exists from a prior occupant but whose current-generation slots were not rewritten would gather stale bytes). Correctly out of Phase 1: Global block-pool tensor storage #118 scope; the only Phase 1: Global block-pool tensor storage #118 caller (unit tests) always writes every visible token first.
• CachePool::memory_usage_bytes adds pool_tensor_bytes() additively and contributes 0 until a live writer is wired (Phase 3: Paged prefill into the block pool #120); it does not perturb the layout-derived reserved_bytes/used_bytes scheduling budgets. Accurate.
• PR body says "15 unit tests"; the file actually contains 12 #[test] functions (all passing). Cosmetic description mismatch only.