diff --git a/src/exo/worker/engines/mlx/builder.py b/src/exo/worker/engines/mlx/builder.py index af7c75bb9d..b7e4697f7c 100644 --- a/src/exo/worker/engines/mlx/builder.py +++ b/src/exo/worker/engines/mlx/builder.py @@ -80,7 +80,7 @@ def build( self.tokenizer.tool_parser, # type: ignore ) - kv_prefix_cache = KVPrefixCache(self.group) + kv_prefix_cache = KVPrefixCache(self.group, model_id=self.model_id) device_rank = 0 if self.group is None else self.group.rank() if os.environ.get("EXO_NO_BATCH"): diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index 7cdcc77fbe..167bb36236 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/src/exo/worker/engines/mlx/cache.py @@ -1,14 +1,25 @@ +import contextlib import gc + +# KV disk persistence (offload) helpers +import hashlib +import json import os +import time as _time +from collections.abc import Callable from copy import deepcopy -from typing import TYPE_CHECKING +from pathlib import Path as _Path +from typing import TYPE_CHECKING, Any, cast import mlx.core as mx import numpy as np import psutil +from mlx.utils import tree_flatten, tree_unflatten from mlx_lm.models.cache import ( ArraysCache, CacheList, + ChunkedKVCache, + ConcatenateKVCache, KVCache, QuantizedKVCache, RotatingKVCache, @@ -21,6 +32,7 @@ ) from mlx_lm.tokenizer_utils import TokenizerWrapper +from exo.shared.constants import EXO_CACHE_HOME from exo.shared.types.memory import Memory from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS from exo.worker.engines.mlx.types import KVCacheType, Model @@ -229,8 +241,130 @@ def has_non_kv_caches(cache: KVCacheType) -> bool: return any(is_non_trimmable_cache_entry(c) for c in cache) +# ── KV disk persistence: state (de)serialization helpers ── +# +# Cache states are not pure array trees: GLM/DeepSeek-V3.2 DSA indexer caches +# hold zero-width value arrays (safetensors rejects size-0 tensors) and +# DeepseekV4Cache branch state holds Nones, ints and int-lists. Such leaves are +# recorded as JSON "placeholders" in the slot's _meta.json and re-inserted on +# load, so the .safetensors file itself stays byte-compatible with stock +# mlx-lm save_prompt_cache for standard architectures. + +_CACHE_CLASS_REGISTRY: dict[str, type] = { + "KVCache": KVCache, + "RotatingKVCache": RotatingKVCache, + "QuantizedKVCache": QuantizedKVCache, + "ArraysCache": ArraysCache, + "CacheList": CacheList, + "ConcatenateKVCache": ConcatenateKVCache, + "ChunkedKVCache": ChunkedKVCache, + "DeepseekV4Cache": DeepseekV4Cache, +} + + +def _is_int_list_leaf(x: object) -> bool: + # mlx tree_flatten silently drops empty lists, so int-lists (DeepseekV4Cache + # buffer/pool lengths, possibly []) must be kept whole as single leaves. + return isinstance(x, list) and all( + isinstance(i, int) for i in cast("list[object]", x) + ) + + +def _partition_cache_state( + cache: KVCacheType, +) -> tuple[dict[str, mx.array], dict[str, dict[str, object]]]: + """Split flattened per-layer cache state into safetensors-storable arrays + plus JSON placeholder specs for the leaves safetensors cannot hold.""" + states: list[object] = [cast(object, c.state) for c in cache] + flat = cast( + "list[tuple[str, object]]", + tree_flatten(states, is_leaf=_is_int_list_leaf), + ) + arrays: dict[str, mx.array] = {} + placeholders: dict[str, dict[str, object]] = {} + for path, leaf in flat: + if isinstance(leaf, mx.array): + if leaf.size > 0: + arrays[path] = leaf + else: + placeholders[path] = { + "kind": "empty", + "shape": list(leaf.shape), + "dtype": str(leaf.dtype).split(".")[-1], + } + elif leaf is None: + placeholders[path] = {"kind": "none"} + elif isinstance(leaf, (bool, int, float, str, list)): + placeholders[path] = {"kind": "json", "value": leaf} + else: + raise TypeError(f"unserializable cache state leaf at {path}: {type(leaf)}") + return arrays, placeholders + + +def _materialize_placeholder(spec: dict[str, object]) -> object: + kind = spec.get("kind") + if kind == "none": + return None + if kind == "empty": + shape = cast("list[int]", spec["shape"]) + dtype = cast("mx.Dtype", getattr(mx, cast(str, spec["dtype"]))) + return mx.zeros(shape, dtype) + if kind == "json": + return spec["value"] + raise ValueError(f"unknown cache placeholder kind: {kind}") + + +def _reconstruct_cache_entry(class_name: str, state: object, meta: object) -> object: + cls = _CACHE_CLASS_REGISTRY[class_name] # KeyError -> caller fail-safes + if cls is DeepseekV4Cache: + # No usable from_state: __init__ needs the sliding window, which is the + # wrapped RotatingKVCache's max_size = meta_state[1]. + meta_strs = cast("list[str]", meta) + v4 = DeepseekV4Cache(int(meta_strs[1])) + v4.state = state + v4.meta_state = meta + return v4 + if cls is CacheList: + # Resolve inner classes through this registry — mlx-lm's + # CacheList.from_state only looks up names in its own module globals. + names, metas = cast("tuple[list[str], list[object]]", meta) + inner_states = cast("list[object]", state) + return CacheList( + *( + _reconstruct_cache_entry(n, s, m) + for n, s, m in zip(names, inner_states, metas, strict=True) + ) + ) + from_state = cast("Callable[[object, object], object]", cls.from_state) + return from_state(state, meta) + + +def _load_slot_cache( + cache_file: str, placeholders: dict[str, dict[str, object]] +) -> list[object]: + """exo-side load_prompt_cache: re-inserts placeholder leaves and rebuilds + cache classes, including ones mlx-lm's loader cannot (DeepseekV4Cache).""" + arrays, metadata = cast( + "tuple[dict[str, mx.array], dict[str, str]]", + cast(object, mx.load(cache_file, return_metadata=True)), + ) + items: list[tuple[str, object]] = list(arrays.items()) + items += [ + (path, _materialize_placeholder(spec)) for path, spec in placeholders.items() + ] + state_tree = cast("list[object]", tree_unflatten(items)) + info, _, classes = cast( + "tuple[list[object], object, list[str]]", + tree_unflatten(list(metadata.items())), + ) + return [ + _reconstruct_cache_entry(c, s, m) + for c, s, m in zip(classes, state_tree, info, strict=True) + ] + + class KVPrefixCache: - def __init__(self, group: mx.distributed.Group | None): + def __init__(self, group: mx.distributed.Group | None, model_id: str = ""): self.prompts: list[mx.array] = [] # mx array of tokens (ints) self.caches: list[KVCacheType] = [] self._snapshots: list[list[CacheSnapshot] | None] = [] @@ -239,6 +373,18 @@ def __init__(self, group: mx.distributed.Group | None): self.prefill_tps: list[float] = [] self._access_counter: int = 0 self._group = group + # KV disk persistence (offload). Opt-in: + # set EXO_KV_DISK_PERSISTENCE=1 to enable. + self._model_id = model_id + self._disk_enabled = os.environ.get("EXO_KV_DISK_PERSISTENCE", "0") == "1" + self._disk_dir = ( + self._init_disk_dir() if model_id and self._disk_enabled else None + ) + self._disk_dirty = False + self._flush_requested_at: float = 0.0 + self._hot_slot_disk_id: int | None = None + if self._disk_dir is not None: + self._janitor_sweep_stale_slots() def clear(self): """Clear all cached prompts and caches.""" @@ -257,8 +403,19 @@ def add_kv_cache( media_regions: list["MediaRegion"] | None = None, prefill_tps: float = 0.0, ): - """Add a new cache entry. Evicts LRU entries if memory is high.""" - self._evict_if_needed() + """Add a new cache entry. With disk on: single hot slot (flush current to disk first). Else: LRU evict.""" + if self._disk_dir and len(self.caches) > 0: + if self._disk_dirty: + self._flush_hot_slot() + self.prompts.clear() + self.caches.clear() + self._snapshots.clear() + self._media_regions.clear() + self._last_used.clear() + self.prefill_tps.clear() + self._hot_slot_disk_id = None + else: + self._evict_if_needed() self.prompts.append(prompt_tokens) self.caches.append(deepcopy(cache)) self._snapshots.append(ssm_snapshots) @@ -266,6 +423,9 @@ def add_kv_cache( self.prefill_tps.append(prefill_tps) self._access_counter += 1 self._last_used.append(self._access_counter) + self._disk_dirty = True + if not self._flush_requested_at: + self._flush_requested_at = _time.time() logger.info(f"KV cache added: {len(prompt_tokens)} tokens") def update_kv_cache( @@ -293,8 +453,34 @@ def update_kv_cache( self.prefill_tps[index] = prefill_tps self._access_counter += 1 self._last_used[index] = self._access_counter + self._disk_dirty = True + if not self._flush_requested_at: + self._flush_requested_at = _time.time() logger.info(f"KV cache updated (index {index}): {len(prompt_tokens)} tokens") + def should_update_entry( + self, + matched_index: int | None, + prefix_hit_length: int, + min_prefix_hit_length: int = 1000, + ) -> bool: + """True only when the new prompt EXTENDS the matched cached entry. + + A sibling conversation that merely shares a long prefix (e.g. two + agent sessions with a common system prompt + bootstrap) must NOT be + saved as an update: with disk persistence the update inherits the + matched entry's disk slot id, so the next flush overwrites the OTHER + conversation's slot and destroys its cache (observed live as two + sessions ping-ponging a single slot). Extension test: the hit covers + the stored prompt up to the prefill-rollback slack. + """ + if matched_index is None or matched_index >= len(self.prompts): + return False + if prefix_hit_length < min_prefix_hit_length: + return False + cached_len = int(self.prompts[matched_index].shape[0]) + return prefix_hit_length >= cached_len - 8 + def _get_snapshot( self, entry_index: int, target_token_count: int ) -> tuple[int, CacheSnapshot | None]: @@ -357,7 +543,21 @@ def get_kv_cache( if length > best_length: best_index, best_length = i, length + # Disk fallback: in-memory hit is only a partial prefix of a different slot + if best_index is not None and self._disk_dir: + _cached_len = len(self.prompts[best_index]) + if best_length < _cached_len - 1: + _disk_result = self._try_load_from_disk( + model, prompt_tokens, min_prefix=best_length + ) + if _disk_result is not None: + return _disk_result + if best_index is None: + if self._disk_dir: + _disk_result = self._try_load_from_disk(model, prompt_tokens) + if _disk_result is not None: + return _disk_result return make_kv_cache(model), prompt_tokens, None, False # For exact match: trim to max_length-1 so remaining has the last token @@ -455,6 +655,329 @@ def _evict_if_needed(self): gc.collect() mx.clear_cache() + # ── KV disk persistence (offload) ── + + def _init_disk_dir(self): + h = hashlib.sha256(self._model_id.encode()).hexdigest()[:16] + base = _Path( + os.environ.get("EXO_KV_DISK_PATH", str(EXO_CACHE_HOME / "kv-cache")) + ) + d = base / h + d.mkdir(parents=True, exist_ok=True) + logger.info(f"KV cache disk dir: {d}") + return d + + def _list_disk_slots(self) -> list[int]: + disk_dir = self._disk_dir + if disk_dir is None: + return [] + slots: list[int] = [] + for f in disk_dir.glob("slot_*_tokens.safetensors"): + try: + slots.append(int(f.stem.split("_")[1])) + except (IndexError, ValueError): + continue + return sorted(slots) + + def _next_disk_slot_id(self) -> int: + existing = self._list_disk_slots() + return max(existing) + 1 if existing else 0 + + def _flush_hot_slot(self) -> None: + """Save the current hot slot to disk immediately (atomic write).""" + disk_dir = self._disk_dir + if disk_dir is None or len(self.caches) == 0: + return + try: + disk_dir.mkdir(parents=True, exist_ok=True) + slot_id = ( + self._hot_slot_disk_id + if self._hot_slot_disk_id is not None + else self._next_disk_slot_id() + ) + base = disk_dir / f"slot_{slot_id}" + # Inline save_prompt_cache, but state leaves safetensors cannot hold + # (zero-width arrays, Nones, ints — GLM DSA / DeepSeek-V4) are + # recorded as placeholders in _meta.json instead of being dropped. + cache = list(self.caches[0]) + cache_data, placeholders = _partition_cache_state(cache) + cache_info: list[Any] = [c.meta_state for c in cache] # pyright: ignore[reportUnknownMemberType] + cache_classes = [type(c).__name__ for c in cache] + cache_metadata: dict[str, str] = cast( + "dict[str, str]", dict(tree_flatten([cache_info, {}, cache_classes])) + ) + tmp_cache = str(base) + "_tmp_cache.safetensors" + mx.save_safetensors(tmp_cache, cache_data, cache_metadata) # pyright: ignore[reportUnknownMemberType] + os.rename(tmp_cache, str(base) + "_cache.safetensors") + meta = { + "model_id": self._model_id, + "token_count": int(len(self.prompts[0])), + "timestamp": _time.time(), + "format": 2, + "placeholders": placeholders, + } + # ".json.tmp" suffix, NOT "slot_N_tmp_meta.json": the latter would + # match the slot_*_meta.json globs in _evict_stale_disk_slots. + tmp_meta = str(base) + "_meta.json.tmp" + with open(tmp_meta, "w") as f: + json.dump(meta, f) + os.rename(tmp_meta, str(base) + "_meta.json") + # Tokens written last: slot discovery globs *_tokens.safetensors, + # so the slot only becomes visible once all three files exist. + tmp_tokens = str(base) + "_tmp_tokens.safetensors" + mx.save_safetensors(tmp_tokens, {"tokens": self.prompts[0]}) # pyright: ignore[reportUnknownMemberType] + os.rename(tmp_tokens, str(base) + "_tokens.safetensors") + self._hot_slot_disk_id = slot_id + self._disk_dirty = False + self._flush_requested_at = 0.0 + logger.info( + f"KV cache flushed to disk: slot_{slot_id} ({len(self.prompts[0])} tokens)" + ) + except Exception as e: + logger.warning(f"KV cache disk flush failed: {e}") + + def _janitor_sweep_stale_slots(self) -> None: + """TTL-sweep OTHER models' slot dirs under the kv-cache root. + + Runs at init and from per-flush eviction: the own-dir TTL pass only + covers this model, so without this sweep, slots of models that are + never loaded again would only fall to the global size cap. + """ + disk_dir = self._disk_dir + if disk_dir is None: + return + max_age_hours = float(os.environ.get("EXO_KV_DISK_TTL_HOURS", "24")) + cutoff = _time.time() - (max_age_hours * 3600) + for meta_file in disk_dir.parent.glob("*/slot_*_meta.json"): + try: + # Own dir is handled by _evict_stale_disk_slots, which knows + # the hot slot; other models' dirs are swept on timestamp only. + if meta_file.parent == disk_dir: + continue + with open(meta_file) as f: + meta = cast("dict[str, Any]", json.load(f)) + if cast(float, meta.get("timestamp", 0)) >= cutoff: + continue + base = str(meta_file)[: -len("_meta.json")] + for ext in [ + "_cache.safetensors", + "_tokens.safetensors", + "_tokens.npy", # legacy slot format + "_meta.json", + ]: + with contextlib.suppress(FileNotFoundError): + os.remove(base + ext) + logger.info( + f"KV cache janitor: evicted stale " + f"{meta_file.parent.name}/{meta_file.stem.removesuffix('_meta')}" + ) + except Exception: + continue + + def _evict_stale_disk_slots(self) -> None: + """Delete stale own-dir slots (TTL), then enforce the GLOBAL size cap: + while the whole kv-cache root exceeds EXO_KV_DISK_MAX_SIZE_GB, evict + the globally oldest slot regardless of owning model (cross-model LRU). + """ + disk_dir = self._disk_dir + if disk_dir is None: + return + max_age_hours = float(os.environ.get("EXO_KV_DISK_TTL_HOURS", "24")) + max_size_gb = float(os.environ.get("EXO_KV_DISK_MAX_SIZE_GB", "500")) + cutoff = _time.time() - (max_age_hours * 3600) + for meta_file in disk_dir.glob("slot_*_meta.json"): + try: + with open(meta_file) as f: + meta = cast("dict[str, Any]", json.load(f)) + if cast(float, meta.get("timestamp", 0)) < cutoff: + slot_id = int(meta_file.stem.split("_")[1]) + if slot_id == self._hot_slot_disk_id: + continue + base = disk_dir / f"slot_{slot_id}" + for ext in [ + "_cache.safetensors", + "_tokens.safetensors", + "_meta.json", + ]: + with contextlib.suppress(FileNotFoundError): + os.remove(str(base) + ext) + logger.info(f"KV cache evicted stale disk slot_{slot_id}") + except Exception: + continue + # TTL for everyone else too — a single model kept loaded for weeks + # would otherwise leave other models' stale slots to the size cap only. + self._janitor_sweep_stale_slots() + max_bytes = max_size_gb * 1024 * 1024 * 1024 + while True: + slots: list[tuple[float, _Path, int]] = [] + total_size = 0 + for meta_file in disk_dir.parent.glob("*/slot_*_meta.json"): + try: + with open(meta_file) as f: + meta = cast("dict[str, Any]", json.load(f)) + slot_id = int(meta_file.stem.split("_")[1]) + if ( + meta_file.parent == disk_dir + and slot_id == self._hot_slot_disk_id + ): + continue + base = meta_file.parent / f"slot_{slot_id}" + slot_size = sum( + f.stat().st_size + for ext in [ + "_cache.safetensors", + "_tokens.safetensors", + "_meta.json", + ] + if (f := _Path(str(base) + ext)).exists() + ) + total_size += slot_size + slots.append( + (cast(float, meta.get("timestamp", 0)), base, slot_size) + ) + except Exception: + continue + if total_size <= max_bytes or not slots: + break + slots.sort() + _oldest_ts, oldest_base, oldest_size = slots[0] + for ext in ["_cache.safetensors", "_tokens.safetensors", "_meta.json"]: + with contextlib.suppress(FileNotFoundError): + os.remove(str(oldest_base) + ext) + logger.info( + f"KV cache evicted disk slot " + f"{oldest_base.parent.name}/{oldest_base.name} " + f"({oldest_size / 1024 / 1024 / 1024:.1f} GB) — " + f"kv-cache over {max_size_gb} GB total" + ) + + def flush_to_disk(self, force: bool = False) -> None: + """Flush the hot slot if dirty and idle for 15s (or force=True).""" + if not self._disk_dir or not self._disk_dirty: + return + if not force and (_time.time() - self._flush_requested_at) < 15: + return + self._flush_hot_slot() + self._evict_stale_disk_slots() + + def _search_disk(self, prompt_tokens: mx.array) -> tuple[int | None, int]: + """Return (slot_id, prefix_length) of the best on-disk prefix match.""" + disk_dir = self._disk_dir + if disk_dir is None: + return None, 0 + best_id: int | None = None + best_length = 0 + for slot_id in self._list_disk_slots(): + if slot_id == self._hot_slot_disk_id: + continue + token_file = disk_dir / f"slot_{slot_id}_tokens.safetensors" + try: + cached_tokens = mx.load(str(token_file))["tokens"] + prefix_len = get_prefix_length(prompt_tokens, cached_tokens) + except Exception: + continue + if prefix_len > best_length: + best_length = prefix_len + best_id = slot_id + return best_id, best_length + + def _try_load_from_disk( + self, model: Model, prompt_tokens: mx.array, min_prefix: int = 0 + ) -> tuple[KVCacheType, mx.array, int, bool] | None: + """Swap in a matching disk slot. Returns (cache, remaining, index, is_exact) or None.""" + disk_dir = self._disk_dir + disk_id, prefix_len = self._search_disk(prompt_tokens) + if ( + disk_id is None + or disk_dir is None + or prefix_len <= min_prefix + or prefix_len < 1000 + ): + return None + try: + base = disk_dir / f"slot_{disk_id}" + # Load into locals FIRST. Do NOT mutate in-memory state until the load + # succeeds — otherwise an incompatible/corrupt slot (e.g. an architecture + # whose cache doesn't round-trip, like GLM) would leave the cache empty + # and crash the caller. This way a bad slot fails safe (recompute). + meta_path = str(base) + "_meta.json" + placeholders: dict[str, dict[str, Any]] = {} + if os.path.exists(meta_path): + with open(meta_path) as f: + slot_meta = cast("dict[str, Any]", json.load(f)) + # Legacy (format-1) slots have no placeholders — same load path. + placeholders = cast( + "dict[str, dict[str, Any]]", slot_meta.get("placeholders") or {} + ) + cache: KVCacheType = cast( + KVCacheType, + _load_slot_cache(str(base) + "_cache.safetensors", placeholders), + ) + tokens = mx.load(str(base) + "_tokens.safetensors")["tokens"] + cached_length = cache_length(cache) + if cached_length <= 0 or cached_length > int(tokens.shape[0]): + logger.warning( + f"KV cache disk slot_{disk_id} inconsistent " + f"(cache {cached_length} vs tokens {int(tokens.shape[0])}) — skipping" + ) + return None + # Slots are stored with the cache a couple of tokens SHORTER than + # the token file (prefill rollback), so cap the reusable prefix at + # what the cache actually holds — same min() as get_kv_cache. + prefix_len = min(prefix_len, cached_length) + if prefix_len <= min_prefix: + # After capping, no better than the in-memory match that + # triggered this lookup — keep the in-memory entry. + return None + tokens_to_trim = cached_length - prefix_len + if tokens_to_trim > 0 and has_non_kv_caches(cache): + # Partial prefix needs trimming, but this architecture's cache + # (e.g. DeepSeek-V4, SSM) can't be trimmed and disk slots carry + # no snapshots — recompute instead of corrupting/crashing. + logger.info( + f"KV cache disk slot_{disk_id}: partial prefix on " + f"non-trimmable cache — recomputing" + ) + return None + prompt_cache = deepcopy(cache) + except Exception as e: + logger.warning(f"KV cache disk load failed for slot_{disk_id}: {e}") + return None + + # Load succeeded — now it is safe to flush the current hot slot and swap in. + if self._disk_dirty and len(self.caches) > 0: + self._flush_hot_slot() + self.prompts.clear() + self.caches.clear() + self._snapshots.clear() + self._media_regions.clear() + self._last_used.clear() + self.prefill_tps.clear() + self.prompts.append(tokens) + self.caches.append(cache) + self._snapshots.append(None) + self._media_regions.append([]) + self._access_counter += 1 + self._last_used.append(self._access_counter) + self.prefill_tps.append(0.0) + self._hot_slot_disk_id = disk_id + self._disk_dirty = False + self._flush_requested_at = 0.0 + logger.info(f"KV cache loaded from disk: slot_{disk_id} ({len(tokens)} tokens)") + if tokens_to_trim > 0: + trim_cache(prompt_cache, tokens_to_trim) + for c in prompt_cache: + # Mirrors get_kv_cache: these types either can't be trimmed or + # expose offset as a read-only property (DeepseekV4Cache) — + # unreachable here thanks to the non-trimmable guard above, + # kept as defense. + if isinstance(c, (ArraysCache, RotatingKVCache, DeepseekV4Cache)): + continue + if hasattr(c, "offset"): + c.offset = prefix_len + remaining = prompt_tokens[prefix_len:] + return prompt_cache, remaining, 0, False + def get_memory_used_percentage(self) -> float: local_pressure: float = get_memory_used_percentage() diff --git a/src/exo/worker/engines/mlx/disaggregated/serve.py b/src/exo/worker/engines/mlx/disaggregated/serve.py index 5220e9d1fd..f437a7429a 100644 --- a/src/exo/worker/engines/mlx/disaggregated/serve.py +++ b/src/exo/worker/engines/mlx/disaggregated/serve.py @@ -60,8 +60,10 @@ def run_prefill_for_request( if kv_prefix_cache is not None: try: cache_snapshots = [snapshot_ssm_states(cache)] - hit_ratio = prefix_hit_length / n_tokens if n_tokens > 0 else 0.0 - if matched_index is not None and hit_ratio >= 0.5: + if kv_prefix_cache.should_update_entry( + matched_index, prefix_hit_length, min_prefix_hit_length=0 + ): + assert matched_index is not None kv_prefix_cache.update_kv_cache( matched_index, prompt_tokens, diff --git a/src/exo/worker/engines/mlx/generator/batch_generate.py b/src/exo/worker/engines/mlx/generator/batch_generate.py index 5c7394ddec..495f70a3ce 100644 --- a/src/exo/worker/engines/mlx/generator/batch_generate.py +++ b/src/exo/worker/engines/mlx/generator/batch_generate.py @@ -58,7 +58,6 @@ ) from exo.worker.runner.bootstrap import logger -_MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 REMOTE_PREFILL_MIN_TOKENS = 1000 @@ -505,15 +504,10 @@ def _save_prefix_cache( return try: - hit_ratio = ( - prefix_hit_length / len(all_prompt_tokens) - if len(all_prompt_tokens) > 0 - else 0.0 - ) - if matched_index is not None and ( - prefix_hit_length >= min_prefix_hit_length - and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE + if self.kv_prefix_cache.should_update_entry( + matched_index, prefix_hit_length, min_prefix_hit_length ): + assert matched_index is not None self.kv_prefix_cache.update_kv_cache( matched_index, all_prompt_tokens, diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 2e3d051251..f7aa928ff3 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -76,8 +76,6 @@ generation_stream = mx.new_stream(mx.default_device()) -_MIN_PREFIX_HIT_RATIO_TO_UPDATE = 0.5 - @contextlib.contextmanager def patch_embed_tokens( @@ -679,15 +677,10 @@ def mlx_generate( prefill_tps = kv_prefix_cache.prefill_tps[matched_index] if kv_prefix_cache is not None: - hit_ratio = ( - prefix_hit_length / len(all_prompt_tokens) - if len(all_prompt_tokens) > 0 - else 0.0 - ) - if matched_index is not None and ( - prefix_hit_length >= min_prefix_hit_length - and hit_ratio >= _MIN_PREFIX_HIT_RATIO_TO_UPDATE + if kv_prefix_cache.should_update_entry( + matched_index, prefix_hit_length, min_prefix_hit_length ): + assert matched_index is not None kv_prefix_cache.update_kv_cache( matched_index, all_prompt_tokens, diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index ac5d054808..d73bb17a0b 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -3,7 +3,10 @@ import time from dataclasses import dataclass from enum import Enum -from typing import BinaryIO +from typing import TYPE_CHECKING, BinaryIO, cast + +if TYPE_CHECKING: + from exo.worker.engines.mlx.cache import KVPrefixCache from anyio import ClosedResourceError, EndOfStream @@ -310,8 +313,17 @@ def handle_first_task(self, task: Task): f"Received {task.__class__.__name__} outside of state machine in {self.current_status=}" ) + def _flush_kv_cache_to_disk(self) -> None: + """Force-flush the hot KV slot to disk; no-op for engines without one.""" + kv = cast( + "KVPrefixCache | None", getattr(self.generator, "kv_prefix_cache", None) + ) + if kv is not None: + kv.flush_to_disk(force=True) + def shutdown(self, task: Task): logger.info("runner shutting down") + self._flush_kv_cache_to_disk() self.update_status(RunnerShuttingDown()) self.acknowledge_task(task) self.generator.close() @@ -380,6 +392,11 @@ def handle_generation_tasks(self, starting_task: GenerationTask): f"Received {item.__class__.__name__} outside of state machine in {self.current_status=}" ) + # Generation queue drained — persist the hot KV slot while idle so a + # later crash doesn't lose it (otherwise it is only flushed on + # conversation switch). + self._flush_kv_cache_to_disk() + self.update_status(RunnerReady(prefill_server_port=self._prefill_server_port)) logger.info("runner ready") diff --git a/src/exo/worker/tests/unittests/test_mlx/test_kv_disk_persistence.py b/src/exo/worker/tests/unittests/test_mlx/test_kv_disk_persistence.py new file mode 100644 index 0000000000..4c2a252657 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_mlx/test_kv_disk_persistence.py @@ -0,0 +1,409 @@ +# type: ignore +"""KV-cache disk persistence round-trips for exotic architectures. + +GLM-5.1 / DeepSeek-V3.2 (DSA) caches hold zero-width arrays and DeepSeek-V4 +caches hold Nones/ints/int-lists — none of which safetensors can store. These +tests cover the placeholder mechanism that preserves them across save/load, +plus backward compatibility with legacy (pre-placeholder) slots. + +No model weights needed: caches are constructed directly. +""" + +import hashlib +import json +import time + +import mlx.core as mx +import pytest +from mlx.utils import tree_flatten +from mlx_lm.models.cache import ( + ArraysCache, + CacheList, + KVCache, + load_prompt_cache, + save_prompt_cache, +) +from mlx_lm.models.deepseek_v4 import DeepseekV4Cache + +from exo.worker.engines.mlx.cache import ( + KVPrefixCache, + _is_int_list_leaf, +) + +MODEL_ID = "test/disk-model" +N_TOKENS = 1100 # disk load requires a >=1000-token prefix + + +@pytest.fixture +def kvp_factory(tmp_path, monkeypatch): + """Build KVPrefixCache instances that all share tmp_path as the disk dir.""" + monkeypatch.setenv("EXO_KV_DISK_PATH", str(tmp_path)) + monkeypatch.setenv("EXO_KV_DISK_PERSISTENCE", "1") + + def factory(): + return KVPrefixCache(None, model_id=MODEL_ID) + + return factory + + +def _slot_dir(tmp_path): + return tmp_path / hashlib.sha256(MODEL_ID.encode()).hexdigest()[:16] + + +def test_disk_persistence_is_opt_in(tmp_path, monkeypatch): + """Without EXO_KV_DISK_PERSISTENCE=1 nothing touches the disk.""" + monkeypatch.setenv("EXO_KV_DISK_PATH", str(tmp_path)) + monkeypatch.delenv("EXO_KV_DISK_PERSISTENCE", raising=False) + kv = KVPrefixCache(None, model_id=MODEL_ID) + assert kv._disk_dir is None + assert list(tmp_path.iterdir()) == [] + + +def test_janitor_sweeps_other_models_stale_slots(tmp_path, kvp_factory): + """Init must TTL-sweep OTHER models' dirs (their runners may never come + back to clean up), while leaving fresh slots and our own dir alone.""" + other = tmp_path / "deadbeef00112233" + other.mkdir() + + def write_slot(i, ts): + (other / f"slot_{i}_meta.json").write_text( + json.dumps({"model_id": "other/model", "token_count": 5, "timestamp": ts}) + ) + (other / f"slot_{i}_tokens.safetensors").write_text("stub") + (other / f"slot_{i}_cache.safetensors").write_text("stub") + + write_slot(0, 1.0) # ancient -> swept + write_slot(1, time.time()) # fresh -> kept + # A stale slot in OUR OWN dir is left to the hot-slot-aware eviction. + own = _slot_dir(tmp_path) + own.mkdir(parents=True) + (own / "slot_7_meta.json").write_text( + json.dumps({"model_id": MODEL_ID, "token_count": 5, "timestamp": 1.0}) + ) + + kvp_factory() # __init__ runs the janitor + + assert not (other / "slot_0_meta.json").exists() + assert not (other / "slot_0_cache.safetensors").exists() + assert not (other / "slot_0_tokens.safetensors").exists() + assert (other / "slot_1_meta.json").exists() + assert (own / "slot_7_meta.json").exists() + + +def test_should_update_entry_only_for_extensions(tmp_path, kvp_factory): + """Update-in-place is for continuations only. A sibling conversation + sharing a long prefix (two agent sessions with a common bootstrap) must + get its own entry — updating steals the matched entry's disk slot.""" + kv = kvp_factory() + kv.add_kv_cache(_tokens(2000), [_filled_kv(2000)]) + + # Continuation: hit covers the stored prompt (modulo prefill rollback) + assert kv.should_update_entry(0, 1998) is True + assert kv.should_update_entry(0, 2000) is True + # Sibling: long shared prefix but far short of the stored conversation + assert kv.should_update_entry(0, 1500) is False + # Below the minimum hit threshold + assert kv.should_update_entry(0, 500) is False + # No match / stale index + assert kv.should_update_entry(None, 1998) is False + assert kv.should_update_entry(5, 1998) is False + + +def test_sibling_conversation_gets_own_slot(tmp_path, kvp_factory): + """Slot-theft regression: saving a sibling conversation must not + overwrite the other conversation's disk slot (observed live as two + sessions ping-ponging slot_10, each switch destroying the other's + cache back to the shared prefix).""" + kv = kvp_factory() + conv_a = _tokens(2000) + kv.add_kv_cache(conv_a, [_filled_kv(2000)]) + kv._flush_hot_slot() # conv A -> slot_0 + + # Sibling: shares the first 1500 tokens, then diverges. + conv_b = mx.concatenate([conv_a[:1500], mx.array([9000 + i for i in range(600)])]) + assert kv.should_update_entry(0, 1500) is False # decision: add, not update + kv.add_kv_cache(conv_b, [_filled_kv(2100)]) + kv._flush_hot_slot() # conv B -> its own slot_1 + + d = _slot_dir(tmp_path) + tok0 = mx.load(str(d / "slot_0_tokens.safetensors"))["tokens"] + tok1 = mx.load(str(d / "slot_1_tokens.safetensors"))["tokens"] + assert mx.array_equal(tok0, conv_a).item() # A's slot intact + assert mx.array_equal(tok1, conv_b).item() + + +def test_per_flush_eviction_also_ttl_sweeps_other_models(tmp_path, kvp_factory): + """The TTL must reach other models' dirs from the per-flush path too — + not only at init — so one model kept loaded for weeks still cleans up.""" + kv = kvp_factory() # init janitor runs before the fixtures exist + + other = tmp_path / "cafebabe00112233" + other.mkdir() + (other / "slot_0_meta.json").write_text( + json.dumps({"model_id": "other/model", "token_count": 5, "timestamp": 1.0}) + ) + (other / "slot_0_tokens.safetensors").write_text("stub") + (other / "slot_0_cache.safetensors").write_text("stub") + + kv._evict_stale_disk_slots() # what flush_to_disk calls after each flush + + assert not (other / "slot_0_meta.json").exists() + assert not (other / "slot_0_cache.safetensors").exists() + + +def test_size_cap_is_global_and_evicts_oldest_first(tmp_path, kvp_factory, monkeypatch): + """The size cap spans ALL model dirs: the globally oldest slot is evicted + first (cross-model LRU), eviction stops once back under the cap, and the + flushing model's own hot slot is never touched.""" + kv = kvp_factory() # created first: the init janitor must not see fixtures + kv.add_kv_cache(_tokens(), [_filled_kv()]) + kv._flush_hot_slot() # own hot slot + + other = tmp_path / "feedface00112233" + other.mkdir() + + def write_slot(i, ts, payload): + (other / f"slot_{i}_meta.json").write_text( + json.dumps({"model_id": "other/model", "token_count": 5, "timestamp": ts}) + ) + (other / f"slot_{i}_tokens.safetensors").write_text(payload) + (other / f"slot_{i}_cache.safetensors").write_text(payload) + + write_slot(0, 2.0, "x" * 500) # globally oldest, and big + write_slot(1, time.time(), "y" * 10) # fresh and small + + # Cap sits between (old + fresh) and (fresh alone); the hot slot is + # excluded from the accounting, so only the foreign slots count. + monkeypatch.setenv("EXO_KV_DISK_MAX_SIZE_GB", str(600 / 1024**3)) + kv._evict_stale_disk_slots() + + assert not (other / "slot_0_meta.json").exists() # oldest evicted + assert (other / "slot_1_meta.json").exists() # under cap again -> kept + assert (_slot_dir(tmp_path) / "slot_0_cache.safetensors").exists() # hot slot safe + + +def _filled_kv(n=N_TOKENS, heads=2, dim=8): + c = KVCache() + c.update_and_fetch( + mx.random.normal([1, heads, n, dim]), mx.random.normal([1, heads, n, dim]) + ) + return c + + +def _tokens(n=N_TOKENS): + return mx.arange(n) + + +def _flush(kvp, cache, tokens): + kvp.add_kv_cache(tokens, cache) + kvp._flush_hot_slot() + + +def _states_equal(a_cache, b_cache): + a = tree_flatten([c.state for c in a_cache], is_leaf=_is_int_list_leaf) + b = tree_flatten([c.state for c in b_cache], is_leaf=_is_int_list_leaf) + assert [p for p, _ in a] == [p for p, _ in b], "state tree paths differ" + for (path, la), (_, lb) in zip(a, b, strict=True): + if isinstance(la, mx.array): + assert isinstance(lb, mx.array), f"{path}: array vs {type(lb)}" + assert la.shape == lb.shape, f"{path}: shape {la.shape} != {lb.shape}" + assert la.dtype == lb.dtype, f"{path}: dtype {la.dtype} != {lb.dtype}" + assert mx.array_equal(la, lb).item(), f"{path}: values differ" + else: + assert la == lb, f"{path}: {la!r} != {lb!r}" + + +def test_standard_kvcache_roundtrip(tmp_path, kvp_factory): + """Llama/Kimi regression: plain KVCache slots survive, exact and partial.""" + cache = [_filled_kv(), _filled_kv()] + tokens = _tokens() + _flush(kvp_factory(), cache, tokens) + + # Exact prefix + new suffix + query = mx.concatenate([tokens, mx.array([7000, 7001])]) + result = kvp_factory()._try_load_from_disk(None, query) + assert result is not None + loaded, remaining, _, _ = result + _states_equal(cache, loaded) + assert all(c.offset == N_TOKENS for c in loaded) + assert mx.array_equal(remaining, query[N_TOKENS:]).item() + + # Partial prefix: diverges at 1050 -> trim to 1050 + query2 = mx.concatenate([tokens[:1050], mx.array([9999] * 5)]) + result2 = kvp_factory()._try_load_from_disk(None, query2) + assert result2 is not None + loaded2, remaining2, _, _ = result2 + assert all(c.offset == 1050 for c in loaded2) + assert mx.array_equal(remaining2, query2[1050:]).item() + + +def test_legacy_format_slot_loads(tmp_path, kvp_factory): + """Slots in stock mlx-lm format (no placeholders key in meta.json) load + through the same path — the placeholder mechanism is strictly additive.""" + cache = [_filled_kv(), _filled_kv()] + tokens = _tokens() + d = _slot_dir(tmp_path) + d.mkdir(parents=True) + save_prompt_cache(str(d / "slot_0_cache.safetensors"), cache) + mx.save_safetensors(str(d / "slot_0_tokens.safetensors"), {"tokens": tokens}) + (d / "slot_0_meta.json").write_text( + json.dumps({"model_id": MODEL_ID, "token_count": N_TOKENS, "timestamp": 1.0}) + ) + + result = kvp_factory()._try_load_from_disk( + None, mx.concatenate([tokens, mx.array([7000])]) + ) + assert result is not None + _states_equal(cache, result[0]) + + +def test_cache_shorter_than_tokens_loads(tmp_path, kvp_factory): + """Real slots store the cache ~2 tokens shorter than the token file + (prefill rollback). The reusable prefix must cap at the cache length — + a strict equality check would reject every production slot.""" + cache = [_filled_kv(n=N_TOKENS - 2)] + tokens = _tokens() # 1100 tokens, cache holds 1098 + _flush(kvp_factory(), cache, tokens) + + query = mx.concatenate([tokens, mx.array([7000])]) + result = kvp_factory()._try_load_from_disk(None, query) + assert result is not None + loaded, remaining, _, _ = result + assert all(c.offset == N_TOKENS - 2 for c in loaded) + assert mx.array_equal(remaining, query[N_TOKENS - 2 :]).item() + + +def test_new_save_readable_by_stock_mlx_lm(tmp_path, kvp_factory): + """For standard archs the new save stays byte-compatible with mlx-lm.""" + cache = [_filled_kv()] + _flush(kvp_factory(), cache, _tokens()) + loaded = load_prompt_cache(str(_slot_dir(tmp_path) / "slot_0_cache.safetensors")) + _states_equal(cache, loaded) + + +def _glm_style_cache(n_layers=2): + """CacheList(main KVCache, indexer KVCache) per layer; the indexer stores a + zero-width values array, mirroring deepseek_v32.py's Indexer.""" + layers = [] + for _ in range(n_layers): + main = _filled_kv() + indexer = KVCache() + indexer.update_and_fetch( + mx.random.normal([1, 1, N_TOKENS, 8]), mx.zeros([1, 1, N_TOKENS, 0]) + ) + layers.append(CacheList(main, indexer)) + return layers + + +def test_glm_cachelist_zero_width_roundtrip(tmp_path, kvp_factory): + cache = _glm_style_cache() + tokens = _tokens() + _flush(kvp_factory(), cache, tokens) + + meta = json.loads((_slot_dir(tmp_path) / "slot_0_meta.json").read_text()) + assert meta["format"] == 2 + empties = [s for s in meta["placeholders"].values() if s["kind"] == "empty"] + assert len(empties) == 2 # one zero-width values array per layer + assert empties[0]["shape"] == [1, 1, N_TOKENS, 0] + + # Exact prefix: previously failed with "not enough values to unpack" + result = kvp_factory()._try_load_from_disk( + None, mx.concatenate([tokens, mx.array([7000])]) + ) + assert result is not None + loaded = result[0] + _states_equal(cache, loaded) + assert loaded[0][1].values.shape == (1, 1, N_TOKENS, 0) + assert all(cl[0].offset == N_TOKENS and cl[1].offset == N_TOKENS for cl in loaded) + + # Partial prefix: CacheList of KVCaches IS trimmable + query = mx.concatenate([tokens[:1050], mx.array([9999] * 5)]) + result2 = kvp_factory()._try_load_from_disk(None, query) + assert result2 is not None + loaded2, remaining2, _, _ = result2 + assert all(cl[0].offset == 1050 and cl[1].offset == 1050 for cl in loaded2) + assert mx.array_equal(remaining2, query[1050:]).item() + + +def _v4_cache(sliding_window=64): + """DeepseekV4Cache with a rotated local window and mixed branch state + (arrays, Nones, int-lists, ints) — previously crashed the flush.""" + c = DeepseekV4Cache(sliding_window) + c.local.update_and_fetch( + mx.random.normal([1, 1, N_TOKENS, 8]), mx.random.normal([1, 1, N_TOKENS, 8]) + ) + comp = c._branches["compressor"] + comp.pool = mx.random.normal([1, 17, 8]) + comp.buffer_kv = mx.random.normal([1, 4, 16]) + comp.pool_lengths = [17] + comp.buffer_count = 3 + idx = c._branches["indexer"] + idx.pool = mx.random.normal([1, 17, 8]) + idx.buffer_lengths = [] + return c + + +def test_v4_cache_roundtrip(tmp_path, kvp_factory): + cache = [_v4_cache(), _v4_cache()] + tokens = _tokens() + _flush(kvp_factory(), cache, tokens) + + # The flush must actually have produced a slot (old code always failed + # with "'NoneType' object has no attribute 'size'"). + assert (_slot_dir(tmp_path) / "slot_0_cache.safetensors").exists() + meta = json.loads((_slot_dir(tmp_path) / "slot_0_meta.json").read_text()) + kinds = {s["kind"] for s in meta["placeholders"].values()} + assert {"none", "json"} <= kinds + + # Exact-prefix continuation + result = kvp_factory()._try_load_from_disk( + None, mx.concatenate([tokens, mx.array([7000, 7001])]) + ) + assert result is not None + loaded, remaining, _, _ = result + _states_equal(cache, loaded) + for orig, new in zip(cache, loaded, strict=True): + assert new.offset == orig.offset == N_TOKENS + assert new.local.max_size == orig.local.max_size + assert new.local._idx == orig.local._idx + assert new._branches["compressor"].pool_lengths == [17] + assert new._branches["compressor"].buffer_count == 3 + assert new._branches["compressor"].prev_kv is None + assert new._branches["indexer"].buffer_lengths == [] + assert mx.array_equal(remaining, mx.array([7000, 7001])).item() + + from copy import deepcopy + + deepcopy(loaded) # _try_load_from_disk deepcopies — must not raise + + +def test_v4_partial_prefix_guard(tmp_path, kvp_factory): + """A partial-prefix hit on a non-trimmable cache must be skipped safely + (previously: AttributeError on the read-only offset, outside the try).""" + tokens = _tokens() + _flush(kvp_factory(), [_v4_cache()], tokens) + + kv = kvp_factory() + query = mx.concatenate([tokens[:1050], mx.array([9999] * 10)]) + assert kv._try_load_from_disk(None, query) is None + assert len(kv.caches) == 0 # in-memory state untouched + + +def test_arrays_cache_none_roundtrip(tmp_path, kvp_factory): + """SSM-style ArraysCache with None entries round-trips via placeholders + (hybrid model: one KVCache layer + one ArraysCache layer).""" + kv_layer = _filled_kv() + ssm = ArraysCache(2) + ssm[0] = mx.random.normal([1, 4]) + assert ssm.cache[1] is None + cache = [kv_layer, ssm] + tokens = _tokens() + _flush(kvp_factory(), cache, tokens) + + result = kvp_factory()._try_load_from_disk( + None, mx.concatenate([tokens, mx.array([7000])]) + ) + assert result is not None + loaded = result[0] + _states_equal(cache, loaded) + assert loaded[1].cache[1] is None diff --git a/src/exo/worker/tests/unittests/test_runner/test_kv_flush_wiring.py b/src/exo/worker/tests/unittests/test_runner/test_kv_flush_wiring.py new file mode 100644 index 0000000000..a0e1391935 --- /dev/null +++ b/src/exo/worker/tests/unittests/test_runner/test_kv_flush_wiring.py @@ -0,0 +1,25 @@ +# type: ignore +"""Runner wiring: the idle/shutdown KV flush reaches the engine's prefix cache.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from exo.worker.runner.runner import Runner + + +def test_flush_calls_engine_prefix_cache(): + kv = MagicMock() + fake_self = SimpleNamespace(generator=SimpleNamespace(kv_prefix_cache=kv)) + Runner._flush_kv_cache_to_disk(fake_self) + kv.flush_to_disk.assert_called_once_with(force=True) + + +def test_flush_noop_for_engine_without_cache(): + # Builder before build(), or engines without a prefix cache (e.g. image). + fake_self = SimpleNamespace(generator=object()) + Runner._flush_kv_cache_to_disk(fake_self) # must not raise + + +def test_flush_noop_when_cache_is_none(): + fake_self = SimpleNamespace(generator=SimpleNamespace(kv_prefix_cache=None)) + Runner._flush_kv_cache_to_disk(fake_self) # must not raise