fix(scheduler): read the prefill chunk cap from the pool, not a snapshot - #439
Open
gberasmus87 wants to merge 1 commit into
Open
fix(scheduler): read the prefill chunk cap from the pool, not a snapshot#439gberasmus87 wants to merge 1 commit into
gberasmus87 wants to merge 1 commit into
Conversation
`CacheManager.__init__` copied `swa_pool.prefill_chunk_budget` into an instance attribute, so the cap kept its construction-time value for the life of the manager and `rebuild` never refreshed it. `Scheduler.rebuild_cache` recomputes `prefill_budget` from `cache_manager.prefill_chunk_budget` after a runtime resize, so reading a frozen value meant it wrote back the number it already had. Two consequences: * Growing the pool leaves prefill chunking where it was. Measured on DSV4-Flash through `POST /v1/cache/rebuild`: the window pool went 100 -> 215 pages and the chunk budget stayed at 4864, so an 11.7k prompt still took three whole-layer expert streams (32.5 s) instead of the one it had just been sized for. * Shrinking is worse, and is the hazard `rebuild_cache`'s own comment warns about: the stale cap is then too LARGE for the pool, so the next long prompt is chunked past what `_alloc_window` can satisfy. Making it a property that re-reads the pool on every access fixes both. Nothing assigns to it -- `scheduler.py` reads it in two places and `DSV4PagedKVCache` owns the value -- so a read-only property is the whole change. `test_rebuild_cache_refreshes_prefill_budget` already covers the scheduler half, but with a `SimpleNamespace` manager whose cap the test sets by hand, so it cannot see this: the gap is between a real manager and its real pool. The new test closes exactly that gap. It drives `DSV4PagedKVCache.rebuild` (whose `_init_paged_state` recomputes `_chunk_budget` from the new window-slot count) rather than assigning the attribute, and checks both directions -- grown and shrunk -- asserting the pool's cap actually moved first so the test cannot pass vacuously. Verified: the new test fails on the unfixed tree and passes with the fix; the rest of tests/scheduler/ is unchanged either way.
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.
fix(scheduler): read the prefill chunk cap from the pool, not a snapshot
CacheManager.__init__copiedswa_pool.prefill_chunk_budgetinto an instanceattribute, so the cap kept its construction-time value for the life of the
manager and
rebuildnever refreshed it.Scheduler.rebuild_cacherecomputesprefill_budgetfromcache_manager.prefill_chunk_budgetafter a runtime resize, so reading a frozenvalue meant it wrote back the number it already had. Two consequences:
DSV4-Flash through
POST /v1/cache/rebuild: the window pool went 100 -> 215pages and the chunk budget stayed at 4864, so an 11.7k prompt still took three
whole-layer expert streams (32.5 s) instead of the one it had just been sized
for.
rebuild_cache's own comment warnsabout: the stale cap is then too LARGE for the pool, so the next long prompt is
chunked past what
_alloc_windowcan satisfy.Making it a property that re-reads the pool on every access fixes both. Nothing
assigns to it --
scheduler.pyreads it in two places andDSV4PagedKVCacheowns the value -- so a read-only property is the whole change.
test_rebuild_cache_refreshes_prefill_budgetalready covers the scheduler half,but with a
SimpleNamespacemanager whose cap the test sets by hand, so itcannot see this: the gap is between a real manager and its real pool. The new
test closes exactly that gap. It drives
DSV4PagedKVCache.rebuild(whose_init_paged_staterecomputes_chunk_budgetfrom the new window-slot count)rather than assigning the attribute, and checks both directions -- grown and
shrunk -- asserting the pool's cap actually moved first so the test cannot pass
vacuously.
Verified: the new test fails on the unfixed tree and passes with the fix; the
rest of tests/scheduler/ is unchanged either way.