feat(prefill): enforce decode-only primary mode

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

@@ -17,8 +17,8 @@
     <string>--skip-cache-check</string>
     <string>--enable-prefill-cache</string>
     <string>--prefill-cache-gb</string><string>1</string>
-    <string>--peer</string><string>169.254.27.104:52051</string>
-    <string>--cache-peer</string><string>169.254.27.104:52051</string>
+    <string>--peer</string><string>169.254.27.104:53051</string>
+    <string>--cache-peer</string><string>169.254.27.104:53051</string>
     <string>--cache-model-id</string><string>gemma-4-26B-A4B-it-mlx-4bit</string>
     <string>--model-revision</string><string>local-4bit-v1</string>
     <string>--tokenizer-revision</string><string>gemma4-v1</string>
@@ -37,6 +37,9 @@
     <string>--measured-rtt-ms</string><string>0.55</string>
     <string>--cache-link-mbps</string><string>10000</string>
     <string>--cache-default-rtt-ms</string><string>0.55</string>
+    <string>--remote-prefill-min-tokens</string><string>0</string>
+    <string>--prefill-worker-timeout-s</string><string>300</string>
+    <string>--prefill-policy</string><string>remote-required</string>
     <string>--network-http-host</string><string>127.0.0.1</string>
     <string>--network-http-port</string><string>8090</string>
     <string>--network-api-key</string><string>__NETWORK_KEY__</string>

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

@@ -0,0 +1,43 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE plist PUBLIC "-//Apple//DTD PLIST 1.0//EN"
+  "http://www.apple.com/DTDs/PropertyList-1.0.dtd">
+<plist version="1.0"><dict>
+  <key>Label</key><string>ai.kakeya.prefill-worker</string>
+  <key>ProgramArguments</key><array>
+    <string>/Users/allen/.venv-distwan/bin/python</string>
+    <string>/Users/allen/Kakeya-LLM-Inference-engine/scripts/start_prefill_worker_node.py</string>
+    <string>--node-id</string><string>allens-mini</string>
+    <string>--bind</string><string>169.254.27.104:53051</string>
+    <string>--advertise</string><string>169.254.27.104:53051</string>
+    <string>--model-id</string><string>/Users/allen/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit</string>
+    <string>--cache-model-id</string><string>gemma-4-26B-A4B-it-mlx-4bit</string>
+    <string>--model-revision</string><string>local-4bit-v1</string>
+    <string>--tokenizer-revision</string><string>gemma4-v1</string>
+    <string>--cache-format-version</string><string>kakeya-prefill-v3-kl-d4-q38</string>
+    <string>--cache-compression</string><string>kakeyalattice-d4</string>
+    <string>--quantization</string><string>4bit-mlx</string>
+    <string>--layer-geometry-hash</string><string>93d9585b0f06b60bac8e1cadf50b29df1adbf086c862e61720b6127d22c30e2b</string>
+    <string>--tenant-id</string><string>private-fleet</string>
+    <string>--fleet-psk-file</string><string>/Users/allen/.kakeya/fleet.psk</string>
+    <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>--prefill-tps</string><string>1</string>
+    <string>--max-concurrent-jobs</string><string>1</string>
+    <string>--network</string><string>thunderbolt</string>
+    <string>--priority</string><string>100</string>
+    <string>--rtt-ms</string><string>0.55</string>
+    <string>--log-level</string><string>INFO</string>
+  </array>
+  <key>WorkingDirectory</key><string>/Users/allen/Kakeya-LLM-Inference-engine</string>
+  <key>EnvironmentVariables</key><dict>
+    <key>PATH</key><string>/Users/allen/.venv-distwan/bin:/usr/bin:/bin:/usr/sbin:/sbin</string>
+    <key>PYTHONPATH</key><string>/Users/allen/Kakeya-LLM-Inference-engine:/Users/allen/Kakeya-LLM-Inference-engine/sdks/python</string>
+  </dict>
+  <key>RunAtLoad</key><true/>
+  <key>KeepAlive</key><true/>
+  <key>ProcessType</key><string>Interactive</string>
+  <key>StandardOutPath</key><string>/Users/allen/.kakeya/prefill-worker.log</string>
+  <key>StandardErrorPath</key><string>/Users/allen/.kakeya/prefill-worker.log</string>
+</dict></plist>

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

@@ -49,6 +49,12 @@ Every remote error (lookup, job, lease, fetch, checksum, decompress, import)
 resets the verifier and falls back to full local prefill. Cache availability
 must never determine request correctness.
 
+For deployments that require a strictly decode-only primary,
+`--prefill-policy remote-required` changes this failure contract: only a
+complete cache hit or completed remote worker job is accepted. Partial hits are
+not extended on Primary, cost gating is bypassed, and worker failure returns
+`UNAVAILABLE` instead of silently running local prefill.
+
 ### Discovery and placement
 
 Capability gossip is the only membership source. Static `--cache-peer` and
@@ -72,6 +78,9 @@ operator-configured defaults.
 
 - Decode KV remains local to the primary.
 - Peer memory is a pre-decode snapshot tier, not coherent remote attention RAM.
+- 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.
 - 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

@@ -92,10 +92,9 @@ Use an isolated venv and copy/sync the repository package. The peer plist is:
 deploy/launchd/ai.kakeya.prefill-network-peer.plist
 ```
 
-In the two-Mac profile, `allens-mini` runs this role only. It does not load the
-model or apply a chat template. Snapshots arrive from the primary's fallback
-prefill today, and from separately deployed compute workers when the fleet has
-additional machines.
+The cache-only plist remains available for rollback or additional RAM-only
+nodes. In the strict two-Mac decode/prefill profile, allens instead runs the
+prefill-worker plist below with a co-located cache.
 
 Check from the head over Thunderbolt:
 
@@ -114,10 +113,10 @@ The worker loads the exact same MLX model as the primary, accepts queued
 prefill-only jobs, writes immutable snapshots into its co-located RAM cache and
 never serves user decode.
 
-This is an additional fleet role, not the `allens-mini` cache-only role in the
-two-Mac profile. Deploy it on Worker A/B/C addresses when those machines exist.
-Workers receive canonical token IDs from the scheduler; they do not construct
-their own chat template.
+In the strict two-Mac profile, `allens-mini` runs this role and Primary uses
+`--prefill-policy remote-required`. Workers receive canonical token IDs from
+the scheduler; they do not construct their own chat template. The worker stores
+the resulting snapshots in its co-located content-addressed cache.
 
 Create a fleet PSK once and copy it to every trusted node:
 
@@ -233,6 +232,16 @@ separate Worker A/B/C path; that mode additionally requires `remote_jobs`.
 Decode throughput is reported separately because all autoregressive decode
 remains on the primary.
 
+The logical cross-node KV namespace is available at:
+
+```bash
+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.
+
 ## Maintenance cache saturation
 
 Enable the bounded, memory-only first-append capture queue on the primary:

inference_engine/distributed/kv_namespace.py

@@ -0,0 +1,53 @@
+"""Logical content-addressed namespace over physically separate KV stores."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Sequence
+
+from inference_engine.distributed.capability import CacheCompatibility
+from inference_engine.distributed.prefill_cache import compatibility_fingerprint
+
+
+@dataclass(frozen=True)
+class VirtualKVMount:
+    node_id: str
+    address: str
+    bytes_used: int
+    bytes_free: int
+    entry_count: int
+    network: str
+
+
+class VirtualKVNamespace:
+    """Present cache-node RAM as one lookup namespace, never as coherent RAM."""
+
+    def __init__(self, compatibility: CacheCompatibility) -> None:
+        self.compatibility = compatibility
+        fingerprint = compatibility_fingerprint(compatibility).hex()
+        tenant = compatibility.tenant_namespace or "default"
+        self.uri = f"kv://{tenant}/{compatibility.model_id}/{fingerprint}"
+
+    def describe(self, nodes: Sequence[dict[str, Any]]) -> dict[str, Any]:
+        mounts = []
+        for node in nodes:
+            cache = node.get("cache")
+            if not cache or cache.get("model_id") != self.compatibility.model_id:
+                continue
+            endpoint = node.get("endpoint") or {}
+            mounts.append(VirtualKVMount(
+                node_id=node["id"],
+                address=endpoint.get("address", ""),
+                bytes_used=int(cache.get("bytes_used", 0)),
+                bytes_free=int(cache.get("bytes_free", 0)),
+                entry_count=int(cache.get("entry_count", 0)),
+                network=endpoint.get("network", "default"),
+            ))
+        return {
+            "uri": self.uri,
+            "access": "content-addressed-lookup-fetch-import",
+            "coherent_shared_memory": False,
+            "mounts": [mount.__dict__ for mount in mounts],
+            "bytes_used": sum(mount.bytes_used for mount in mounts),
+            "bytes_free": sum(mount.bytes_free for mount in mounts),
+            "entry_count": sum(mount.entry_count for mount in mounts),
+        }

inference_engine/distributed/prefill_cache_runtime.py

@@ -43,6 +43,7 @@
 from inference_engine.distributed.prefill_scheduler import (
     PrefillCostConfig,
     choose_prefill_worker,
+    compatible_prefill_workers,
     remote_import_wins,
     select_cache_replicas,
 )
@@ -80,6 +81,10 @@ class PrefillReuseStats:
     last_publish_error: str = ""
 
 
+class RemotePrefillRequiredError(RuntimeError):
+    """Raised when decode-only primary policy cannot obtain a complete KV."""
+
+
 @dataclass(frozen=True)
 class _Hit:
     source: str
@@ -114,6 +119,7 @@ def __init__(
         cost_config: PrefillCostConfig | None = None,
         auth: FleetAuthConfig | None = None,
         on_reuse=None,
+        require_remote_compute: bool = False,
     ) -> None:
         if min(
             lookup_timeout_s,
@@ -143,6 +149,7 @@ def __init__(
         self.replication_factor = int(replication_factor)
         self.cost_config = cost_config or PrefillCostConfig()
         self.auth = auth
+        self.require_remote_compute = bool(require_remote_compute)
         self._hash_key = auth.tenant_hash_key() if auth is not None else b""
         self.stats = PrefillReuseStats()
         self._stats_lock = threading.Lock()
@@ -167,14 +174,29 @@ def prepare(self, verifier: Any, token_ids: Sequence[int]) -> int:
         )
         hit = self._best_hit(hashes)
         reused = 0
+        if (
+            self.require_remote_compute
+            and hit is not None
+            and hit.hit_tokens != len(tokens)
+        ):
+            hit = None
         if hit is not None:
             reused = self._try_import(verifier, tokens, hit)
-        elif len(tokens) >= self.remote_compute_min_tokens:
+        if reused == 0 and (
+            self.require_remote_compute
+            or len(tokens) >= self.remote_compute_min_tokens
+        ):
             remote_hit = self._compute_remote(tokens, hashes)
             if remote_hit is not None:
                 reused = self._try_import(verifier, tokens, remote_hit)
         if reused == 0:
             self.stats.misses += 1
+        if self.require_remote_compute and reused != len(tokens):
+            verifier.reset()
+            reason = self.stats.last_fallback_reason or (
+                "no compatible remote prefill worker completed the request"
+            )
+            raise RemotePrefillRequiredError(reason)
 
         self._compute_and_publish(verifier, tokens, hashes, reused)
         return reused
@@ -267,15 +289,28 @@ def _compute_remote(
             card for card in self._cards()
             if card.node_id != self.local_store.node_id
         )
-        target = choose_prefill_worker(
-            cards,
-            self.compatibility,
-            prompt_tokens=len(tokens),
-            estimated_snapshot_bytes=(
-                len(tokens) * self.estimated_snapshot_bytes_per_token
-            ),
-            config=self.cost_config,
-        )
+        if self.require_remote_compute:
+            candidates = compatible_prefill_workers(cards, self.compatibility)
+            target = min(
+                candidates,
+                key=lambda item: (
+                    item.capability.load,
+                    item.capability.queued_tokens,
+                    -item.capability.tokens_per_second_prefill,
+                    item.node_id,
+                ),
+                default=None,
+            )
+        else:
+            target = choose_prefill_worker(
+                cards,
+                self.compatibility,
+                prompt_tokens=len(tokens),
+                estimated_snapshot_bytes=(
+                    len(tokens) * self.estimated_snapshot_bytes_per_token
+                ),
+                config=self.cost_config,
+            )
         if target is None:
             return None
         request = distributed_pb2.SubmitPrefillJobRequest(

inference_engine/network/api.py

@@ -107,6 +107,10 @@ def create_group(request: CreateGroupRequest):
     def topology():
         return state.topology()
 
+    @app.get("/v1/network/kvfs")
+    def virtual_kv_file():
+        return state.virtual_kv_file()
+
     @app.get("/v1/network/tokens")
     def tokens():
         summary = state.summary()

inference_engine/network/state.py

@@ -11,6 +11,7 @@
 from typing import Any, Callable
 
 from inference_engine.distributed.capability import CapabilityRegistry
+from inference_engine.distributed.kv_namespace import VirtualKVNamespace
 from inference_engine.distributed.prefill_cache import PrefixCacheStore
 
 
@@ -25,6 +26,7 @@ def __init__(
     ) -> None:
         self.registry = registry
         self.cache_store = cache_store
+        self.kv_namespace = VirtualKVNamespace(cache_store.compatibility)
         self.state_path = Path(state_path).expanduser()
         self.prefill_stats_provider = prefill_stats_provider
         self._lock = threading.RLock()
@@ -257,6 +259,9 @@ def topology(self) -> dict[str, Any]:
                 } for target in ids[1:])
         return {"nodes": nodes, "edges": edges}
 
+    def virtual_kv_file(self) -> dict[str, Any]:
+        return self.kv_namespace.describe(self.nodes())
+
     def _load(self) -> dict[str, Any]:
         if self.state_path.exists():
             try:

inference_engine/server/grpc_app.py

@@ -36,6 +36,9 @@
 import grpc
 
 from inference_engine.memory.pool import PoolExhausted
+from inference_engine.distributed.prefill_cache_runtime import (
+    RemotePrefillRequiredError,
+)
 from inference_engine.server.proto_gen.kakeya.v1 import (
     runtime_pb2,
     runtime_pb2_grpc,
@@ -210,6 +213,8 @@ async def AppendTokens(  # noqa: N802 — gRPC-generated method casing
             )
         except SessionNotFoundError as exc:
             await context.abort(grpc.StatusCode.NOT_FOUND, str(exc))
+        except RemotePrefillRequiredError as exc:
+            await context.abort(grpc.StatusCode.UNAVAILABLE, str(exc))
         except ValueError as exc:
             await context.abort(grpc.StatusCode.INVALID_ARGUMENT, str(exc))
         except InvariantViolation as exc:

scripts/start_grpc_runtime_server.py

@@ -453,6 +453,7 @@ async def _serve(args: argparse.Namespace) -> int:
                 primary_compute_penalty_ms=args.primary_prefill_penalty_ms,
             ),
             auth=prefill_auth,
+            require_remote_compute=(args.prefill_policy == "remote-required"),
             on_reuse=(
                 (lambda count: telemetry_callback(count, count))
                 if telemetry_callback is not None else None
@@ -767,6 +768,13 @@ def main() -> int:
                          "with prefill; drives work to compute peers.")
     ap.add_argument("--remote-prefill-min-tokens", type=int, default=128)
     ap.add_argument("--prefill-worker-timeout-s", type=float, default=120.0)
+    ap.add_argument(
+        "--prefill-policy",
+        choices=["local-fallback", "remote-required"],
+        default="local-fallback",
+        help="Use remote workers opportunistically, or require complete remote "
+             "prefill so the primary remains decode-only.",
+    )
     ap.add_argument("--network-label", default="lan",
                     help="Advertised interface: thunderbolt|lan|tailscale|public.")
     ap.add_argument("--network-priority", type=int, default=50)