perf(core): paged pool slab growth reallocates whole layer tensors and transiently doubles server footprint

[inureyes]

Measurement context

Found during a cross-version server memory-footprint measurement (v0.1.4 vs main, 2026-06-12) using scripts/bench_memory_footprint.py (commit 18c88c1). Method: one fresh mlxcel serve process per (config, scenario), fixed HTTP workloads at temperature 0, sampling the process phys-footprint via /usr/bin/footprint at 2.5 s (RSS does not see MLX Metal buffers at all). Model: llama-3.2-1b-4bit on M1 Ultra, --max-batch-size 8. Scenarios: idle, shared_prefix_burst (8 concurrent requests sharing a ~3.7k-token system prompt), seqshare (same prefix, 8 sequential requests), multiturn (one conversation, 8 turns), churn (32 distinct prompts). Run-to-run noise: idle/multiturn within 0.5 percent, burst/seqshare/churn peaks within 10 to 25 percent.

Problem

PagedBlockPool::grow_pool (src/lib/mlxcel-core/src/cache/paged.rs:1587) grows a layer's slab by allocating a fresh [capacity, block_size, n_kv_heads, head_dim] tensor with ffi::zeros and copying the old contents via slice_update, every POOL_GROW_CHUNK_BLOCKS = 32 blocks. The grow ops are recorded into the lazy graph, so during a long prefill many layers' old and new slabs coexist until the next eval. The transient cost approaches 2x the entire pool.

Measured impact

• shared_prefix_burst peak: default paged backend 2815 MiB vs forced dense 2356 MiB (llama-3.2-1b).
• qwen3-0.6b with a ~4.2k-token prefix, 8 concurrent: paged 8752 MiB vs dense 3833 MiB (2.3x). The number matches pool (~3.5 GiB) x 2 plus weights.
• v0.1.4 baseline on the identical workload: 1653 MiB peak. The default config regressed about +70 percent.

Candidate fixes (decide during implementation)

1. Pre-size on prefill admission: write_prefill knows the full token count up front, so grow once to the exact target instead of 32-block steps.
2. Exponential growth (for example double-or-needed) to cut realloc frequency.
3. Eager eval/free discipline per layer grow so old slabs do not pile up inside one lazy graph.
4. Budget-driven pre-reservation when --kv-cache-budget is set.

Acceptance criteria

• Byte-identical decode parity vs current main (existing paged parity tests stay green).
• shared_prefix_burst fp_peak for the default paged backend within about 1.15x of the dense backend on scripts/bench_memory_footprint.py.
• Before/after bench numbers recorded in the PR.
• The fix is live in the default serve path, not behind a flag.

Related

Filed together from the same measurement: #225, #226, #227, #228.