diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh index f55c0bf0..f027fb0b 100755 --- a/deploy/install_prefill_worker_launchd.sh +++ b/deploy/install_prefill_worker_launchd.sh @@ -18,6 +18,16 @@ MODEL_REVISION="${KAKEYA_MODEL_REVISION:-}" TOKENIZER_REVISION="${KAKEYA_TOKENIZER_REVISION:-}" QUANTIZATION="${KAKEYA_CACHE_QUANTIZATION:-4bit-mlx}" ROPE_HASH="${KAKEYA_ROPE_HASH:-}" +SINK="${KAKEYA_WORKER_SINK:-4}" +WINDOW="${KAKEYA_WORKER_WINDOW:-64}" +BLOCK_TOKENS="${KAKEYA_CACHE_BLOCK_TOKENS:-64}" +PREFILL_TPS="${KAKEYA_WORKER_PREFILL_TPS:-20}" +NETWORK="${KAKEYA_WORKER_NETWORK:-lan}" +PRIORITY="${KAKEYA_WORKER_PRIORITY:-50}" +RTT_MS="${KAKEYA_WORKER_RTT_MS:-1.0}" +PEER="${KAKEYA_WORKER_PEER:-}" +MAX_CONCURRENT_JOBS="${KAKEYA_WORKER_MAX_CONCURRENT_JOBS:-1}" +MAX_PROMPT_TOKENS="${KAKEYA_WORKER_MAX_PROMPT_TOKENS:-131072}" LABEL="ai.kakeya.prefill-worker" PLIST="$HOME/Library/LaunchAgents/$LABEL.plist" LOG_DIR="$HOME/.kakeya" @@ -27,6 +37,10 @@ psk_xml="" if [[ -n "$PSK_FILE" ]]; then psk_xml="--fleet-psk-file$PSK_FILE" fi +peer_xml="" +if [[ -n "$PEER" ]]; then + peer_xml="--peer$PEER" +fi cat > "$PLIST" < @@ -49,6 +63,16 @@ cat > "$PLIST" <--layer-geometry-hash$KAKEYA_LAYER_GEOMETRY_HASH --tenant-id$TENANT --cache-gb$CACHE_GB + --sink$SINK + --window$WINDOW + --block-size-tokens$BLOCK_TOKENS + --prefill-tps$PREFILL_TPS + --network$NETWORK + --priority$PRIORITY + --rtt-ms$RTT_MS + --max-concurrent-jobs$MAX_CONCURRENT_JOBS + --max-prompt-tokens$MAX_PROMPT_TOKENS + $peer_xml $psk_xml WorkingDirectory$KAKEYA_WORKER_REPO diff --git a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist index c1237a60..7a5a3a58 100644 --- a/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist +++ b/deploy/launchd/ai.kakeya.grpc-runtime-prefill.plist @@ -17,19 +17,31 @@ --skip-cache-check --enable-prefill-cache --prefill-cache-gb1 - --cache-peer169.254.27.104:52051 + --peer169.254.27.104:53051 + --cache-peer169.254.27.104:53051 --cache-model-idgemma-4-26B-A4B-it-mlx-4bit --model-revisionlocal-4bit-v1 --tokenizer-revisiongemma4-v1 --cache-quantization4bit-mlx --cache-kv-dtypebfloat16 --cache-block-tokens64 + --cache-tenant-idprivate-fleet + --fleet-psk-file/Users/fluffy314/.kakeya/fleet.psk --node-idhead-runtime --advertise169.254.187.239:51051 --cache-advertise169.254.187.239:51051 --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__ + --network-state/Users/fluffy314/.kakeya/inference_network.json --network-telemetry-urlhttp://127.0.0.1:8090/v1/network/telemetry/tokens --network-telemetry-api-key__NETWORK_KEY__ --log-levelINFO diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 93f26b7d..baa090f8 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -126,8 +126,16 @@ 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_LAYER_GEOMETRY_HASH="" +export KAKEYA_WORKER_SINK="4" +export KAKEYA_WORKER_WINDOW="2048" +export KAKEYA_CACHE_BLOCK_TOKENS="64" +export KAKEYA_WORKER_NETWORK="thunderbolt" +export KAKEYA_WORKER_PRIORITY="100" +export KAKEYA_WORKER_RTT_MS="0.55" +export KAKEYA_WORKER_PEER="169.254.187.239:51051" export KAKEYA_FLEET_PSK_FILE="$HOME/.kakeya/fleet.psk" export KAKEYA_TENANT_ID="private-fleet" bash deploy/install_prefill_worker_launchd.sh @@ -158,7 +166,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:52051","region":"Private","role":"cache"}' + -d '{"alias":"peer-mini","address":"169.254.27.104:53051","region":"Private","role":"hybrid"}' ``` Create a paired group: @@ -192,8 +200,20 @@ Minimal acceptance: curl -fsS https://kakeya.ai/healthz curl -fsS https://kakeya.ai/v1/network/summary curl -fsS https://kakeya.ai/v1/network/tokens +curl -fsS https://kakeya.ai/v1/network/prefill + +PYTHONPATH=.:sdks/python python scripts/verify_remote_prefill_e2e.py \ + --address 127.0.0.1:51051 \ + --dashboard http://127.0.0.1:8090 \ + --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. + ## Rollback The cache is an optimization; inference correctness does not depend on it. diff --git a/inference_engine/network/api.py b/inference_engine/network/api.py index 79d8eb26..0a56213a 100644 --- a/inference_engine/network/api.py +++ b/inference_engine/network/api.py @@ -102,6 +102,10 @@ def tokens(): "hit_rate": summary["kv_hit_rate"], } + @app.get("/v1/network/prefill") + def prefill(): + return state.prefill_stats() + @app.post( "/v1/network/telemetry/tokens", dependencies=[Depends(require_key)], diff --git a/inference_engine/network/dashboard.py b/inference_engine/network/dashboard.py index 2802e920..eb4bb8c9 100644 --- a/inference_engine/network/dashboard.py +++ b/inference_engine/network/dashboard.py @@ -16,7 +16,7 @@ def dashboard_html() -> str: h1{font-size:24px;margin:0}h2{font-size:17px;margin:22px 0 10px}.sub,.muted{color:var(--muted)} button{background:transparent;border:1px solid var(--line);color:var(--text);border-radius:7px;padding:8px 12px;cursor:pointer} button.active,.primary{background:var(--accent);color:#0b1020;border-color:var(--accent);font-weight:650} -.tabs{display:flex;gap:7px}.stats,.grid2,.grid3{display:grid;gap:12px}.stats{grid-template-columns:repeat(5,1fr);margin-top:22px}.grid2{grid-template-columns:1.2fr .8fr}.grid3{grid-template-columns:repeat(3,1fr)} +.tabs{display:flex;gap:7px}.stats,.grid2,.grid3{display:grid;gap:12px}.stats{grid-template-columns:repeat(auto-fit,minmax(140px,1fr));margin-top:22px}.grid2{grid-template-columns:1.2fr .8fr}.grid3{grid-template-columns:repeat(3,1fr)} .card{background:var(--panel);border:1px solid var(--line);border-radius:9px;padding:15px}.stat b{font-size:23px;display:block}.stat span{color:var(--muted);font-size:12px} .map{height:280px;position:relative;overflow:hidden}.map:before{content:"";position:absolute;inset:34px;border:1px dashed var(--line);border-radius:45%} .dot{position:absolute;transform:translate(-50%,-50%)}.dot i{display:block;width:12px;height:12px;background:var(--ok);border:2px solid var(--bg);border-radius:50%}.dot.head i{width:16px;height:16px;background:var(--accent)}.dot small{white-space:nowrap} @@ -31,8 +31,8 @@ def dashboard_html() -> str:

Kakeya Inference Network

P2P Prefill KV sharing across trusted inference nodes
- -
0Online nodes
0Inference groups
0Completed tokens
0%KV-assisted tokens
0 GBShared cache online
+ +
0Online nodes
0Inference groups
0Completed tokens
0%KV-assisted tokens
0 GBShared cache online
0Remote prefill jobs
0Remote KV imports
0Tokens reused

Online node distribution

Live KV discovery

Waiting for node telemetry
Capability gossip and prefix lookups appear here.

Cache capacity

0 / 0 GB
@@ -49,7 +49,7 @@ def dashboard_html() -> str: $('createGroup').onclick=async()=>{await fetch('/v1/network/groups',{method:'POST',headers:writeHeaders(),body:JSON.stringify({name:$('groupName').value,node_ids:$('groupNodes').value.split(',').map(x=>x.trim()).filter(Boolean)})});load()}; function nodePosition(i,total){let a=(i/Math.max(total,1))*Math.PI*2;return {x:50+38*Math.cos(a),y:53+35*Math.sin(a)}} async function load(){let [s,n,g]=await Promise.all([fetch('/v1/network/summary').then(r=>r.json()),fetch('/v1/network/nodes').then(r=>r.json()),fetch('/v1/network/groups').then(r=>r.json())]); -$('online').textContent=s.online_nodes;$('groupCount').textContent=s.groups;$('tokens').textContent=fmt(s.completed_tokens);$('hitRate').textContent=(s.kv_hit_rate*100).toFixed(0)+'%';$('cache').textContent=gb(s.cache_bytes_used+s.cache_bytes_free)+' GB'; +$('online').textContent=s.online_nodes;$('groupCount').textContent=s.groups;$('tokens').textContent=fmt(s.completed_tokens);$('hitRate').textContent=(s.kv_hit_rate*100).toFixed(0)+'%';$('cache').textContent=gb(s.cache_bytes_used+s.cache_bytes_free)+' GB';let p=s.prefill||{};$('remoteJobs').textContent=fmt(p.remote_jobs);$('remoteHits').textContent=fmt(p.remote_hits);$('reusedTokens').textContent=fmt(p.tokens_reused); let total=s.cache_bytes_used+s.cache_bytes_free,pct=total?s.cache_bytes_used/total*100:0;$('capacityLabel').textContent=`${gb(s.cache_bytes_used)} / ${gb(total)} GB`;$('capacityBar').style.width=pct+'%'; $('map').innerHTML=n.map((x,i)=>{let p=nodePosition(i,n.length);return `
${x.region}
${x.alias}
`}).join(''); $('nodesBody').innerHTML=n.map(x=>`${x.alias}${x.role}${x.region}${x.cache?(x.cache.model_id+' / '+x.cache.format):'—'}${x.endpoint.network} ${x.endpoint.rtt_ms?x.endpoint.rtt_ms+'ms':''}${x.status}`).join(''); diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py index 0e194804..39695d4a 100644 --- a/inference_engine/network/state.py +++ b/inference_engine/network/state.py @@ -6,8 +6,9 @@ import secrets import threading import time +from dataclasses import asdict, is_dataclass from pathlib import Path -from typing import Any +from typing import Any, Callable from inference_engine.distributed.capability import CapabilityRegistry from inference_engine.distributed.prefill_cache import PrefixCacheStore @@ -20,10 +21,12 @@ def __init__( cache_store: PrefixCacheStore, *, state_path: str | Path, + prefill_stats_provider: Callable[[], Any] | None = None, ) -> None: self.registry = registry self.cache_store = cache_store self.state_path = Path(state_path).expanduser() + self.prefill_stats_provider = prefill_stats_provider self._lock = threading.RLock() self._data = self._load() @@ -218,8 +221,17 @@ def summary(self) -> dict[str, Any]: "local_lookup_hits": cache_stats.lookup_hits, "local_lookup_misses": cache_stats.lookup_misses, "local_tokens_served": cache_stats.tokens_served, + "prefill": self.prefill_stats(), } + def prefill_stats(self) -> dict[str, Any]: + if self.prefill_stats_provider is None: + return {} + stats = self.prefill_stats_provider() + if is_dataclass(stats) and not isinstance(stats, type): + return asdict(stats) + return dict(stats) + def topology(self) -> dict[str, Any]: nodes = self.nodes() edges = [] diff --git a/scripts/start_grpc_runtime_server.py b/scripts/start_grpc_runtime_server.py index acda21bd..2059e050 100755 --- a/scripts/start_grpc_runtime_server.py +++ b/scripts/start_grpc_runtime_server.py @@ -557,6 +557,10 @@ async def _serve(args: argparse.Namespace) -> int: registry, prefill_store, state_path=args.network_state, + prefill_stats_provider=( + (lambda: prefill_hook.stats) + if prefill_hook is not None else None + ), ) http_server = uvicorn.Server(uvicorn.Config( create_network_app( diff --git a/scripts/verify_remote_prefill_e2e.py b/scripts/verify_remote_prefill_e2e.py new file mode 100644 index 00000000..f2a02533 --- /dev/null +++ b/scripts/verify_remote_prefill_e2e.py @@ -0,0 +1,101 @@ +#!/usr/bin/env python3 +"""Run one cold-prefix request and prove that a remote prefill worker served it.""" +from __future__ import annotations + +import argparse +import json +import secrets +import time +import urllib.request + + +def _get_json(url: str): + with urllib.request.urlopen(url, timeout=5) as response: + return json.load(response) + + +def _delta(before: dict, after: dict, key: str) -> int: + return int(after.get(key, 0)) - int(before.get(key, 0)) + + +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) + args = parser.parse_args() + + from kakeya import Client + from transformers import AutoTokenizer + + from scripts.chat_grpc import _resolve_eos_token_ids + + tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_id) + nonce = secrets.token_hex(8) + sentence = ( + "Kakeya remote prefill verification context. " + f"Unique run {nonce}. " + ) + text = sentence + token_ids = tokenizer.encode(text, add_special_tokens=True) + while len(token_ids) < args.minimum_prefix_tokens: + text += sentence + token_ids = tokenizer.encode(text, add_special_tokens=True) + + nodes = _get_json(f"{args.dashboard}/v1/network/nodes") + workers = [ + node for node in nodes + if node.get("prefill_worker") and node.get("status") == "online" + ] + if not workers: + print(json.dumps({ + "ok": False, + "reason": "no online prefill worker capability", + "nodes": nodes, + }, indent=2)) + return 2 + + before = _get_json(f"{args.dashboard}/v1/network/prefill") + started = time.perf_counter() + with Client(args.address) as client: + with client.create_session( + eos_token_ids=_resolve_eos_token_ids(tokenizer), + client_label="remote-prefill-e2e", + ) as session: + session.append(token_ids) + list(session.generate(max_tokens=1)) + elapsed = time.perf_counter() - started + after = _get_json(f"{args.dashboard}/v1/network/prefill") + + 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 + ), + "worker_nodes": [node["id"] for node in workers], + "prefix_tokens": len(token_ids), + "wall_seconds": elapsed, + "delta": { + key: _delta(before, after, key) + for key in ( + "remote_jobs", + "remote_hits", + "tokens_reused", + "tokens_computed", + "bytes_received", + "remote_job_failures", + "fallbacks", + ) + }, + "before": before, + "after": after, + } + print(json.dumps(result, indent=2)) + return 0 if result["ok"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py new file mode 100644 index 00000000..4efd689b --- /dev/null +++ b/tests/inference_engine/bridge/test_prefill_worker_launchd.py @@ -0,0 +1,42 @@ +from pathlib import Path + + +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" + + +def test_worker_installer_emits_full_cache_compatibility_contract(): + source = INSTALLER.read_text() + for flag in ( + "--sink", + "--window", + "--block-size-tokens", + "--prefill-tps", + "--network", + "--priority", + "--rtt-ms", + "--max-concurrent-jobs", + "--max-prompt-tokens", + ): + assert f"{flag}" in source + assert 'PEER="${KAKEYA_WORKER_PEER:-}"' in source + assert "--peer" in source + + +def test_head_runtime_discovers_and_uses_worker_cache_port(): + plist = HEAD_PLIST.read_text() + assert ( + "--peer169.254.27.104:53051" + in plist + ) + assert ( + "--cache-peer169.254.27.104:53051" + in plist + ) + assert "--primary-prefill-penalty-ms" in plist + assert ( + "--cache-tenant-idprivate-fleet" + in plist + ) + assert "--fleet-psk-file" in plist diff --git a/tests/inference_engine/network/test_network_api.py b/tests/inference_engine/network/test_network_api.py index 97957ba6..51a69402 100644 --- a/tests/inference_engine/network/test_network_api.py +++ b/tests/inference_engine/network/test_network_api.py @@ -18,6 +18,11 @@ def _client(tmp_path): CapabilityRegistry(NodeCapability(node_id="head", grpc_address="head:1")), PrefixCacheStore(compatibility, max_bytes=100, node_id="head"), state_path=tmp_path / "state.json", + prefill_stats_provider=lambda: { + "remote_jobs": 3, + "remote_hits": 2, + "tokens_reused": 192, + }, ) client = TestClient(create_network_app(state, api_key="secret")) client.network_state = state @@ -27,13 +32,16 @@ def _client(tmp_path): def test_dashboard_health_and_read_apis(tmp_path): client = _client(tmp_path) assert client.get("/").status_code == 200 - assert "Kakeya Inference Network" in client.get("/network").text + dashboard = client.get("/network").text + assert "Kakeya Inference Network" in dashboard + assert "Remote prefill jobs" in dashboard assert client.get("/healthz").json()["status"] == "ok" assert client.get("/v1/network/summary").json()["online_nodes"] == 1 assert len(client.get("/v1/network/nodes").json()) == 1 assert client.get("/v1/network/groups").json() == [] assert "nodes" in client.get("/v1/network/topology").json() assert client.get("/v1/network/tokens").json()["completed"] == 0 + assert client.get("/v1/network/prefill").json()["remote_jobs"] == 3 events = client.get("/v1/network/events?once=true") assert events.status_code == 200 assert "event: summary" in events.text diff --git a/tests/inference_engine/network/test_network_state.py b/tests/inference_engine/network/test_network_state.py index 418671eb..8c969a9d 100644 --- a/tests/inference_engine/network/test_network_state.py +++ b/tests/inference_engine/network/test_network_state.py @@ -9,6 +9,7 @@ PrefillWorkerCapability, ) from inference_engine.distributed.prefill_cache import PrefixCacheStore +from inference_engine.distributed.prefill_cache_runtime import PrefillReuseStats from inference_engine.network.state import NetworkState @@ -39,6 +40,11 @@ def _state(tmp_path): CapabilityRegistry(self_card=card), store, state_path=tmp_path / "network.json", + prefill_stats_provider=lambda: { + "remote_jobs": 2, + "remote_hits": 1, + "tokens_reused": 128, + }, ) @@ -57,6 +63,8 @@ def test_registration_groups_tokens_and_persistence(tmp_path): assert summary["registered_nodes"] == 2 assert summary["completed_tokens"] == 100 assert summary["kv_hit_rate"] == 0.7 + assert summary["prefill"]["remote_jobs"] == 2 + assert state.prefill_stats()["tokens_reused"] == 128 assert state.groups()[0]["id"] == group["id"] assert state.topology()["edges"][0]["target"] == "peer" @@ -102,3 +110,14 @@ def test_invalid_persisted_state_falls_back_to_empty(tmp_path): ) assert state.groups() == [] assert state.summary()["completed_tokens"] == 0 + assert state.prefill_stats() == {} + + +def test_prefill_stats_serializes_runtime_dataclass(tmp_path): + state = _state(tmp_path) + state.prefill_stats_provider = lambda: PrefillReuseStats( + remote_jobs=4, + remote_hits=3, + tokens_reused=256, + ) + assert state.prefill_stats()["remote_jobs"] == 4