diff --git a/deploy/install_prefill_worker_launchd.sh b/deploy/install_prefill_worker_launchd.sh index 8c62a5d3..82df7d14 100755 --- a/deploy/install_prefill_worker_launchd.sh +++ b/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="--peer$PEER" fi +adaptive_xml="" +if [[ "$ADAPTIVE_CACHE" == "1" ]]; then + adaptive_xml="--adaptive-cache" +fi cat > "$PLIST" < @@ -67,6 +74,9 @@ cat > "$PLIST" <--layer-geometry-hash$KAKEYA_LAYER_GEOMETRY_HASH --tenant-id$TENANT --cache-gb$CACHE_GB + --cache-min-gb$CACHE_MIN_GB + --memory-reserve-gb$MEMORY_RESERVE_GB + $adaptive_xml --sink$SINK --window$WINDOW --block-size-tokens$BLOCK_TOKENS diff --git a/deploy/launchd/ai.kakeya.prefill-worker-peer.plist b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist index cbe2ebb5..7bbe0dd3 100644 --- a/deploy/launchd/ai.kakeya.prefill-worker-peer.plist +++ b/deploy/launchd/ai.kakeya.prefill-worker-peer.plist @@ -22,7 +22,10 @@ --sink4 --window2048 --block-size-tokens64 - --cache-gb0.25 + --cache-gb8 + --cache-min-gb0.25 + --adaptive-cache + --memory-reserve-gb2 --prefill-tps1 --max-concurrent-jobs1 --networkthunderbolt diff --git a/docs/adr/0017-prefill-compute-worker-orchestration.md b/docs/adr/0017-prefill-compute-worker-orchestration.md index 944c3590..533913bd 100644 --- a/docs/adr/0017-prefill-compute-worker-orchestration.md +++ b/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 diff --git a/docs/ops/distributed-prefill-kv-network.md b/docs/ops/distributed-prefill-kv-network.md index 6e28d42b..f2e9c0e8 100644 --- a/docs/ops/distributed-prefill-kv-network.md +++ b/docs/ops/distributed-prefill-kv-network.md @@ -141,6 +141,10 @@ export KAKEYA_WORKER_ADVERTISE=":53051" export KAKEYA_LAYER_GEOMETRY_HASH="" 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 diff --git a/inference_engine/distributed/kv_namespace.py b/inference_engine/distributed/kv_namespace.py index 41e24ce0..287c8149 100644 --- a/inference_engine/distributed/kv_namespace.py +++ b/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, diff --git a/inference_engine/distributed/prefill_cache.py b/inference_engine/distributed/prefill_cache.py index a7ea7b91..638bd151 100644 --- a/inference_engine/distributed/prefill_cache.py +++ b/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) diff --git a/inference_engine/distributed/prefill_cache_runtime.py b/inference_engine/distributed/prefill_cache_runtime.py index 81c47f4e..94b1c03b 100644 --- a/inference_engine/distributed/prefill_cache_runtime.py +++ b/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], diff --git a/inference_engine/network/state.py b/inference_engine/network/state.py index 18133437..43a97594 100644 --- a/inference_engine/network/state.py +++ b/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() diff --git a/scripts/start_prefill_worker_node.py b/scripts/start_prefill_worker_node.py index b5ef96ea..e658719f 100644 --- a/scripts/start_prefill_worker_node.py +++ b/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") diff --git a/tests/inference_engine/bridge/test_prefill_worker_launchd.py b/tests/inference_engine/bridge/test_prefill_worker_launchd.py index 4d37b571..3044b677 100644 --- a/tests/inference_engine/bridge/test_prefill_worker_launchd.py +++ b/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"{flag}" 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 "--cache-gb0.25" in worker + assert "--cache-gb8" in worker + assert "--cache-min-gb0.25" in worker + assert "--adaptive-cache" in worker assert "--window2048" in worker assert "--prefill-tps1" in worker assert "scripts/start_prefill_cache_node.py" in PEER_PLIST.read_text() diff --git a/tests/inference_engine/bridge/test_prefill_worker_memory.py b/tests/inference_engine/bridge/test_prefill_worker_memory.py new file mode 100644 index 00000000..767a62be --- /dev/null +++ b/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, + ) diff --git a/tests/inference_engine/distributed/test_kv_namespace.py b/tests/inference_engine/distributed/test_kv_namespace.py index 60059b7d..03ef04b4 100644 --- a/tests/inference_engine/distributed/test_kv_namespace.py +++ b/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" diff --git a/tests/inference_engine/distributed/test_prefill_cache.py b/tests/inference_engine/distributed/test_prefill_cache.py index 01455a2e..f6fcb9f9 100644 --- a/tests/inference_engine/distributed/test_prefill_cache.py +++ b/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") diff --git a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py index bc2a0d39..8c07c613 100644 --- a/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py +++ b/tests/inference_engine/distributed/test_prefill_cache_runtime_fallback.py @@ -231,6 +231,23 @@ def test_remote_required_forces_compatible_worker_despite_cost(monkeypatch): hook.close() +def test_hot_promotion_failure_does_not_break_remote_import(): + hook = DistributedPrefillCacheHook(PrefixCacheStore( + CacheCompatibility(model_id="m"), + max_bytes=1, + node_id="head", + )) + hook._promote_remote_hit( + _Hit("peer", "lease", 1, 1, 2, block_hash=b"h" * 32), + b"too-large", + 1, + ) + assert hook.stats.hot_promotion_failures == 1 + hook._promote_remote_hit(_Hit("peer", "lease", 1, 1, 0), b"", 1) + assert hook.stats.hot_promotions == 0 + hook.close() + + def test_successful_local_import_suffix_and_on_reuse(monkeypatch): compatibility = CacheCompatibility(model_id="m", block_size_tokens=2) store = PrefixCacheStore(compatibility, max_bytes=1024, node_id="head") diff --git a/tests/inference_engine/distributed/test_prefill_orchestrator_e2e.py b/tests/inference_engine/distributed/test_prefill_orchestrator_e2e.py index d284e5fd..a31e9c92 100644 --- a/tests/inference_engine/distributed/test_prefill_orchestrator_e2e.py +++ b/tests/inference_engine/distributed/test_prefill_orchestrator_e2e.py @@ -158,6 +158,9 @@ async def test_dynamic_worker_computes_remote_prefill_and_head_imports( assert verifier.next_token_logits == "logits" assert hook.stats.remote_jobs == 1 assert hook.stats.remote_hits == 1 + assert hook.stats.hot_promotions == 1 + assert hook.stats.hot_promotion_bytes > 0 + assert len(head_store.block_hashes()) == 1 finally: hook.close() jobs.close()