From 906d0d5f6a2ec78490a7b7001130bf125c28b4e9 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Sat, 18 Jul 2026 23:49:27 +0800 Subject: [PATCH] fix(prefill): cap reservation at sliding window Estimate final snapshot capacity from retained sink and window KV instead of unbounded prompt history so long auto-loop contexts remain admissible. Co-authored-by: Cursor --- docs/ops/distributed-prefill-kv-network.md | 4 +++ .../distributed/prefill_worker.py | 29 +++++++++++++++++-- .../distributed/test_prefill_worker.py | 20 +++++++++++++ 3 files changed, 50 insertions(+), 3 deletions(-) diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 76228d0..76a21eb 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/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 diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index f13ef01..80f0859 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/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) diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 3efa4e0..74be658 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/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