From 78b197716a36c97941afeb4f432707e19c3e0417 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Sun, 19 Jul 2026 23:15:17 +0800 Subject: [PATCH 1/2] fix(prefill): inject immutable retained-token cap Make snapshot capacity checks use the worker's configured sink and sliding window directly so missing runtime compatibility fields cannot restore full-prompt estimates. Co-authored-by: Cursor --- docs/ops/distributed-prefill-kv-network.md | 4 ++++ inference_engine/distributed/prefill_worker.py | 17 +++++++++++++++-- scripts/start_prefill_worker_node.py | 1 + .../bridge/test_prefill_worker_launchd.py | 7 +++++++ .../distributed/test_prefill_worker.py | 6 ++++++ 5 files changed, 33 insertions(+), 2 deletions(-) diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 9209fd7..e4b9f87 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -417,6 +417,10 @@ hash boundary. The default segment is 256 tokens, keeping each allens step below the five-minute research budget at the measured Prefill rate. Worker job status updates `tokens_computed` after every segment; Terminal heartbeats display tokens, percentage, and ETA. +The worker injects `sink + window` as an explicit immutable retained-token cap +into `PrefillJobStore`. Snapshot capacity checks therefore remain window-aware +even if an RPC compatibility object omits its window field; long prompts cannot +fall back to full-length estimates after deployment or process restart. Karpathy-style optimization lives in `autoresearch/prefill/`. Humans edit `program.md`, the research agent edits only `candidate.py`, and immutable diff --git a/inference_engine/distributed/prefill_worker.py b/inference_engine/distributed/prefill_worker.py index 80f0859..0a2a0e5 100644 --- a/inference_engine/distributed/prefill_worker.py +++ b/inference_engine/distributed/prefill_worker.py @@ -60,15 +60,22 @@ def estimate_final_snapshot_bytes( token_count: int, compatibility: CacheCompatibility, bytes_per_token: int, + *, + max_retained_tokens: int = 0, ) -> 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_limit = int(max_retained_tokens) + if retained_limit <= 0 and compatibility.window_size > 0: + retained_limit = ( + compatibility.sink_size + compatibility.window_size + ) + if retained_limit > 0: retained_tokens = min( retained_tokens, - compatibility.sink_size + compatibility.window_size, + retained_limit, ) return retained_tokens * int(bytes_per_token) @@ -111,6 +118,7 @@ def __init__( completed_ttl_s: float = 600.0, max_prompt_tokens: int = 131_072, estimated_snapshot_bytes_per_token: int = 16, + max_retained_tokens: int = 0, ) -> None: if min( max_concurrent_jobs, @@ -120,6 +128,8 @@ def __init__( estimated_snapshot_bytes_per_token, ) <= 0: raise ValueError("worker limits must be > 0") + if max_retained_tokens < 0: + raise ValueError("max_retained_tokens must be >= 0") if (engine is None) == (engine_factory is None): raise ValueError("provide exactly one of engine or engine_factory") self.engine = engine @@ -132,6 +142,7 @@ def __init__( self.estimated_snapshot_bytes_per_token = int( estimated_snapshot_bytes_per_token, ) + self.max_retained_tokens = int(max_retained_tokens) self._jobs: dict[str, PrefillJob] = {} self._requests: dict[tuple[str, str], str] = {} self._lock = threading.RLock() @@ -216,6 +227,7 @@ def submit( len(job.token_ids), self.cache_store.compatibility, self.estimated_snapshot_bytes_per_token, + max_retained_tokens=self.max_retained_tokens, ) if estimated_bytes > self.cache_store.max_bytes: raise ValueError( @@ -296,6 +308,7 @@ def _run(self, job_id: str) -> None: len(job.token_ids), self.cache_store.compatibility, self.estimated_snapshot_bytes_per_token, + max_retained_tokens=self.max_retained_tokens, ), ) engine = self._engine_for_current_thread() diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py index 49653d0..64da9cd 100644 --- a/scripts/start_prefill_worker_node.py +++ b/scripts/start_prefill_worker_node.py @@ -143,6 +143,7 @@ def engine_factory() -> MLXPrefillComputeEngine: estimated_snapshot_bytes_per_token=( args.estimated_snapshot_bytes_per_token ), + max_retained_tokens=args.sink + args.window, ) jobs.warmup() diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index 2445056..0e2b7f8 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -85,3 +85,10 @@ def test_two_mac_deployment_uses_allens_as_prefill_only(): assert "--window2048" in worker assert "--prefill-tps1" in worker assert "scripts/start_prefill_cache_node.py" in PEER_PLIST.read_text() + + +def test_worker_injects_explicit_retained_token_cap(): + source = ( + ROOT / "scripts" / "start_prefill_worker_node.py" + ).read_text() + assert "max_retained_tokens=args.sink + args.window" in source diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 74be658..092bba8 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/tests/inference_engine/distributed/test_prefill_worker.py @@ -55,6 +55,12 @@ def test_snapshot_estimate_caps_at_sink_plus_sliding_window(): ) no_window = CacheCompatibility(model_id="m", window_size=0) assert estimate_final_snapshot_bytes(2731, no_window, 10) == 27_310 + assert estimate_final_snapshot_bytes( + 2731, + no_window, + 400_000, + max_retained_tokens=2052, + ) == 820_800_000 with pytest.raises(ValueError, match="must be > 0"): estimate_final_snapshot_bytes(0, compatibility, 400_000) From 623580f124aded1118547f124f9e017b4939ac71 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Sun, 19 Jul 2026 23:19:13 +0800 Subject: [PATCH 2/2] test(prefill): cover invalid retained-token caps Exercise constructor validation so the explicit window cap remains fully covered by the strict CI gate. Co-authored-by: Cursor --- tests/inference_engine/distributed/test_prefill_worker.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/inference_engine/distributed/test_prefill_worker.py b/tests/inference_engine/distributed/test_prefill_worker.py index 092bba8..6e1cdf7 100644 --- a/tests/inference_engine/distributed/test_prefill_worker.py +++ b/tests/inference_engine/distributed/test_prefill_worker.py @@ -252,6 +252,8 @@ def test_job_store_validation_queue_stats_and_gc(): PrefillJobStore(None, cache) with pytest.raises(ValueError, match="exactly one"): PrefillJobStore(_Engine(), cache, engine_factory=_Engine) + with pytest.raises(ValueError, match="max_retained_tokens"): + PrefillJobStore(_Engine(), cache, max_retained_tokens=-1) blocking = _Engine() blocking.block.set() jobs = PrefillJobStore(blocking, cache, max_jobs=1, max_prompt_tokens=4)