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..6e1cdf7 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)
@@ -246,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)