diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 44adde42f..f24e4ba75 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -50,11 +50,9 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ # lifecycle (alloc_swa / out-of-window free / free-on-finish). is_swa gates only the extra # SWARadixCache reuse machinery (tree match/insert/evict_swa/swa_uuid lock). self.swa_paged = swa_pool is not None and getattr(swa_pool, "swa_paged", False) - # Owned-pool capability pickup: a plugged-in swa pool may cap the prefill chunk (DSV4: - # ~half the window working set). Instance attrs shadow the class defaults; absent - # attributes leave the defaults untouched (Gemma4). - if swa_pool is not None: - self.prefill_chunk_budget = getattr(swa_pool, "prefill_chunk_budget", None) + # Owned-pool capability pickup is a PROPERTY (below), not a snapshot taken here. + # This used to read `getattr(swa_pool, "prefill_chunk_budget", None)` into an + # instance attribute, which froze the cap at its construction-time value. self.prefix_cache = self._make_prefix_cache(device, page_size, type) self.device = device self.num_pages = num_pages @@ -64,7 +62,34 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ # ----- capability hooks (defaults; plugged-in pools may narrow them) ----- supports_runtime_rebuild = True - prefill_chunk_budget = None # generic shared page pool: no per-model prefill chunk cap + + @property + def prefill_chunk_budget(self) -> int | None: + """Live prefill-chunk cap from the owned window pool, re-read on every access. + + None for a generic shared page pool (no per-model cap) and for a pool that + does not publish one (Gemma4). + + This MUST NOT be snapshotted. It used to be an int copied in ``__init__``, + so it kept its construction-time value for the life of the manager, and + ``rebuild`` never refreshed it. Two consequences, both real: + + * ``Scheduler.rebuild_cache`` recomputes ``prefill_budget`` from this after a + runtime resize. Reading a frozen value meant it wrote back the number it + already had, so resizing the DSV4 window pool through + ``POST /v1/cache/rebuild`` moved the pool but left prefill chunking exactly + where it was. Measured on DSV4-Flash: the pool went 100 -> 215 window pages + and the chunk budget stayed at 4864, so an 11.7k prompt still took three + whole-layer expert streams (32.5s) instead of the one it was sized for. + * Worse in the shrink direction, and the hazard ``rebuild_cache``'s own comment + warns about: after shrinking the window pool the stale cap is too LARGE, so + the next long prompt is chunked past what the pool can hold and crashes + ``_alloc_window``. + """ + pool = self.swa_pool + if pool is None: + return None + return getattr(pool, "prefill_chunk_budget", None) @property def prefill_chunk_align(self) -> int: diff --git a/tests/scheduler/test_dsv4_generic_manager.py b/tests/scheduler/test_dsv4_generic_manager.py index ad69b89ec..25fd241ed 100644 --- a/tests/scheduler/test_dsv4_generic_manager.py +++ b/tests/scheduler/test_dsv4_generic_manager.py @@ -151,6 +151,38 @@ def test_capability_surface_matches_scheduler_expectations(): assert pool.sliding_window_size == P +def test_prefill_chunk_budget_follows_the_pool_across_a_rebuild(): + """The manager must report the pool's LIVE chunk cap, not the one it had at construction. + + ``Scheduler.rebuild_cache`` recomputes ``prefill_budget`` from + ``cache_manager.prefill_chunk_budget`` after a runtime resize, so a value snapshotted in + ``CacheManager.__init__`` makes the rebuild write back the number it already had. Growing + the pool then leaves prompts chunked against the old, smaller cap; shrinking it leaves a cap + too LARGE for the pool, and the next long prompt is chunked past what ``_alloc_window`` can + satisfy. + + ``test_rebuild_cache_refreshes_prefill_budget`` covers the scheduler half 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. + """ + cm, pool, _ = _stack(num_pages=32) + before = cm.prefill_chunk_budget + assert before == pool.prefill_chunk_budget > 0 + + # What POST /v1/cache/rebuild does: resize the window pool in place. _init_paged_state + # recomputes the pool's cap from the new window-slot count. + pool.rebuild(dsv4_pool_sizes(num_pages=96 + 1, args=_args(), swa_ratio=1.0, P=P)) + grown = pool.prefill_chunk_budget + assert grown > before, "inert test: the rebuild did not move the pool's cap" + assert cm.prefill_chunk_budget == grown + + # Shrink is the dangerous direction: a stale cap is then too large for the pool. + pool.rebuild(dsv4_pool_sizes(num_pages=24 + 1, args=_args(), swa_ratio=1.0, P=P)) + shrunk = pool.prefill_chunk_budget + assert shrunk < before, "inert test: the shrink did not move the pool's cap" + assert cm.prefill_chunk_budget == shrunk + + def test_chunk_boundaries_stay_page_aligned_under_unaligned_budget(): """Chunk continuations resume the compressor carry, so every minted chunk must END page-aligned -- even when the binding cap is an unaligned token-budget leftover. A leftover