feat(prefill): add tiered KV hot promotion and offload

deploy/install_prefill_worker_launchd.sh

@@ -12,6 +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}"
+ADAPTIVE_CACHE="${KAKEYA_WORKER_ADAPTIVE_CACHE:-0}"
 PSK_FILE="${KAKEYA_FLEET_PSK_FILE:-}"
 CACHE_MODEL_ID="${KAKEYA_CACHE_MODEL_ID:-$KAKEYA_WORKER_MODEL}"
 MODEL_REVISION="${KAKEYA_MODEL_REVISION:-}"
@@ -43,6 +46,10 @@ peer_xml=""
 if [[ -n "$PEER" ]]; then
   peer_xml="<string>--peer</string><string>$PEER</string>"
 fi
+adaptive_xml=""
+if [[ "$ADAPTIVE_CACHE" == "1" ]]; then
+  adaptive_xml="<string>--adaptive-cache</string>"
+fi
 
 cat > "$PLIST" <<EOF
 <?xml version="1.0" encoding="UTF-8"?>
@@ -67,6 +74,9 @@ cat > "$PLIST" <<EOF
     <string>--layer-geometry-hash</string><string>$KAKEYA_LAYER_GEOMETRY_HASH</string>
     <string>--tenant-id</string><string>$TENANT</string>
     <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>
+    $adaptive_xml
     <string>--sink</string><string>$SINK</string>
     <string>--window</string><string>$WINDOW</string>
     <string>--block-size-tokens</string><string>$BLOCK_TOKENS</string>

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

@@ -22,7 +22,10 @@
     <string>--sink</string><string>4</string>
     <string>--window</string><string>2048</string>
     <string>--block-size-tokens</string><string>64</string>
-    <string>--cache-gb</string><string>0.25</string>
+    <string>--cache-gb</string><string>8</string>
+    <string>--cache-min-gb</string><string>0.25</string>
+    <string>--adaptive-cache</string>
+    <string>--memory-reserve-gb</string><string>2</string>
     <string>--prefill-tps</string><string>1</string>
     <string>--max-concurrent-jobs</string><string>1</string>
     <string>--network</string><string>thunderbolt</string>

docs/adr/0017-prefill-compute-worker-orchestration.md

@@ -81,6 +81,11 @@ operator-configured defaults.
 - Cache mounts are exposed as one content-addressed `kv://` namespace for
   management. This virtualizes naming and location only; fetch/import still
   copies the selected snapshot into Primary memory.
+- A successful remote import promotes the complete snapshot into Primary's
+  bounded hot LRU. Primary eviction removes only that hot copy; the worker's
+  cold/offload copy remains available.
+- Worker cache capacity is adaptive: physical memory minus active MLX model
+  bytes and an operator reserve, bounded by configured minimum and ceiling.
 - Snapshot payloads support zlib framing and retain SHA-256 of the uncompressed
   bytes.
 - Replication uses rendezvous hashing and a bounded replication factor instead

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

@@ -141,6 +141,10 @@ export KAKEYA_WORKER_ADVERTISE="<worker-ip>:53051"
 export KAKEYA_LAYER_GEOMETRY_HASH="<same-value-as-primary>"
 export KAKEYA_WORKER_SINK="4"
 export KAKEYA_WORKER_WINDOW="2048"
+export KAKEYA_WORKER_CACHE_GB="8"
+export KAKEYA_WORKER_CACHE_MIN_GB="0.25"
+export KAKEYA_WORKER_MEMORY_RESERVE_GB="2"
+export KAKEYA_WORKER_ADAPTIVE_CACHE="1"
 export KAKEYA_CACHE_BLOCK_TOKENS="64"
 export KAKEYA_CACHE_FORMAT_VERSION="kakeya-prefill-v3-kl-d4-q38"
 export KAKEYA_CACHE_COMPRESSION="kakeyalattice-d4"
@@ -240,7 +244,9 @@ curl -fsS http://127.0.0.1:8090/v1/network/kvfs
 
 The returned `kv://` URI and mount table virtualize naming and management only.
 Payloads remain in each Mac's physical RAM and are copied into Primary once
-before decode; `coherent_shared_memory` is always false.
+before decode; `coherent_shared_memory` is always false. Mounts are marked
+`hot` for Primary and `cold-offload` for worker/cache peers. Remote imports are
+promoted into the hot LRU, while eviction leaves the cold copy untouched.
 
 ## Maintenance cache saturation
 

inference_engine/distributed/kv_namespace.py

@@ -16,13 +16,20 @@ class VirtualKVMount:
     bytes_free: int
     entry_count: int
     network: str
+    tier: str
 
 
 class VirtualKVNamespace:
     """Present cache-node RAM as one lookup namespace, never as coherent RAM."""
 
-    def __init__(self, compatibility: CacheCompatibility) -> None:
+    def __init__(
+        self,
+        compatibility: CacheCompatibility,
+        *,
+        primary_node_id: str = "head-runtime",
+    ) -> None:
         self.compatibility = compatibility
+        self.primary_node_id = primary_node_id
         fingerprint = compatibility_fingerprint(compatibility).hex()
         tenant = compatibility.tenant_namespace or "default"
         self.uri = f"kv://{tenant}/{compatibility.model_id}/{fingerprint}"
@@ -41,6 +48,11 @@ def describe(self, nodes: Sequence[dict[str, Any]]) -> dict[str, Any]:
                 bytes_free=int(cache.get("bytes_free", 0)),
                 entry_count=int(cache.get("entry_count", 0)),
                 network=endpoint.get("network", "default"),
+                tier=(
+                    "hot"
+                    if node["id"] == self.primary_node_id
+                    else "cold-offload"
+                ),
             ))
         return {
             "uri": self.uri,

inference_engine/distributed/prefill_cache.py

@@ -274,6 +274,19 @@ def stats(self) -> CacheStats:
                 put_failures=self._put_failures,
             )
 
+    def resize(self, max_bytes: int) -> bool:
+        """Resize the LRU budget, evicting cold unleased blocks when shrinking."""
+        if max_bytes <= 0:
+            raise ValueError("max_bytes must be > 0")
+        with self._lock:
+            previous = self.max_bytes
+            self.max_bytes = int(max_bytes)
+            self._evict_to_budget()
+            if self._bytes_used > self.max_bytes:
+                self.max_bytes = max(previous, self._bytes_used)
+                return False
+            return True
+
     def block_hashes(self) -> tuple[bytes, ...]:
         with self._lock:
             return tuple(self._blocks)

inference_engine/distributed/prefill_cache_runtime.py

@@ -79,6 +79,9 @@ class PrefillReuseStats:
     publish_failures: int = 0
     bytes_published: int = 0
     last_publish_error: str = ""
+    hot_promotions: int = 0
+    hot_promotion_bytes: int = 0
+    hot_promotion_failures: int = 0
 
 
 class RemotePrefillRequiredError(RuntimeError):
@@ -269,6 +272,7 @@ def _try_import(
                 self.stats.local_hits += 1
             else:
                 self.stats.remote_hits += 1
+                self._promote_remote_hit(hit, payload, reused)
             return reused
         except Exception as exc:
             # Cache is an optimization. A corrupt/expired/unreachable hit must
@@ -280,6 +284,27 @@ def _try_import(
             verifier.reset()
             return 0
 
+    def _promote_remote_hit(
+        self,
+        hit: _Hit,
+        payload: bytes,
+        token_count: int,
+    ) -> None:
+        if not hit.block_hash:
+            return
+        try:
+            stored = self.local_store.put(CacheBlock.create(
+                hit.block_hash,
+                token_count,
+                payload,
+            ))
+        except ValueError:
+            self.stats.hot_promotion_failures += 1
+            return
+        if stored:
+            self.stats.hot_promotions += 1
+            self.stats.hot_promotion_bytes += len(payload)
+
     def _compute_remote(
         self,
         tokens: list[int],

inference_engine/network/state.py

@@ -26,7 +26,10 @@ def __init__(
     ) -> None:
         self.registry = registry
         self.cache_store = cache_store
-        self.kv_namespace = VirtualKVNamespace(cache_store.compatibility)
+        self.kv_namespace = VirtualKVNamespace(
+            cache_store.compatibility,
+            primary_node_id=registry.self_card.node_id,
+        )
         self.state_path = Path(state_path).expanduser()
         self.prefill_stats_provider = prefill_stats_provider
         self._lock = threading.RLock()

scripts/start_prefill_worker_node.py

@@ -68,6 +68,28 @@ def physical_memory_bytes() -> int:
         return 0
 
 
+def adaptive_cache_budget(
+    *,
+    total_bytes: int,
+    active_model_bytes: int,
+    ceiling_bytes: int,
+    minimum_bytes: int,
+    reserve_bytes: int,
+) -> int:
+    if min(total_bytes, ceiling_bytes, minimum_bytes) <= 0 or reserve_bytes < 0:
+        raise ValueError("adaptive cache budget inputs are invalid")
+    available = max(0, total_bytes - active_model_bytes - reserve_bytes)
+    return max(minimum_bytes, min(ceiling_bytes, available))
+
+
+def mlx_active_memory_bytes() -> int:
+    try:
+        import mlx.core as mx
+        return int(mx.get_active_memory())
+    except (AttributeError, RuntimeError):
+        return 0
+
+
 async def serve(args) -> None:
     compatibility = CacheCompatibility(
         model_id=args.cache_model_id or args.model_id,
@@ -96,9 +118,16 @@ async def serve(args) -> None:
             "MLX prefill workers require --max-concurrent-jobs 1 so the "
             "model and its stream remain on one compute thread",
         )
+    minimum_cache_bytes = int(args.cache_min_gb * (1 << 30))
+    cache_ceiling_bytes = int(args.cache_gb * (1 << 30))
+    if minimum_cache_bytes > cache_ceiling_bytes:
+        raise SystemExit("--cache-min-gb must be <= --cache-gb")
     store = PrefixCacheStore(
         compatibility,
-        max_bytes=int(args.cache_gb * (1 << 30)),
+        max_bytes=(
+            minimum_cache_bytes
+            if args.adaptive_cache else cache_ceiling_bytes
+        ),
         node_id=args.node_id,
     )
 
@@ -123,7 +152,29 @@ def engine_factory() -> MLXPrefillComputeEngine:
     )
     jobs.warmup()
 
+    def refresh_cache_budget() -> tuple[int, int]:
+        active = mlx_active_memory_bytes()
+        if args.adaptive_cache:
+            target = adaptive_cache_budget(
+                total_bytes=physical_memory_bytes(),
+                active_model_bytes=active,
+                ceiling_bytes=cache_ceiling_bytes,
+                minimum_bytes=minimum_cache_bytes,
+                reserve_bytes=int(args.memory_reserve_gb * (1 << 30)),
+            )
+            store.resize(target)
+        return active, store.stats().max_bytes
+
+    active_model_bytes, cache_budget_bytes = refresh_cache_budget()
+    _LOG.info(
+        "worker memory tiers: model_active=%d cache_budget=%d reserve=%.2fGiB",
+        active_model_bytes,
+        cache_budget_bytes,
+        args.memory_reserve_gb,
+    )
+
     def card() -> NodeCapability:
+        active_model_bytes, _ = refresh_cache_budget()
         inflight, queued, load, queued_tokens = jobs.stats()
         worker = PrefillWorkerCapability(
             compatibility=compatibility,
@@ -135,7 +186,9 @@ def card() -> NodeCapability:
             tokens_per_second_prefill=args.prefill_tps,
             ram_bytes_free=max(
                 0,
-                physical_memory_bytes() - store.stats().bytes_used,
+                physical_memory_bytes()
+                - active_model_bytes
+                - store.stats().bytes_used,
             ),
             queued_tokens=queued_tokens,
         )
@@ -240,6 +293,9 @@ def main() -> None:
     parser.add_argument("--sink", type=int, default=4)
     parser.add_argument("--window", type=int, default=64)
     parser.add_argument("--cache-gb", type=float, default=4.0)
+    parser.add_argument("--cache-min-gb", type=float, default=0.25)
+    parser.add_argument("--adaptive-cache", action="store_true")
+    parser.add_argument("--memory-reserve-gb", type=float, default=2.0)
     parser.add_argument("--cache-compression",
                         choices=["none", "zlib", "kakeyalattice-d4"],
                         default="zlib")

tests/inference_engine/bridge/test_prefill_worker_launchd.py

@@ -22,6 +22,8 @@ def test_worker_installer_emits_full_cache_compatibility_contract():
         "--rtt-ms",
         "--max-concurrent-jobs",
         "--max-prompt-tokens",
+        "--cache-min-gb",
+        "--memory-reserve-gb",
     ):
         assert f"<string>{flag}</string>" in source
     assert 'PEER="${KAKEYA_WORKER_PEER:-}"' in source
@@ -59,7 +61,9 @@ def test_two_mac_deployment_uses_allens_as_prefill_only():
     worker = WORKER_PLIST.read_text()
     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>0.25</string>" 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>--adaptive-cache</string>" in worker
     assert "<string>--window</string><string>2048</string>" in worker
     assert "<string>--prefill-tps</string><string>1</string>" in worker
     assert "scripts/start_prefill_cache_node.py" in PEER_PLIST.read_text()

tests/inference_engine/bridge/test_prefill_worker_memory.py

@@ -0,0 +1,40 @@
+import pytest
+
+from scripts.start_prefill_worker_node import adaptive_cache_budget
+
+
+def test_adaptive_budget_uses_only_model_headroom():
+    gib = 1 << 30
+    assert adaptive_cache_budget(
+        total_bytes=16 * gib,
+        active_model_bytes=10 * gib,
+        ceiling_bytes=8 * gib,
+        minimum_bytes=1 * gib,
+        reserve_bytes=2 * gib,
+    ) == 4 * gib
+    assert adaptive_cache_budget(
+        total_bytes=16 * gib,
+        active_model_bytes=15 * gib,
+        ceiling_bytes=8 * gib,
+        minimum_bytes=1 * gib,
+        reserve_bytes=2 * gib,
+    ) == 1 * gib
+
+
+def test_adaptive_budget_caps_spare_memory_and_validates():
+    gib = 1 << 30
+    assert adaptive_cache_budget(
+        total_bytes=32 * gib,
+        active_model_bytes=1 * gib,
+        ceiling_bytes=8 * gib,
+        minimum_bytes=1 * gib,
+        reserve_bytes=2 * gib,
+    ) == 8 * gib
+    with pytest.raises(ValueError):
+        adaptive_cache_budget(
+            total_bytes=0,
+            active_model_bytes=0,
+            ceiling_bytes=1,
+            minimum_bytes=1,
+            reserve_bytes=0,
+        )

tests/inference_engine/distributed/test_kv_namespace.py

@@ -6,7 +6,7 @@ def test_virtual_namespace_aggregates_matching_cache_mounts():
     namespace = VirtualKVNamespace(CacheCompatibility(
         model_id="gemma",
         tenant_namespace="private",
-    ))
+    ), primary_node_id="head")
     result = namespace.describe([
         {
             "id": "head",
@@ -42,6 +42,7 @@ def test_virtual_namespace_aggregates_matching_cache_mounts():
         "bytes_free": 20,
         "entry_count": 2,
         "network": "thunderbolt",
+        "tier": "hot",
     }]
 
 
@@ -54,3 +55,4 @@ def test_virtual_namespace_defaults_tenant_and_endpoint():
     assert result["uri"].startswith("kv://default/m/")
     assert result["mounts"][0]["address"] == ""
     assert result["mounts"][0]["network"] == "default"
+    assert result["mounts"][0]["tier"] == "cold-offload"

tests/inference_engine/distributed/test_prefill_cache.py

@@ -123,3 +123,23 @@ def test_put_rejects_when_active_lease_pins_capacity():
         store.put(second)
     assert store.block_hashes() == (first.block_hash,)
     assert store.stats().put_failures == 1
+
+
+def test_resize_evicts_cold_blocks_and_preserves_pinned_budget():
+    store = PrefixCacheStore(_compat(), max_bytes=10, node_id="x")
+    first = CacheBlock.create(bytes(32), 1, b"12345")
+    second = CacheBlock.create(bytes.fromhex("01" * 32), 1, b"abc")
+    store.put(first)
+    store.put(second)
+    assert store.resize(4)
+    assert store.block_hashes() == (second.block_hash,)
+    assert store.stats().max_bytes == 4
+    store.lookup([second.block_hash])
+    assert not store.resize(1)
+    assert store.stats().max_bytes == 4
+    try:
+        store.resize(0)
+    except ValueError:
+        pass
+    else:
+        raise AssertionError("expected resize validation")