From 953b5865d8ab9b8299cab464054918d9d31dd7a2 Mon Sep 17 00:00:00 2001 From: fluffy314 Date: Thu, 16 Jul 2026 20:32:09 +0800 Subject: [PATCH] feat(prefill): add tiered KV hot promotion and offload Promote successful remote snapshots into Primary's bounded hot LRU while retaining cold copies on allens, and adapt worker offload capacity to real MLX model memory headroom with explicit reserve and safety bounds. Co-authored-by: Cursor --- deploy/install_prefill_worker_launchd.sh | 10 ++++ .../ai.kakeya.prefill-worker-peer.plist | 5 +- ...17-prefill-compute-worker-orchestration.md | 5 ++ docs/ops/distributed-prefill-kv-network.md | 8 ++- inference_engine/distributed/kv_namespace.py | 14 ++++- inference_engine/distributed/prefill_cache.py | 13 ++++ .../distributed/prefill_cache_runtime.py | 25 ++++++++ inference_engine/network/state.py | 5 +- scripts/start_prefill_worker_node.py | 60 ++++++++++++++++++- .../bridge/test_prefill_worker_launchd.py | 6 +- .../bridge/test_prefill_worker_memory.py | 40 +++++++++++++ .../distributed/test_kv_namespace.py | 4 +- .../distributed/test_prefill_cache.py | 20 +++++++ .../test_prefill_cache_runtime_fallback.py | 17 ++++++ .../test_prefill_orchestrator_e2e.py | 3 + 15 files changed, 227 insertions(+), 8 deletions(-) create mode 100644 tests/inference_engine/bridge/test_prefill_worker_memory.py 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()