Skip to content

fix(scheduler): read the prefill chunk cap from the pool, not a snapshot - #439

Open
gberasmus87 wants to merge 1 commit into
FlashML-org:mainfrom
gberasmus87:scheduler-prefill-chunk-budget-live
Open

fix(scheduler): read the prefill chunk cap from the pool, not a snapshot#439
gberasmus87 wants to merge 1 commit into
FlashML-org:mainfrom
gberasmus87:scheduler-prefill-chunk-budget-live

Conversation

@gberasmus87

Copy link
Copy Markdown

fix(scheduler): read the prefill chunk cap from the pool, not a snapshot

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.

`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.
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.

1 participant