perf: presize paged pool per prefill span and eval grown slabs eagerly#229
Merged
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes the paged-pool slab growth footprint spikes measured in #224.
write_prefillsizes a layer's pool once for the whole prefill span: the newpresize_for_spancounts 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. Previouslyassign_rowregrew the slab everyPOOL_GROW_CHUNK_BLOCKS = 32rows and every step reallocated and copied the entire layer tensor inside one lazy graph.grow_poolevaluates 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)
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_tests25/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.paged_scheduler_parity(4),paged_prefix_share_parity(2),paged_real_model_parity(1).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:
grow_eventscounter now pins the behavior: 0 slab-copy growths for a presized multi-chunk span, exactly 1 for the growth-past-capacity case. Removingpresize_for_spanfails the suite.write_prefillalso serves single-token decode writes (KVCache::write_paged), so the eager grow eval fires there too, once per32 x block_sizedecode 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.