fix(prefill): cap reservation at sliding window

docs/ops/distributed-prefill-kv-network.md

@@ -390,6 +390,10 @@ prevents adaptive shrink from consuming active reservations, then atomically
 publishes and leases the final snapshot before adding optional intermediate
 boundaries. The 16GB allens deployment uses a 1 GiB cache floor and 0.5 GiB
 memory reserve. Capacity failures are rejected before Prefill starts.
+Snapshot capacity estimation uses retained KV tokens, not unbounded history:
+`min(prompt_tokens, sink + window) * estimated_bytes_per_token`. Long
+auto-loop histories therefore stop increasing the reservation once the MLX
+sliding window is full.
 The MLX worker exports and compresses only the final chained-prefix snapshot.
 The final hash commits every preceding token block, so exporting a growing full
 snapshot at every 64-token boundary is redundant and creates quadratic

inference_engine/distributed/prefill_worker.py

@@ -56,6 +56,23 @@ def compute_prefill(
     ) -> Sequence[CacheBlock]: ...
 
 
+def estimate_final_snapshot_bytes(
+    token_count: int,
+    compatibility: CacheCompatibility,
+    bytes_per_token: int,
+) -> int:
+    """Estimate retained KV, capped by sink + sliding-window capacity."""
+    if token_count <= 0 or bytes_per_token <= 0:
+        raise ValueError("token_count and bytes_per_token must be > 0")
+    retained_tokens = int(token_count)
+    if compatibility.window_size > 0:
+        retained_tokens = min(
+            retained_tokens,
+            compatibility.sink_size + compatibility.window_size,
+        )
+    return retained_tokens * int(bytes_per_token)
+
+
 @dataclass
 class PrefillJob:
     job_id: str
@@ -195,8 +212,10 @@ def submit(
                     if deadline_ms > 0 else 0.0
                 ),
             )
-            estimated_bytes = (
-                len(job.token_ids) * self.estimated_snapshot_bytes_per_token
+            estimated_bytes = estimate_final_snapshot_bytes(
+                len(job.token_ids),
+                self.cache_store.compatibility,
+                self.estimated_snapshot_bytes_per_token,
             )
             if estimated_bytes > self.cache_store.max_bytes:
                 raise ValueError(
@@ -273,7 +292,11 @@ def _run(self, job_id: str) -> None:
             # content-addressed snapshots needed by later restore requests.
             self.cache_store.reserve(
                 job.job_id,
-                len(job.token_ids) * self.estimated_snapshot_bytes_per_token,
+                estimate_final_snapshot_bytes(
+                    len(job.token_ids),
+                    self.cache_store.compatibility,
+                    self.estimated_snapshot_bytes_per_token,
+                ),
             )
             engine = self._engine_for_current_thread()
             set_progress = getattr(engine, "set_progress_callback", None)

tests/inference_engine/distributed/test_prefill_worker.py

@@ -23,6 +23,7 @@
     PrefillJobState,
     PrefillJobStore,
     add_prefill_worker_service,
+    estimate_final_snapshot_bytes,
     get_prefill_job_sync,
     submit_prefill_job_sync,
 )
@@ -39,6 +40,25 @@
 AUTH = FleetAuthConfig(b"k" * 32, "tenant", "client")
 
 
+def test_snapshot_estimate_caps_at_sink_plus_sliding_window():
+    compatibility = CacheCompatibility(
+        model_id="m",
+        block_size_tokens=64,
+        sink_size=4,
+        window_size=2048,
+    )
+    assert estimate_final_snapshot_bytes(2731, compatibility, 400_000) == (
+        820_800_000
+    )
+    assert estimate_final_snapshot_bytes(1000, compatibility, 400_000) == (
+        400_000_000
+    )
+    no_window = CacheCompatibility(model_id="m", window_size=0)
+    assert estimate_final_snapshot_bytes(2731, no_window, 10) == 27_310
+    with pytest.raises(ValueError, match="must be > 0"):
+        estimate_final_snapshot_bytes(0, compatibility, 400_000)
+
+
 class _Engine:
     def __init__(self):
         self.calls = 0