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

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:

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