perf(core): chunked slab storage for the paged pool to eliminate growth reallocation entirely #235

Problem

A paged pool layer's physical storage is one tensor [capacity_blocks, block_size, n_kv_heads, head_dim] per K and V. Growth past capacity reallocates the whole slab: grow_pool (src/lib/mlxcel-core/src/cache/paged.rs) allocates a larger tensor with ffi::zeros and copies the old contents via slice_update. #224 reduced the damage (one presize per prefill span plus eager eval bounding the transient to one layer's old+new pair), but two structural costs remain whenever the pool grows across successive non-shared long prefills:

1. Every growth still copies the entire layer slab (O(pool) memcpy per grow per layer).
2. The freed old slabs form a ladder of strictly increasing sizes in the MLX buffer cache that no later allocation can reuse (the next grow always needs a bigger buffer), so the cache accumulates roughly the sum of the ladder. Measured before the sharing fixes: qwen3-0.6b, 8 sequential pool-extending ~4.2k-token prefills peaked at 5930 MiB vs 3833 MiB on the dense backend (1.55x), with the gap attributable to ladder retention.

Cross-request sharing (#227) removed pool growth for shared-prefix traffic, so the regression no longer shows on the shared scenarios. It still reproduces with DISTINCT long prompts (each prefill extends the pool), e.g. 8 concurrent unique ~4k-token system prompts.

Proposed solution: chunked slabs

Replace the single per-layer pool tensor with a vector of fixed-size slab tensors:

• pool_k[layer], pool_v[layer] become Vec<UniquePtr<MlxArray>>, each slab shaped [POOL_SLAB_BLOCKS, block_size, n_kv_heads, head_dim]. A physical row maps as slab = row / POOL_SLAB_BLOCKS, row_in_slab = row % POOL_SLAB_BLOCKS.
• Growth appends a fresh zeros slab. No copy, no free, no transient old+new pair, no ladder: allocator high-water becomes exactly the live capacity. presize_for_span appends however many slabs the span needs in one step; the perf(core): paged pool slab growth reallocates whole layer tensors and transiently doubles server footprint #224 pool_grow_events counter becomes a slab-append counter (no slab is ever copied).
• POOL_SLAB_BLOCKS starts at the current POOL_GROW_CHUNK_BLOCKS = 32 (1024 tokens per layer-slab at block size 32, about 8 MiB per layer for qwen3-0.6b), keeping allocation granularity identical to today. Tunable later.

All direct pool-tensor access already lives in paged.rs (the fields are private), so the change is localized. Touch points and their adaptations:

1. write_block: slice_update targets the row's slab instead of the single tensor. Same O(block) donation discipline.
2. copy_on_write_block: source slice comes from the source row's slab.
3. gather_visible (the live pooled read): one ffi::take cannot span tensors. Split the block-table row list into maximal same-slab runs in logical order, take each run from its slab, and concatenate the per-run results in order before the existing reshape/slice/transpose. Rows for one sequence are allocated mostly monotonically, so a 4k-token sequence (128 blocks) spans at most about 4 runs; the concat output copy is the same magnitude as the gather copy the current code already materializes. Worst-case fragmentation (freed-row reuse interleaving sequences within slabs) degrades to one take per run, which the benchmark gate below bounds.
4. paged_decode_fused (Phase 6: Fused Metal paged-attention kernel #123, gated off by default, env MLXCEL_PAGED_ATTENTION_NATIVE): the fused kernel reads a single pool buffer. Either teach it per-slab base pointers or have it decline to the gather fallback when a layer has more than one slab; declining is acceptable for the first PR since the kernel is gated off and documented as not beating gather at long context.
5. pool_tensor_bytes, stats bytes_reserved: sum over slabs.
6. The Distributed block-table serialization (kv_cache_serde) #125 serde block-content extraction and the restore path: per-block reads/writes localize to the row's slab via the same mapping helper.
7. Introduce one private helper (slab_for_row(layer, row) -> (&MlxArray, i32) plus a mutable variant) so every touch point shares the mapping.

Behavioral invariants: block ids, refcounts, free-row reuse, budget accounting, detach/adopt/clone, and real-bytes accounting (#226) are all row-index-level concepts and stay unchanged.

Acceptance criteria

• All paged parity suites stay byte-identical: paged_pool unit tests, paged_detach tests, real-model paged_scheduler_parity, paged_prefix_share_parity (7 cases), paged_real_model_parity, serde handoff parity.
• New unit tests: a span crossing multiple slabs round-trips byte-identically against the dense reference; growth appends slabs without copying (counter-pinned); gather over an interleaved (fragmented) block table is byte-identical.
• Footprint gate: 8 concurrent DISTINCT ~4k-token prompts (qwen3-0.6b, scripts/bench_memory_footprint.py with per-request unique system prompts) peak within about 1.15x of the dense backend on the same workload.
• Decode throughput gate: examples/page_gather_microbench.rs and the bench harness scenario wall times within run noise of main.
• The fused-kernel gate behavior (decline or multi-slab support) is documented in turbo-kv-cache.md and ADR 0001 if the decline route is taken.

Related

Follow-up from the #224-#228 memory-footprint queue (see #224 for the partial mitigation this replaces / see the queue for context). Filed together with #234.