Skip to content

perf: presize paged pool per prefill span and eval grown slabs eagerly#229

Merged
inureyes merged 2 commits into
mainfrom
perf/issue-224-paged-pool-grow-spikes
Jun 12, 2026
Merged

perf: presize paged pool per prefill span and eval grown slabs eagerly#229
inureyes merged 2 commits into
mainfrom
perf/issue-224-paged-pool-grow-spikes

Conversation

@inureyes

@inureyes inureyes commented Jun 12, 2026

Copy link
Copy Markdown
Member

Summary

Fixes the paged-pool slab growth footprint spikes measured in #224.

  • write_prefill sizes a layer's pool once for the whole prefill span: the new presize_for_span counts span blocks that still need a physical row, nets out reusable free rows, and either allocates the pool at the right size (first write) or grows it once. Previously assign_row regrew the slab every POOL_GROW_CHUNK_BLOCKS = 32 rows and every step reallocated and copied the entire layer tensor inside one lazy graph.
  • grow_pool evaluates the grown K/V copies immediately, so each old slab returns to the MLX buffer cache before the next layer grows. A multi-layer growth episode now transiently holds one layer's old+new pair instead of every layer's pair.

Measurements (scripts/bench_memory_footprint.py, shared_prefix_burst, M1 Ultra)

case before after dense reference
llama-3.2-1b, ~3.7k-token prefix x8 2815 MiB 2225 MiB 2326 MiB (paged now 0.96x)
qwen3-0.6b, ~4.2k-token prefix x8 8752 MiB 5930 MiB 3833 MiB

The acceptance target (paged within 1.15x of dense) is met on the harness model. The residual qwen3 gap comes from old-slab ladders accumulating in the MLX buffer cache across eight successive pool-extending prefills; that regrowth disappears once cross-request prefix sharing stops re-prefilling shared content (#225, #227), so it is not chased further here. Burst wall times moved 12.6 to 12.9 s (llama) and 21.8 to 23.4 s (qwen3 burst phase), within run noise.

Validation

  • cache::paged_pool_tests 25/25 including two new tests: single-step presize for a multi-chunk span (capacity lands at the rounded target), and incremental growth past a presized capacity staying byte-identical.
  • Real-model gates, all byte-identical: paged_scheduler_parity (4), paged_prefix_share_parity (2), paged_real_model_parity (1).
  • Cold clippy clean, fmt clean. Full mlxcel-core lib suite 865 passed; the single failure (cache::turbo::pack3::tests::packed_bytes_rejects_non_multiple_of_8) is pre-existing on main in release mode (debug_assert-based should_panic) and unrelated.

Closes #224

Review follow-up

Two reviewer findings addressed before merge:

  • Test efficacy: the original tests also passed on the old incremental-growth code (final capacity converges identically). A grow_events counter now pins the behavior: 0 slab-copy growths for a presized multi-chunk span, exactly 1 for the growth-past-capacity case. Removing presize_for_span fails the suite.
  • Decode-path note: write_prefill also serves single-token decode writes (KVCache::write_paged), so the eager grow eval fires there too, once per 32 x block_size decode tokens per layer in lockstep across layers. It is a bounded, rare sync (and it bounds the transient there as well); scenario wall times stayed within run noise.

inureyes added 2 commits June 12, 2026 11:04
write_prefill now sizes a layer's pool once for the whole span
(presize_for_span nets unassigned span blocks against reusable free rows
and grows or allocates in a single step) instead of letting assign_row
regrow the slab every POOL_GROW_CHUNK_BLOCKS rows, each step copying the
entire layer tensor inside one lazy graph. grow_pool also evaluates the
grown K/V copies immediately so each old slab returns to the MLX buffer
cache before the next layer grows.

Measured (scripts/bench_memory_footprint.py, shared_prefix_burst):
llama-3.2-1b paged peak 2815 -> 2225 MiB (now 0.96x of dense, was 1.20x);
qwen3-0.6b with a ~4.2k-token prefix 8752 -> 5930 MiB. The qwen3 case
keeps a residual gap from old-slab ladders accumulating in the MLX
buffer cache across successive prefills; eliminating those regrowths
entirely is the cross-request sharing work tracked in #225/#227.

Parity: paged pool suite 25/25 plus two new presize tests; real-model
paged_scheduler_parity, paged_prefix_share_parity, and
paged_real_model_parity all byte-identical (7/7).
Reviewers found the two new tests passed on the old incremental-growth
code as well (final capacity converges identically). Add a grow_events
counter incremented on every slab-copy reallocation and assert 0 grows
for the presized multi-chunk span and exactly 1 for the
growth-past-capacity case, so removing presize_for_span now fails the
suite. Also make the presize target arithmetic saturating.
@inureyes
inureyes merged commit d655f71 into main Jun 12, 2026
5 checks passed
@inureyes
inureyes deleted the perf/issue-224-paged-pool-grow-spikes branch June 12, 2026 02:25
@inureyes inureyes added this to the 0.2 milestone Jun 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant