fix(prefill): atomically reserve and publish snapshots

deploy/install_prefill_worker_launchd.sh

@@ -12,8 +12,9 @@ set -euo pipefail
 BIND="${KAKEYA_WORKER_BIND:-0.0.0.0:53051}"
 TENANT="${KAKEYA_TENANT_ID:-default}"
 CACHE_GB="${KAKEYA_WORKER_CACHE_GB:-4}"
-CACHE_MIN_GB="${KAKEYA_WORKER_CACHE_MIN_GB:-0.25}"
-MEMORY_RESERVE_GB="${KAKEYA_WORKER_MEMORY_RESERVE_GB:-2}"
+CACHE_MIN_GB="${KAKEYA_WORKER_CACHE_MIN_GB:-1}"
+MEMORY_RESERVE_GB="${KAKEYA_WORKER_MEMORY_RESERVE_GB:-0.5}"
+SNAPSHOT_BYTES_PER_TOKEN="${KAKEYA_SNAPSHOT_BYTES_PER_TOKEN:-400000}"
 ADAPTIVE_CACHE="${KAKEYA_WORKER_ADAPTIVE_CACHE:-0}"
 PSK_FILE="${KAKEYA_FLEET_PSK_FILE:-}"
 CACHE_MODEL_ID="${KAKEYA_CACHE_MODEL_ID:-$KAKEYA_WORKER_MODEL}"
@@ -76,6 +77,7 @@ cat > "$PLIST" <<EOF
     <string>--cache-gb</string><string>$CACHE_GB</string>
     <string>--cache-min-gb</string><string>$CACHE_MIN_GB</string>
     <string>--memory-reserve-gb</string><string>$MEMORY_RESERVE_GB</string>
+    <string>--estimated-snapshot-bytes-per-token</string><string>$SNAPSHOT_BYTES_PER_TOKEN</string>
     $adaptive_xml
     <string>--sink</string><string>$SINK</string>
     <string>--window</string><string>$WINDOW</string>

deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist

@@ -9,7 +9,7 @@
     <string>/Users/fluffy314/Documents/Kakeya-LLM-Inference-engine-prefill-kv-network/scripts/start_grpc_runtime_server.py</string>
     <string>--backend</string><string>mlx</string>
     <string>--verifier-id</string><string>/Users/fluffy314/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit</string>
-    <string>--bind</string><string>127.0.0.1:51051</string>
+    <string>--bind</string><string>0.0.0.0:51051</string>
     <string>--capacity</string><string>1</string>
     <string>--sink</string><string>4</string>
     <string>--window</string><string>2048</string>

deploy/launchd/ai.kakeya.prefill-worker-peer.plist

@@ -23,9 +23,10 @@
     <string>--window</string><string>2048</string>
     <string>--block-size-tokens</string><string>64</string>
     <string>--cache-gb</string><string>8</string>
-    <string>--cache-min-gb</string><string>0.25</string>
+    <string>--cache-min-gb</string><string>1</string>
     <string>--adaptive-cache</string>
-    <string>--memory-reserve-gb</string><string>2</string>
+    <string>--memory-reserve-gb</string><string>0.5</string>
+    <string>--estimated-snapshot-bytes-per-token</string><string>400000</string>
     <string>--prefill-tps</string><string>1</string>
     <string>--max-concurrent-jobs</string><string>1</string>
     <string>--network</string><string>thunderbolt</string>

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

@@ -349,6 +349,11 @@ and semantic fallback are forbidden. A global Critic score is valid only when
 `critic_omitted_tokens=0`. Long Prefill operations emit a heartbeat every 30
 seconds; on the 16GB allens worker, full-context Critic Prefill may take 15–25
 minutes.
+The worker reserves estimated final-snapshot capacity before model compute,
+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.
 Interactive prompt templates are deterministic and contain no per-run nonce, so
 repeating the same task can reuse allens cold-tier and Primary hot-tier KV.
 

inference_engine/distributed/prefill_cache.py

@@ -159,15 +159,16 @@ def __init__(
         self._evictions = 0
         self._bytes_evicted = 0
         self._put_failures = 0
+        self._reservations: dict[str, int] = {}
         self._lock = threading.RLock()
 
     def put(self, block: CacheBlock) -> bool:
         """Publish one immutable block. Returns False for an identical hit."""
-        if block.nbytes > self.max_bytes:
-            with self._lock:
-                self._put_failures += 1
-            raise ValueError("block payload exceeds cache capacity")
         with self._lock:
+            available = self.max_bytes - sum(self._reservations.values())
+            if block.nbytes > available:
+                self._put_failures += 1
+                raise ValueError("block payload exceeds cache capacity")
             self._expire_leases(time.time())
             existing = self._blocks.get(block.block_hash)
             if existing is not None:
@@ -179,7 +180,7 @@ def put(self, block: CacheBlock) -> bool:
             self._blocks[block.block_hash] = block
             self._bytes_used += block.nbytes
             self._epoch += 1
-            self._evict_to_budget()
+            self._evict_to_limit(available)
             if block.block_hash not in self._blocks:
                 self._put_failures += 1
                 raise ValueError(
@@ -263,6 +264,80 @@ def fetch(self, lease_id: str, *, now: float | None = None) -> tuple[CacheBlock,
             self._bytes_served += lease.transfer_bytes
             return tuple(blocks)
 
+    def reserve(self, reservation_id: str, byte_count: int) -> None:
+        """Reserve cache capacity before an expensive Prefill job starts."""
+        if not reservation_id or byte_count <= 0:
+            raise ValueError("reservation id and byte count must be positive")
+        with self._lock:
+            if reservation_id in self._reservations:
+                raise ValueError("duplicate cache reservation")
+            requested = int(byte_count)
+            reserved = sum(self._reservations.values()) + requested
+            if reserved > self.max_bytes:
+                raise ValueError(
+                    f"snapshot reservation {requested} exceeds available "
+                    f"cache budget {self.max_bytes - sum(self._reservations.values())}",
+                )
+            self._expire_leases(time.time())
+            self._evict_to_limit(self.max_bytes - reserved)
+            if self._bytes_used > self.max_bytes - reserved:
+                raise ValueError("cache capacity is pinned by active leases")
+            self._reservations[reservation_id] = requested
+
+    def release_reservation(self, reservation_id: str) -> None:
+        with self._lock:
+            self._reservations.pop(reservation_id, None)
+
+    def publish_and_lease(
+        self,
+        blocks: Sequence[CacheBlock],
+        block_hashes: Sequence[bytes],
+        *,
+        reservation_id: str,
+        lease_seconds: float = DEFAULT_LEASE_SECONDS,
+    ) -> PrefixLease:
+        """Atomically publish and pin the final computed snapshot."""
+        if not blocks or len(blocks) != len(block_hashes):
+            raise ValueError("one computed snapshot is required per block hash")
+        if lease_seconds <= 0:
+            raise ValueError("lease_seconds must be > 0")
+        final = blocks[-1]
+        if final.block_hash != bytes(block_hashes[-1]):
+            raise ValueError("final snapshot hash does not match request")
+        now = time.time()
+        with self._lock:
+            reserved = self._reservations.get(reservation_id)
+            if reserved is None:
+                raise ValueError("unknown cache reservation")
+            if final.nbytes > reserved:
+                raise ValueError(
+                    f"final snapshot {final.nbytes} exceeds reservation {reserved}",
+                )
+            self._expire_leases(now)
+            self._evict_to_limit(self.max_bytes - final.nbytes)
+            if self._bytes_used > self.max_bytes - final.nbytes:
+                raise ValueError("cache capacity is pinned by active leases")
+            self._put_locked(final)
+            lease_id = secrets.token_urlsafe(18)
+            lease = PrefixLease(
+                lease_id=lease_id,
+                block_hashes=(final.block_hash,),
+                hit_block_count=len(block_hashes),
+                hit_token_count=final.token_count,
+                transfer_bytes=final.nbytes,
+                cache_epoch=self._epoch,
+                expires_at_unix=now + lease_seconds,
+                payload_sha256=final.payload_sha256,
+            )
+            self._leases[lease_id] = lease
+            del self._reservations[reservation_id]
+            # Preserve longest useful boundaries when spare capacity remains.
+            for block in reversed(blocks[:-1]):
+                if block.nbytes + self._bytes_used > self.max_bytes:
+                    continue
+                self._put_locked(block)
+            return lease
+
     def stats(self) -> CacheStats:
         with self._lock:
             return CacheStats(
@@ -285,9 +360,12 @@ def resize(self, max_bytes: int) -> bool:
             raise ValueError("max_bytes must be > 0")
         with self._lock:
             previous = self.max_bytes
+            reserved = sum(self._reservations.values())
+            if int(max_bytes) < reserved:
+                return False
             self.max_bytes = int(max_bytes)
-            self._evict_to_budget()
-            if self._bytes_used > self.max_bytes:
+            self._evict_to_limit(self.max_bytes - reserved)
+            if self._bytes_used > self.max_bytes - reserved:
                 self.max_bytes = max(previous, self._bytes_used)
                 return False
             return True
@@ -322,8 +400,11 @@ def _pinned_hashes(self) -> set[bytes]:
         }
 
     def _evict_to_budget(self) -> None:
+        self._evict_to_limit(self.max_bytes)
+
+    def _evict_to_limit(self, limit: int) -> None:
         pinned = self._pinned_hashes()
-        while self._bytes_used > self.max_bytes and self._blocks:
+        while self._bytes_used > limit and self._blocks:
             victim = next((h for h in self._blocks if h not in pinned), None)
             if victim is None:
                 break
@@ -333,6 +414,19 @@ def _evict_to_budget(self) -> None:
             self._bytes_evicted += block.nbytes
             self._epoch += 1
 
+    def _put_locked(self, block: CacheBlock) -> bool:
+        existing = self._blocks.get(block.block_hash)
+        if existing is not None:
+            if existing.payload_sha256 != block.payload_sha256:
+                self._put_failures += 1
+                raise ValueError("content-address collision with different payload")
+            self._blocks.move_to_end(block.block_hash)
+            return False
+        self._blocks[block.block_hash] = block
+        self._bytes_used += block.nbytes
+        self._epoch += 1
+        return True
+
 
 def total_payload_bytes(blocks: Iterable[CacheBlock]) -> int:
     return sum(block.nbytes for block in blocks)

inference_engine/distributed/prefill_worker.py

@@ -93,12 +93,14 @@ def __init__(
         max_jobs: int = 128,
         completed_ttl_s: float = 600.0,
         max_prompt_tokens: int = 131_072,
+        estimated_snapshot_bytes_per_token: int = 16,
     ) -> None:
         if min(
             max_concurrent_jobs,
             max_jobs,
             completed_ttl_s,
             max_prompt_tokens,
+            estimated_snapshot_bytes_per_token,
         ) <= 0:
             raise ValueError("worker limits must be > 0")
         if (engine is None) == (engine_factory is None):
@@ -110,6 +112,9 @@ def __init__(
         self.max_jobs = int(max_jobs)
         self.completed_ttl_s = float(completed_ttl_s)
         self.max_prompt_tokens = int(max_prompt_tokens)
+        self.estimated_snapshot_bytes_per_token = int(
+            estimated_snapshot_bytes_per_token,
+        )
         self._jobs: dict[str, PrefillJob] = {}
         self._requests: dict[tuple[str, str], str] = {}
         self._lock = threading.RLock()
@@ -190,6 +195,14 @@ def submit(
                     if deadline_ms > 0 else 0.0
                 ),
             )
+            estimated_bytes = (
+                len(job.token_ids) * self.estimated_snapshot_bytes_per_token
+            )
+            if estimated_bytes > self.cache_store.max_bytes:
+                raise ValueError(
+                    f"estimated final snapshot {estimated_bytes} exceeds "
+                    f"cache capacity {self.cache_store.max_bytes}",
+                )
             self._jobs[job.job_id] = job
             self._requests[request_key] = job.job_id
             job.future = self._executor.submit(self._run, job.job_id)
@@ -216,6 +229,7 @@ def cancel(self, job_id: str, tenant_id: str) -> bool:
             if job.future is not None and job.future.cancel():
                 job.state = PrefillJobState.CANCELLED
                 job.finished_at = time.time()
+                self.cache_store.release_reservation(job.job_id)
             return True
 
     def stats(self) -> tuple[int, int, float, int]:
@@ -239,6 +253,7 @@ def _run(self, job_id: str) -> None:
             if job.cancelled.is_set():
                 job.state = PrefillJobState.CANCELLED
                 job.finished_at = time.time()
+                self.cache_store.release_reservation(job.job_id)
                 return
             job.state = PrefillJobState.RUNNING
         started = time.perf_counter()
@@ -252,6 +267,13 @@ def _run(self, job_id: str) -> None:
                 timer.daemon = True
                 timer.start()
         try:
+            # MLX workers are single-job. Reserve the full current budget
+            # before model compute so adaptive resizing and unrelated
+            # boundaries cannot evict the final snapshot before leasing.
+            self.cache_store.reserve(
+                job.job_id,
+                self.cache_store.max_bytes,
+            )
             blocks = tuple(self._engine_for_current_thread().compute_prefill(
                 job.token_ids,
                 job.block_hashes,
@@ -264,13 +286,11 @@ def _run(self, job_id: str) -> None:
                 raise RuntimeError(
                     "prefill engine must return one snapshot per block hash",
                 )
-            for block in blocks:
-                if job.cancelled.is_set():
-                    raise InterruptedError("prefill job cancelled")
-                self.cache_store.put(block)
-            lease = self.cache_store.lookup(job.block_hashes)
-            if not lease.lease_id:
-                raise RuntimeError("computed snapshot was not discoverable")
+            lease = self.cache_store.publish_and_lease(
+                blocks,
+                job.block_hashes,
+                reservation_id=job.job_id,
+            )
             if job.cancelled.is_set():
                 raise InterruptedError("prefill job cancelled")
             with self._lock:
@@ -292,6 +312,7 @@ def _run(self, job_id: str) -> None:
                 job.state = PrefillJobState.FAILED
                 job.failure_reason = f"{type(exc).__name__}: {exc}"
         finally:
+            self.cache_store.release_reservation(job.job_id)
             if timer is not None:
                 timer.cancel()
             with self._lock:

scripts/start_prefill_worker_node.py

@@ -136,6 +136,9 @@ def engine_factory() -> MLXPrefillComputeEngine:
         max_jobs=args.max_jobs,
         completed_ttl_s=args.job_ttl_s,
         max_prompt_tokens=args.max_prompt_tokens,
+        estimated_snapshot_bytes_per_token=(
+            args.estimated_snapshot_bytes_per_token
+        ),
     )
     jobs.warmup()
 
@@ -290,6 +293,11 @@ def main() -> None:
     parser.add_argument("--max-concurrent-jobs", type=int, default=1)
     parser.add_argument("--max-jobs", type=int, default=128)
     parser.add_argument("--max-prompt-tokens", type=int, default=131072)
+    parser.add_argument(
+        "--estimated-snapshot-bytes-per-token",
+        type=int,
+        default=400_000,
+    )
     parser.add_argument("--job-ttl-s", type=float, default=600.0)
     parser.add_argument("--prefill-tps", type=float, default=20.0)
     parser.add_argument("--max-concurrent-rpcs", type=int, default=32)

tests/inference_engine/bridge/test_prefill_worker_launchd.py

@@ -24,6 +24,7 @@ def test_worker_installer_emits_full_cache_compatibility_contract():
         "--max-prompt-tokens",
         "--cache-min-gb",
         "--memory-reserve-gb",
+        "--estimated-snapshot-bytes-per-token",
     ):
         assert f"<string>{flag}</string>" in source
     assert 'PEER="${KAKEYA_WORKER_PEER:-}"' in source
@@ -48,6 +49,7 @@ def test_two_mac_deployment_uses_allens_as_prefill_only():
         "<string>--prefill-policy</string><string>remote-required</string>"
         in plist
     )
+    assert "<string>--bind</string><string>0.0.0.0:51051</string>" in plist
     assert (
         "<string>--prefill-worker-timeout-s</string><string>3600</string>"
         in plist
@@ -66,7 +68,13 @@ def test_two_mac_deployment_uses_allens_as_prefill_only():
     assert "scripts/start_prefill_worker_node.py" in worker
     assert "scripts/start_prefill_cache_node.py" not in worker
     assert "<string>--cache-gb</string><string>8</string>" in worker
-    assert "<string>--cache-min-gb</string><string>0.25</string>" in worker
+    assert "<string>--cache-min-gb</string><string>1</string>" in worker
+    assert "<string>--memory-reserve-gb</string><string>0.5</string>" in worker
+    assert (
+        "<string>--estimated-snapshot-bytes-per-token</string>"
+        "<string>400000</string>"
+        in worker
+    )
     assert "<string>--adaptive-cache</string>" in worker
     assert "<string>--window</string><string>2048</string>" in worker
     assert "<string>--prefill-tps</string><string>1</string>" in worker