diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist index 2751dc04..560e9085 100644 --- a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist +++ b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist @@ -17,8 +17,8 @@ --skip-cache-check --enable-prefill-cache --prefill-cache-gb1 - --peer169.254.27.104:53051 - --cache-peer169.254.27.104:53051 + --peer169.254.27.104:52051 + --cache-peer169.254.27.104:52051 --cache-model-idgemma-4-26B-A4B-it-mlx-4bit --model-revisionlocal-4bit-v1 --tokenizer-revisiongemma4-v1 @@ -35,11 +35,8 @@ --network-labelthunderbolt --network-priority100 --measured-rtt-ms0.55 - --remote-prefill-min-tokens128 --cache-link-mbps10000 --cache-default-rtt-ms0.55 - --prefill-min-savings-ratio0 - --primary-prefill-penalty-ms1000 --network-http-host127.0.0.1 --network-http-port8090 --network-api-key__NETWORK_KEY__ diff --git a/deploy/launchd/ai.kakeya.prefill-network-peer.plist b/deploy/launchd/ai.kakeya.prefill-network-peer.plist index e6ad1bd1..8a06bea7 100644 --- a/deploy/launchd/ai.kakeya.prefill-network-peer.plist +++ b/deploy/launchd/ai.kakeya.prefill-network-peer.plist @@ -4,26 +4,33 @@ Labelai.kakeya.prefill-network ProgramArguments - /Users/allen/kakeya-prefill-venv/bin/python - /Users/allen/kakeya-prefill-network/start_prefill_cache_node.py + /Users/allen/.venv-distwan/bin/python + /Users/allen/Kakeya-LLM-Inference-engine/scripts/start_prefill_cache_node.py --node-idallens-mini --bind169.254.27.104:52051 --advertise169.254.27.104:52051 --model-idgemma-4-26B-A4B-it-mlx-4bit --model-revisionlocal-4bit-v1 --tokenizer-revisiongemma4-v1 + --cache-format-versionkakeya-prefill-v3-kl-d4-q38 + --cache-compressionkakeyalattice-d4 --quantization4bit-mlx --layer-geometry-hash93d9585b0f06b60bac8e1cadf50b29df1adbf086c862e61720b6127d22c30e2b - --cache-gb2 + --tenant-idprivate-fleet + --fleet-psk-file/Users/allen/.kakeya/fleet.psk + --sink4 + --window2048 + --block-size-tokens64 + --cache-gb8 --networkthunderbolt --priority100 --rtt-ms0.55 --state-path/Users/allen/.kakeya/inference_network_peer.json - WorkingDirectory/Users/allen/kakeya-prefill-network + WorkingDirectory/Users/allen/Kakeya-LLM-Inference-engine EnvironmentVariables - PATH/Users/allen/kakeya-prefill-venv/bin:/usr/bin:/bin:/usr/sbin:/sbin - PYTHONPATH/Users/allen/kakeya-prefill-network + PATH/Users/allen/.venv-distwan/bin:/usr/bin:/bin:/usr/sbin:/sbin + PYTHONPATH/Users/allen/Kakeya-LLM-Inference-engine:/Users/allen/Kakeya-LLM-Inference-engine/sdks/python RunAtLoad KeepAlive diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 1c7df1bb..cbb77ab5 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -92,6 +92,11 @@ 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. + Check from the head over Thunderbolt: ```bash @@ -109,6 +114,11 @@ 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. + Create a fleet PSK once and copy it to every trusted node: ```bash @@ -127,8 +137,8 @@ export KAKEYA_CACHE_MODEL_ID="gemma-4-26B-A4B-it-mlx-4bit" export KAKEYA_MODEL_REVISION="local-4bit-v1" export KAKEYA_TOKENIZER_REVISION="gemma4-v1" export KAKEYA_WORKER_NODE_ID="prefill-mini-1" -export KAKEYA_WORKER_BIND="169.254.27.104:53051" -export KAKEYA_WORKER_ADVERTISE="169.254.27.104:53051" +export KAKEYA_WORKER_BIND=":53051" +export KAKEYA_WORKER_ADVERTISE=":53051" export KAKEYA_LAYER_GEOMETRY_HASH="" export KAKEYA_WORKER_SINK="4" export KAKEYA_WORKER_WINDOW="2048" @@ -149,7 +159,7 @@ The primary must use the same compatibility and auth values: ```text --enable-prefill-cache --enable-capability-exchange ---peer 169.254.27.104:53051 +--peer :53051 --cache-tenant-id private-fleet --fleet-psk-file ~/.kakeya/fleet.psk --cache-compression zlib @@ -175,7 +185,7 @@ KEY="$(cat ~/.kakeya/network_api_key)" curl -fsS -X POST http://127.0.0.1:8090/v1/network/nodes/register \ -H "Content-Type: application/json" \ -H "X-API-Key: $KEY" \ - -d '{"alias":"peer-mini","address":"169.254.27.104:53051","region":"Private","role":"hybrid"}' + -d '{"alias":"allens-mini","address":"169.254.27.104:52051","region":"Private","role":"cache"}' ``` Create a paired group: @@ -217,11 +227,11 @@ PYTHONPATH=.:sdks/python python scripts/verify_remote_prefill_e2e.py \ --tokenizer-id ~/kakeya-models/gemma-4-26B-A4B-it-mlx-4bit ``` -The verifier exits non-zero unless a live worker capability is present and one -cold unique prefix increments `remote_jobs`, `remote_hits`, and -`tokens_reused`. Decode throughput is reported separately; remote prefill is -accepted on lower TTFT/prefill time and higher request throughput, not a change -to single-stream decode tokens/s. +By default the verifier accepts a remote cache hit and requires `remote_hits` +and `tokens_reused` to increase. Use `--require-worker` only when testing the +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. ## Rollback diff --git a/scripts/verify_remote_prefill_e2e.py b/scripts/verify_remote_prefill_e2e.py index f2a02533..2b04afcf 100644 --- a/scripts/verify_remote_prefill_e2e.py +++ b/scripts/verify_remote_prefill_e2e.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -"""Run one cold-prefix request and prove that a remote prefill worker served it.""" +"""Prove a remote cache hit, or explicitly require remote compute plus import.""" from __future__ import annotations import argparse @@ -18,12 +18,30 @@ def _delta(before: dict, after: dict, key: str) -> int: return int(after.get(key, 0)) - int(before.get(key, 0)) +def _acceptance( + before: dict, + after: dict, + *, + minimum_prefix_tokens: int, + require_worker: bool, +) -> tuple[bool, str]: + remote_hit = ( + _delta(before, after, "remote_hits") >= 1 + and _delta(before, after, "tokens_reused") >= minimum_prefix_tokens + ) + remote_compute = remote_hit and _delta(before, after, "remote_jobs") >= 1 + path = "remote_compute" if remote_compute else "remote_cache" if remote_hit else "none" + return (remote_compute if require_worker else remote_hit), path + + def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--address", default="127.0.0.1:51051") parser.add_argument("--dashboard", default="http://127.0.0.1:8090") parser.add_argument("--tokenizer-id", required=True) parser.add_argument("--minimum-prefix-tokens", type=int, default=128) + parser.add_argument("--head-node-id", default="head-runtime") + parser.add_argument("--require-worker", action="store_true") args = parser.parse_args() from kakeya import Client @@ -48,10 +66,22 @@ def main() -> int: node for node in nodes if node.get("prefill_worker") and node.get("status") == "online" ] - if not workers: + cache_nodes = [ + node for node in nodes + if ( + node["id"] != args.head_node_id + and node.get("cache") + and node.get("status") == "online" + ) + ] + if not cache_nodes or (args.require_worker and not workers): print(json.dumps({ "ok": False, - "reason": "no online prefill worker capability", + "reason": ( + "no online prefill worker capability" + if args.require_worker and not workers + else "no online remote cache capability" + ), "nodes": nodes, }, indent=2)) return 2 @@ -68,13 +98,16 @@ def main() -> int: elapsed = time.perf_counter() - started after = _get_json(f"{args.dashboard}/v1/network/prefill") + accepted, path = _acceptance( + before, + after, + minimum_prefix_tokens=args.minimum_prefix_tokens, + require_worker=args.require_worker, + ) result = { - "ok": ( - _delta(before, after, "remote_jobs") >= 1 - and _delta(before, after, "remote_hits") >= 1 - and _delta(before, after, "tokens_reused") - >= args.minimum_prefix_tokens - ), + "ok": accepted, + "accepted_path": path, + "cache_nodes": [node["id"] for node in cache_nodes], "worker_nodes": [node["id"] for node in workers], "prefix_tokens": len(token_ids), "wall_seconds": elapsed, diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index e32d0ced..b1dce440 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -4,6 +4,7 @@ ROOT = Path(__file__).resolve().parents[3] INSTALLER = ROOT / "deploy" / "install_prefill_worker_launchd.sh" HEAD_PLIST = ROOT / "deploy" / "launchd" / "ai.kakeya.grpc-runtime-prefill.plist" +PEER_PLIST = ROOT / "deploy" / "launchd" / "ai.kakeya.prefill-network-peer.plist" def test_worker_installer_emits_full_cache_compatibility_contract(): @@ -30,17 +31,17 @@ def test_worker_installer_emits_full_cache_compatibility_contract(): assert 'launchctl kickstart -k "$DOMAIN/$LABEL"' in source -def test_head_runtime_discovers_and_uses_worker_cache_port(): +def test_two_mac_deployment_uses_allens_as_cache_only(): plist = HEAD_PLIST.read_text() assert ( - "--peer169.254.27.104:53051" + "--peer169.254.27.104:52051" in plist ) assert ( - "--cache-peer169.254.27.104:53051" + "--cache-peer169.254.27.104:52051" in plist ) - assert "--primary-prefill-penalty-ms" in plist + assert "--primary-prefill-penalty-ms" not in plist assert ( "--cache-tenant-idprivate-fleet" in plist @@ -51,3 +52,8 @@ def test_head_runtime_discovers_and_uses_worker_cache_port(): "kakeyalattice-d4" in plist ) + peer = PEER_PLIST.read_text() + assert "scripts/start_prefill_cache_node.py" in peer + assert "scripts/start_prefill_worker_node.py" not in peer + assert "--cache-gb8" in peer + assert "--window2048" in peer diff --git a/tests/inference_engine/bridge/test_remote_prefill_acceptance.py b/tests/inference_engine/bridge/test_remote_prefill_acceptance.py new file mode 100644 index 00000000..521c2cae --- /dev/null +++ b/tests/inference_engine/bridge/test_remote_prefill_acceptance.py @@ -0,0 +1,38 @@ +from scripts.verify_remote_prefill_e2e import _acceptance + + +def test_cache_only_replay_is_accepted_without_worker_job(): + accepted, path = _acceptance( + {"remote_hits": 0, "remote_jobs": 0, "tokens_reused": 0}, + {"remote_hits": 1, "remote_jobs": 0, "tokens_reused": 512}, + minimum_prefix_tokens=512, + require_worker=False, + ) + assert accepted + assert path == "remote_cache" + + +def test_worker_mode_requires_compute_job_and_remote_import(): + before = {"remote_hits": 0, "remote_jobs": 0, "tokens_reused": 0} + cache_only = {"remote_hits": 1, "remote_jobs": 0, "tokens_reused": 512} + assert _acceptance( + before, + cache_only, + minimum_prefix_tokens=512, + require_worker=True, + ) == (False, "remote_cache") + assert _acceptance( + before, + {**cache_only, "remote_jobs": 1}, + minimum_prefix_tokens=512, + require_worker=True, + ) == (True, "remote_compute") + + +def test_insufficient_reuse_is_rejected(): + assert _acceptance( + {"remote_hits": 0, "remote_jobs": 0, "tokens_reused": 0}, + {"remote_hits": 1, "remote_jobs": 0, "tokens_reused": 64}, + minimum_prefix_tokens=128, + require_worker=False, + ) == (False, "none")