Revert "Fix kv prefix cache (#1262)"

This reverts commit cd8c01b7c8.

README.md

@@ -5,7 +5,7 @@
   <img alt="exo logo" src="/docs/imgs/exo-logo-transparent.png" width="50%" height="50%">
 </picture>
 
-exo: Run frontier AI locally. Maintained by [exo labs](https://x.com/exolabs).
+exo: Run your own AI cluster at home with everyday devices. Maintained by [exo labs](https://x.com/exolabs).
 
 <p align="center">
   <a href="https://discord.gg/TJ4P57arEm" target="_blank" rel="noopener noreferrer"><img src="https://img.shields.io/badge/Discord-Join%20Server-5865F2?logo=discord&logoColor=white" alt="Discord"></a>

app/EXO/EXO/Services/NetworkSetupHelper.swift

@@ -18,6 +18,9 @@ enum NetworkSetupHelper {
 
         set -euo pipefail
 
+        # Wait for macOS to finish network setup after boot
+        sleep 30
+
         PREFS="/Library/Preferences/SystemConfiguration/preferences.plist"
 
         # Remove bridge0 interface
@@ -80,7 +83,7 @@ enum NetworkSetupHelper {
                 let alert = NSAlert()
                 alert.messageText = "EXO Network Configuration"
                 alert.informativeText =
-                    "EXO needs to install a system service to automatically disable Thunderbolt Bridge on startup. This prevents network loops when connecting multiple Macs via Thunderbolt.\n\nYou will be prompted for your administrator password."
+                    "EXO needs to install a system service to configure local networking. This will disable Thunderbolt Bridge (preventing packet storms) and install a Network Location.\n\nYou will be prompted for your password."
                 alert.alertStyle = .informational
                 alert.addButton(withTitle: "Install")
                 alert.addButton(withTitle: "Not Now")

justfile

@@ -1,7 +1,7 @@
 export NIX_CONFIG := "extra-experimental-features = nix-command flakes"
 
 fmt:
-    treefmt || nix fmt
+    nix fmt
 
 lint:
     uv run ruff check --fix

pyproject.toml

@@ -17,9 +17,9 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx==0.30.4; sys_platform == 'darwin'",
-    "mlx[cpu]==0.30.4; sys_platform == 'linux'",
-    "mlx-lm",
+    "mlx @ git+https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git; sys_platform == 'darwin'",
+    "mlx[cpu]==0.30.3; sys_platform == 'linux'",
+    "mlx-lm==0.30.5",
     "tiktoken>=0.12.0", # required for kimi k2 tokenizer
     "hypercorn>=0.18.0",
     "openai-harmony>=0.0.8",
@@ -63,7 +63,6 @@ members = [
 
 [tool.uv.sources]
 exo_pyo3_bindings = { workspace = true }
-mlx-lm = { git = "https://github.com/ml-explore/mlx-lm", branch = "main" }
 # Uncomment to use local mlx/mlx-lm development versions:
 # mlx = { path = "/Users/Shared/mlx", editable=true }
 # mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }

src/exo/main.py

@@ -90,6 +90,7 @@ class Node:
             worker = Worker(
                 node_id,
                 session_id,
+                connection_message_receiver=router.receiver(topics.CONNECTION_MESSAGES),
                 global_event_receiver=router.receiver(topics.GLOBAL_EVENTS),
                 local_event_sender=router.sender(topics.LOCAL_EVENTS),
                 command_sender=router.sender(topics.COMMANDS),
@@ -226,6 +227,9 @@ class Node:
                         self.worker = Worker(
                             self.node_id,
                             result.session_id,
+                            connection_message_receiver=self.router.receiver(
+                                topics.CONNECTION_MESSAGES
+                            ),
                             global_event_receiver=self.router.receiver(
                                 topics.GLOBAL_EVENTS
                             ),

src/exo/shared/models/model_cards.py

@@ -121,14 +121,6 @@ MODEL_CARDS: dict[str, ModelCard] = {
         supports_tensor=True,
         tasks=[ModelTask.TextGeneration],
     ),
-    "kimi-k2.5": ModelCard(
-        model_id=ModelId("mlx-community/Kimi-K2.5"),
-        storage_size=Memory.from_gb(617),
-        n_layers=61,
-        hidden_size=7168,
-        supports_tensor=True,
-        tasks=[ModelTask.TextGeneration],
-    ),
     # llama-3.1
     "llama-3.1-8b": ModelCard(
         model_id=ModelId("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit"),

src/exo/shared/types/mlx.py

@@ -1,12 +0,0 @@
-"""Shared types for MLX-related functionality."""
-
-from collections.abc import Sequence
-
-from mlx_lm.models.cache import (
-    KVCache,
-    QuantizedKVCache,
-    RotatingKVCache,
-)
-
-# This list contains one cache entry per transformer layer
-KVCacheType = Sequence[KVCache | RotatingKVCache | QuantizedKVCache]

src/exo/worker/engines/mlx/auto_parallel.py

@@ -23,7 +23,6 @@ from mlx_lm.models.glm4_moe_lite import Glm4MoeLiteDecoderLayer, Glm4MoeLiteMLP
 from mlx_lm.models.glm4_moe_lite import Model as GLM4MoeLiteModel
 from mlx_lm.models.gpt_oss import GptOssMoeModel
 from mlx_lm.models.gpt_oss import Model as GptOssModel
-from mlx_lm.models.kimi_k25 import Model as KimiK25Model
 from mlx_lm.models.llama import Model as LlamaModel
 from mlx_lm.models.minimax import Model as MiniMaxModel
 from mlx_lm.models.ministral3 import Model as Ministral3Model
@@ -345,7 +344,7 @@ def tensor_auto_parallel(
             all_to_sharded_linear_in_place,
             sharded_to_all_linear_in_place,
         )
-    elif isinstance(model, (DeepseekV3Model, DeepseekV32Model, KimiK25Model)):
+    elif isinstance(model, (DeepseekV3Model, DeepseekV32Model)):
         tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
             group,
             all_to_sharded_linear,
@@ -454,7 +453,7 @@ def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
 
         # Update DeepSeek V3 specific parameters when layers are shrunk
         if isinstance(
-            model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel, KimiK25Model)
+            model, (DeepseekV3Model, DeepseekV32Model, Glm4MoeModel)
         ) and hasattr(inner_model_instance, "num_layers"):
             logger.info(
                 f"Setting num_layers to {len(layers)} for model {model.model.__class__.__name__}"
@@ -623,7 +622,6 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
         on_timeout: TimeoutCallback | None,
     ) -> nn.Module:
         model = cast(MiniMaxModel, model)
-        rank = self.group.rank()
         for layer in model.layers:
             eval_with_timeout(
                 layer.parameters(), timeout_seconds / len(model.layers), on_timeout
@@ -633,16 +631,6 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
             layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
             layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
             layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
-
-            # Shard qk_norm weights if present (must match sharded head count)
-            if getattr(layer.self_attn, "use_qk_norm", False):
-                layer.self_attn.q_norm.weight = layer.self_attn.q_norm.weight.split(  # type: ignore
-                    self.N, axis=-1
-                )[rank]
-                layer.self_attn.k_norm.weight = layer.self_attn.k_norm.weight.split(  # type: ignore
-                    self.N, axis=-1
-                )[rank]
-
             layer.self_attn.num_attention_heads //= self.N
             layer.self_attn.num_key_value_heads //= self.N
 

src/exo/worker/engines/mlx/cache.py

@@ -1,81 +1,39 @@
-import os
+# type: ignore
+# TODO: Fix this file, including types!
 from copy import deepcopy
-from typing import Any, cast
+from typing import Callable
 
 import mlx.core as mx
-from mlx_lm.models.cache import (
-    KVCache,
-    QuantizedKVCache,
-    RotatingKVCache,
-    trim_prompt_cache,
-)
-from mlx_lm.models.gpt_oss import Model as GptOssModel
+from mlx_lm import stream_generate
+from mlx_lm.models.cache import _BaseCache, trim_prompt_cache
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
-from exo.shared.types.mlx import KVCacheType
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.constants import CACHE_GROUP_SIZE, KV_CACHE_BITS
-from exo.worker.runner.bootstrap import logger
-
-# Fraction of device memory above which LRU eviction kicks in
-_DEFAULT_MEMORY_THRESHOLD = 0.85
-_MEMORY_THRESHOLD = float(
-    os.environ.get("EXO_MEMORY_THRESHOLD", _DEFAULT_MEMORY_THRESHOLD)
-)
+from exo.worker.engines.mlx.constants import KEEP_KV_SIZE, KV_BITS, KV_GROUP_SIZE
+from exo.worker.engines.mlx.utils_mlx import make_kv_cache
 
 
 class KVPrefixCache:
-    def __init__(self, tokenizer: TokenizerWrapper):
+    def __init__(self):
+        # Only one prefix cache per runner.
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
-        self.caches: list[KVCacheType] = []
-        self._last_used: list[int] = []  # monotonic counter of last access per entry
-        self._access_counter: int = 0
-        self._tokenizer: TokenizerWrapper = tokenizer
+        self.caches: list[list[_BaseCache]] = []
 
-    def clear(self):
-        """Clear all cached prompts and caches."""
-        self.prompts.clear()
-        self.caches.clear()
-        self._last_used.clear()
-
-    def add_kv_cache(self, prompt: str, cache: KVCacheType):
-        """Add a new cache entry. Evicts LRU entries if memory is high."""
-        self._evict_if_needed()
-        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
+    def add_kv_cache(
+        self, tokenizer: TokenizerWrapper, prompt: str, cache: list[_BaseCache]
+    ):
+        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
         self.prompts.append(tokenized_prompt)
         self.caches.append(deepcopy(cache))
-        self._access_counter += 1
-        self._last_used.append(self._access_counter)
-        logger.info(f"KV cache added: {len(tokenized_prompt)} tokens")
-
-    def update_kv_cache(
-        self,
-        index: int,
-        prompt: str,
-        cache: KVCacheType,
-    ):
-        """Update an existing cache entry in-place."""
-        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
-        self.prompts[index] = tokenized_prompt
-        self.caches[index] = deepcopy(cache)
-        self._access_counter += 1
-        self._last_used[index] = self._access_counter
-        logger.info(f"KV cache updated (index {index}): {len(tokenized_prompt)} tokens")
 
     def get_kv_cache(
         self,
         model: Model,
+        tokenizer: TokenizerWrapper,
+        sampler: Callable[[mx.array], mx.array],
         prompt: str,
-    ) -> tuple[KVCacheType, mx.array, int | None]:
-        """Get KV cache for prompt, returning remaining tokens to prefill.
-
-        Returns:
-            Tuple of (cache, remaining_tokens, matched_index) where:
-            - cache: KV cache to use for generation
-            - remaining_tokens: tokens that still need prefilling
-            - matched_index: index of the matched entry (None if no match)
-        """
-        tokenized_prompt = encode_prompt(self._tokenizer, prompt)
+    ) -> list[_BaseCache]:
+        tokenized_prompt = self.encode_prompt(tokenizer, prompt)
         max_length = len(tokenized_prompt)
 
         best_snapshot_index, best_snapshot_length = None, 0
@@ -84,127 +42,63 @@ class KVPrefixCache:
             length = _get_prefix_length(tokenized_prompt, cached_prompt)
 
             if length == max_length:
-                # Exact match - cached prompt starts with our entire prompt
-                # Trim cache to prompt length - 1, return last token for stream_generate
-                prompt_cache = deepcopy(self.caches[i])
-                cached_length = _cache_length(self.caches[i])
-                tokens_to_trim = cached_length - (max_length - 1)
-                if tokens_to_trim > 0:
-                    trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
-                self._access_counter += 1
-                self._last_used[i] = self._access_counter
-                logger.info(f"KV cache exact match: {max_length} tokens (instant)")
-                return prompt_cache, tokenized_prompt[-1:], i
+                return self.caches[i]
 
             if length > best_snapshot_length:
                 best_snapshot_index, best_snapshot_length = i, length
 
         if best_snapshot_index is not None:
-            new_tokens = max_length - best_snapshot_length
-            logger.info(
-                f"KV cache prefix match: {best_snapshot_length}/{max_length} tokens "
-                f"(reusing {best_snapshot_length}, need to prefill {new_tokens})"
-            )
-
             prompt_cache = deepcopy(self.caches[best_snapshot_index])
-
-            # Trim removes tokens from the end, so we trim (cached_length - prefix_length) to keep the prefix
-            cached_length = _cache_length(self.caches[best_snapshot_index])
-            tokens_to_trim = cached_length - best_snapshot_length
-            if tokens_to_trim > 0:
-                trim_prompt_cache(cast(list[Any], prompt_cache), tokens_to_trim)
-
-            self._access_counter += 1
-            self._last_used[best_snapshot_index] = self._access_counter
-            remaining_tokens = tokenized_prompt[best_snapshot_length:]
-            return prompt_cache, remaining_tokens, best_snapshot_index
+            trim_prompt_cache(prompt_cache, max_length - best_snapshot_length)
+            tokenized_prompt = tokenized_prompt[best_snapshot_index:]
 
         else:
-            prompt_cache = make_kv_cache(model)
-            if len(self.prompts) == 0:
-                logger.info(f"KV cache empty, need to prefill {max_length} tokens")
-            else:
-                logger.info(
-                    f"KV cache no prefix match, need to prefill {max_length} tokens"
-                )
-
-            return prompt_cache, tokenized_prompt, None
-
-    def _evict_if_needed(self):
-        """Evict least recently used entries while memory pressure is high."""
-        if len(self.caches) == 0:
-            return
-
-        active: int = mx.metal.get_active_memory()
-        limit = int(mx.metal.device_info()["max_recommended_working_set_size"])
-        if active < limit * _MEMORY_THRESHOLD:
-            return
-
-        # Evict LRU entries until below threshold or only one entry left
-        while len(self.caches) > 0:
-            lru_index = self._last_used.index(min(self._last_used))
-            evicted_tokens = len(self.prompts[lru_index])
-            self.prompts.pop(lru_index)
-            self.caches.pop(lru_index)
-            self._last_used.pop(lru_index)
-            logger.info(
-                f"KV cache evicted LRU entry ({evicted_tokens} tokens) due to memory pressure"
+            prompt_cache = make_kv_cache(
+                model,
+                # max_kv_size=MAX_KV_SIZE,
+                # keep=KEEP_KV_SIZE
             )
 
-            active = mx.metal.get_active_memory()
-            if active < limit * _MEMORY_THRESHOLD:
-                break
+        prefill(model, tokenizer, sampler, tokenized_prompt, prompt_cache)
 
+        return prompt_cache
 
-def encode_prompt(tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
-    """Encode a prompt string to token array.
-
-    For chat-templated prompts (which have their own structure markers like
-    <|im_user|>, <|im_middle|>, etc.), we should NOT add BOS/EOS tokens as
-    that would corrupt the prompt structure.
-    """
-    # Chat templates define their own structure - don't add BOS/EOS
-    tokenized_prompt = tokenizer.encode(prompt, add_special_tokens=False)
-    return mx.array(tokenized_prompt)
-
-
-def _cache_length(cache: KVCacheType) -> int:
-    """Get the number of tokens in a KV cache."""
-    # Use .offset attribute which all cache types have (len() not implemented in older QuantizedKVCache)
-    return max(c.offset for c in cache)  # type: ignore
+    def encode_prompt(self, tokenizer: TokenizerWrapper, prompt: str) -> mx.array:
+        add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
+            tokenizer.bos_token
+        )
+        tokenized_prompt = tokenizer.encode(
+            prompt, add_special_tokens=add_special_tokens
+        )
+        return mx.array(tokenized_prompt)
 
 
 def _get_prefix_length(prompt: mx.array, cached_prompt: mx.array) -> int:
-    """Find the length of the common prefix between two token arrays."""
-    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]))
+    n = min(int(prompt.shape[0]), int(cached_prompt.shape[0]), KEEP_KV_SIZE)
     if n == 0:
         return 0
 
-    equal = mx.equal(prompt[:n], cached_prompt[:n]).astype(mx.int32)
+    equal = (prompt[:n] == cached_prompt[:n]).astype(mx.int32)
     prefix_mask = mx.cumprod(equal)  # stays 1 until first mismatch, then 0 forever
     return int(mx.sum(prefix_mask).item())
 
 
-def make_kv_cache(
-    model: Model, max_kv_size: int | None = None, keep: int = 0
-) -> KVCacheType:
-    assert hasattr(model, "layers")
-
-    # TODO: Do this for all models
-    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
-        logger.info("Using MLX LM's make cache")
-        return model.make_cache()  # type: ignore
-
-    if max_kv_size is None:
-        if KV_CACHE_BITS is None:
-            logger.info("Using default KV cache")
-            return [KVCache() for _ in model.layers]
-        else:
-            logger.info("Using quantized KV cache")
-            return [
-                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
-                for _ in model.layers
-            ]
-    else:
-        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
-        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]
+def prefill(
+    model: Model,
+    tokenizer: TokenizerWrapper,
+    sampler: Callable[[mx.array], mx.array],
+    prompt: mx.array,
+    cache: list[_BaseCache],
+) -> None:
+    for _ in stream_generate(
+        model=model,
+        tokenizer=tokenizer,
+        prompt=prompt,
+        max_tokens=0,
+        sampler=sampler,
+        prompt_cache=cache,
+        prefill_step_size=2048,
+        kv_group_size=KV_GROUP_SIZE,
+        kv_bits=KV_BITS,
+    ):
+        pass

src/exo/worker/engines/mlx/constants.py

@@ -4,7 +4,7 @@
 KV_GROUP_SIZE: int | None = 32
 KV_BITS: int | None = None
 ATTENTION_KV_BITS: int | None = 4
-MAX_TOKENS: int = 32168
+MAX_TOKENS: int = 8192
 MAX_KV_SIZE: int | None = 3200
 KEEP_KV_SIZE: int | None = 1600
 QUANTIZE_MODEL_MODE: str | None = "affine"

src/exo/worker/engines/mlx/generator/generate.py

@@ -1,12 +1,12 @@
-import time
 from typing import Any, Callable, Generator, cast, get_args
 
 import mlx.core as mx
 from mlx_lm.generate import stream_generate
-from mlx_lm.models.cache import trim_prompt_cache
+from mlx_lm.models.cache import KVCache
 from mlx_lm.sample_utils import make_sampler
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
+# from exo.engines.mlx.cache import KVPrefixCache
 from exo.shared.types.api import (
     BenchChatCompletionTaskParams,
     ChatCompletionMessage,
@@ -14,78 +14,35 @@ from exo.shared.types.api import (
     GenerationStats,
 )
 from exo.shared.types.memory import Memory
-from exo.shared.types.mlx import KVCacheType
 from exo.shared.types.tasks import ChatCompletionTaskParams
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
 )
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.cache import KVPrefixCache, encode_prompt, make_kv_cache
 from exo.worker.engines.mlx.constants import KV_BITS, KV_GROUP_SIZE, MAX_TOKENS
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
+    make_kv_cache,
     mx_barrier,
 )
 from exo.worker.runner.bootstrap import logger
 
 generation_stream = mx.new_stream(mx.default_device())
 
-_MIN_PREFIX_HIT_TO_UPDATE = 1000
 
-
-def prefill(
-    model: Model,
-    tokenizer: TokenizerWrapper,
-    sampler: Callable[[mx.array], mx.array],
-    prompt_tokens: mx.array,
-    cache: KVCacheType,
-) -> float:
-    """Prefill the KV cache with prompt tokens.
-
-    This runs the model over the prompt tokens to populate the cache,
-    then trims off the extra generated token.
-
-    Returns:
-        tokens_per_sec
-    """
-    num_tokens = len(prompt_tokens)
-    if num_tokens == 0:
-        return 0.0
-
-    logger.debug(f"Prefilling {num_tokens} tokens...")
-    start_time = time.perf_counter()
-
-    def progress_callback(processed: int, total: int) -> None:
-        elapsed = time.time() - start_time
-        tok_per_sec = processed / elapsed if elapsed > 0 else 0
-        logger.debug(
-            f"Prefill progress: {processed}/{total} tokens ({tok_per_sec:.1f} tok/s)"
-        )
-
-    # Use max_tokens=1 because max_tokens=0 does not work.
-    # We just throw away the generated token - we only care about filling the cache
-    for _ in stream_generate(
-        model=model,
-        tokenizer=tokenizer,
-        prompt=prompt_tokens,
-        max_tokens=1,
-        sampler=sampler,
-        prompt_cache=cache,
-        prefill_step_size=2048,
-        kv_group_size=KV_GROUP_SIZE,
-        kv_bits=KV_BITS,
-        prompt_progress_callback=progress_callback,
-    ):
-        break  # Stop after first iteration - cache is now filled
-    trim_prompt_cache(cast(list[Any], cache), 1)
-
-    elapsed = time.perf_counter() - start_time
-    tokens_per_sec = num_tokens / elapsed if elapsed > 0 else 0.0
-    logger.debug(
-        f"Prefill complete: {num_tokens} tokens in {elapsed:.2f}s "
-        f"({tokens_per_sec:.1f} tok/s)"
-    )
-    return tokens_per_sec
+def maybe_quantize_kv_cache(
+    prompt_cache: list[KVCache | Any],
+    quantized_kv_start: int,
+    kv_group_size: int,
+    kv_bits: int | None,
+) -> None:
+    if kv_bits is None:
+        return
+    for e, c in enumerate(prompt_cache):
+        if (
+            hasattr(c, "to_quantized") and c.offset >= quantized_kv_start  # type: ignore
+        ):
+            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
 
 
 def warmup_inference(
@@ -163,7 +120,6 @@ def mlx_generate(
     tokenizer: TokenizerWrapper,
     task: ChatCompletionTaskParams,
     prompt: str,
-    kv_prefix_cache: KVPrefixCache | None = None,
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
@@ -175,22 +131,7 @@ def mlx_generate(
     if task.seed is not None:
         mx.random.seed(task.seed)
 
-    # Do not use the prefix cache if we are trying to do benchmarks.
-    if is_bench:
-        kv_prefix_cache = None
-
-    # Use prefix cache if available, otherwise create fresh cache
-    prefix_hit_length = 0
-    matched_index: int | None = None
-    if kv_prefix_cache is None:
-        caches = make_kv_cache(model=model)
-        prompt_tokens = encode_prompt(tokenizer, prompt)
-    else:
-        caches, prompt_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, prompt
-        )
-        all_prompt_tokens = encode_prompt(tokenizer, prompt)
-        prefix_hit_length = len(all_prompt_tokens) - len(prompt_tokens)
+    caches = make_kv_cache(model=model)
 
     logits_processors: list[Callable[[mx.array, mx.array], mx.array]] = []
     if is_bench:
@@ -203,19 +144,11 @@ def mlx_generate(
         top_p=task.top_p if task.top_p is not None else 1.0,
     )
 
-    # Prefill cache with all tokens except the last one
-    prefill_tps = prefill(model, tokenizer, sampler, prompt_tokens[:-1], caches)
-
-    # stream_generate starts from the last token
-    last_token = prompt_tokens[-1:]
-
     max_tokens = task.max_tokens or MAX_TOKENS
-    generated_text_parts: list[str] = []
-    generation_start_time = time.perf_counter()
     for out in stream_generate(
         model=model,
         tokenizer=tokenizer,
-        prompt=last_token,
+        prompt=prompt,
         max_tokens=max_tokens,
         sampler=sampler,
         logits_processors=logits_processors,
@@ -225,13 +158,12 @@ def mlx_generate(
         kv_group_size=KV_GROUP_SIZE,
         kv_bits=KV_BITS,
     ):
-        generated_text_parts.append(out.text)
         logger.info(out.text)
 
         stats: GenerationStats | None = None
         if out.finish_reason is not None:
             stats = GenerationStats(
-                prompt_tps=float(prefill_tps or out.prompt_tps),
+                prompt_tps=float(out.prompt_tps),
                 generation_tps=float(out.generation_tps),
                 prompt_tokens=int(out.prompt_tokens),
                 generation_tokens=int(out.generation_tokens),
@@ -253,26 +185,6 @@ def mlx_generate(
         )
 
         if out.finish_reason is not None:
-            # Log generation stats
-            generation_elapsed = time.perf_counter() - generation_start_time
-            generated_tokens = len(generated_text_parts)
-            generation_tps = (
-                generated_tokens / generation_elapsed if generation_elapsed > 0 else 0.0
-            )
-            logger.debug(
-                f"Generation complete: prefill {prompt_tokens} tokens @ "
-                f"{prefill_tps:.1f} tok/s, generated {generated_tokens} tokens @ "
-                f"{generation_tps:.1f} tok/s"
-            )
-            if kv_prefix_cache is not None:
-                full_prompt = prompt + "".join(generated_text_parts)
-                if (
-                    matched_index is not None
-                    and prefix_hit_length >= _MIN_PREFIX_HIT_TO_UPDATE
-                ):
-                    kv_prefix_cache.update_kv_cache(matched_index, full_prompt, caches)
-                else:
-                    kv_prefix_cache.add_kv_cache(full_prompt, caches)
             break
 
         # TODO: Do we want an mx_barrier?

src/exo/worker/engines/mlx/utils_mlx.py

@@ -18,12 +18,15 @@ try:
 except ImportError:
     pass  # transformers < 5.0 or bytes_to_unicode not available
 
-from mlx_lm.models.cache import KVCache
+from mlx_lm.models.cache import KVCache, QuantizedKVCache, RotatingKVCache
 from mlx_lm.models.deepseek_v3 import DeepseekV3Model
+from mlx_lm.models.gpt_oss import Model as GptOssModel
 from mlx_lm.tokenizer_utils import TokenizerWrapper
 
 from exo.shared.models.model_cards import ModelId
 from exo.worker.engines.mlx.constants import (
+    CACHE_GROUP_SIZE,
+    KV_CACHE_BITS,
     TRUST_REMOTE_CODE,
 )
 
@@ -165,6 +168,7 @@ def mlx_distributed_init(
 
                 jaccl_coordinator = jaccl_coordinators[bound_instance.bound_node_id]
 
+                # TODO: update once upstream fixes
                 logger.info(
                     f"rank {rank} MLX_IBV_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
                 )
@@ -258,10 +262,10 @@ def shard_and_load(
 
     logger.info(f"Group size: {group.size()}, group rank: {group.rank()}")
 
-    # Estimate timeout based on model size (5x default for large queued workloads)
-    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "300"))
+    # Estimate timeout based on model size
+    base_timeout = float(os.environ.get("EXO_MODEL_LOAD_TIMEOUT", "60"))
     model_size_gb = get_weights_size(shard_metadata).in_bytes / (1024**3)
-    timeout_seconds = base_timeout + model_size_gb
+    timeout_seconds = base_timeout + model_size_gb / 5
     logger.info(
         f"Evaluating model parameters with timeout of {timeout_seconds:.0f}s "
         f"(model size: {model_size_gb:.1f}GB)"
@@ -338,35 +342,8 @@ def load_tokenizer_for_model_id(
 
     # Kimi uses a custom TikTokenTokenizer that transformers 5.x can't load via AutoTokenizer
     if "kimi-k2" in model_id_lower:
-        import importlib.util
-        import types
-
         sys.path.insert(0, str(model_path))
-
-        # Load tool_declaration_ts first (tokenization_kimi imports it with relative import)
-        tool_decl_path = model_path / "tool_declaration_ts.py"
-        if tool_decl_path.exists():
-            spec = importlib.util.spec_from_file_location(
-                "tool_declaration_ts", tool_decl_path
-            )
-            if spec and spec.loader:
-                tool_decl_module = importlib.util.module_from_spec(spec)
-                sys.modules["tool_declaration_ts"] = tool_decl_module
-                spec.loader.exec_module(tool_decl_module)
-
-        # Load tokenization_kimi with patched source (convert relative to absolute import)
-        tok_path = model_path / "tokenization_kimi.py"
-        source = tok_path.read_text()
-        source = source.replace("from .tool_declaration_ts", "from tool_declaration_ts")
-        spec = importlib.util.spec_from_file_location("tokenization_kimi", tok_path)
-        if spec:
-            tok_module = types.ModuleType("tokenization_kimi")
-            tok_module.__file__ = str(tok_path)
-            sys.modules["tokenization_kimi"] = tok_module
-            exec(compile(source, tok_path, "exec"), tok_module.__dict__)  # noqa: S102
-            TikTokenTokenizer = tok_module.TikTokenTokenizer  # type: ignore[attr-defined]  # noqa: N806
-        else:
-            from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
+        from tokenization_kimi import TikTokenTokenizer  # type: ignore[import-not-found]  # noqa: I001
 
         hf_tokenizer: Any = TikTokenTokenizer.from_pretrained(model_path)  # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType]
 
@@ -489,6 +466,31 @@ class NullKVCache(KVCache):
         raise NotImplementedError("We should not be setting a NullKVCache.")
 
 
+def make_kv_cache(
+    model: Model, max_kv_size: int | None = None, keep: int = 0
+) -> list[KVCache | RotatingKVCache | QuantizedKVCache]:
+    assert hasattr(model, "layers")
+
+    # TODO: Do this for all models
+    if hasattr(model, "make_cache") and isinstance(model, GptOssModel):
+        logger.info("Using MLX LM's make cache")
+        return model.make_cache()  # type: ignore
+
+    if max_kv_size is None:
+        if KV_CACHE_BITS is None:
+            logger.info("Using default KV cache")
+            return [KVCache() for _ in model.layers]
+        else:
+            logger.info("Using quantized KV cache")
+            return [
+                QuantizedKVCache(group_size=CACHE_GROUP_SIZE, bits=KV_CACHE_BITS)
+                for _ in model.layers
+            ]
+    else:
+        logger.info(f"Using rotating KV cache with {max_kv_size=} with {keep=}")
+        return [RotatingKVCache(max_size=max_kv_size, keep=keep) for _ in model.layers]
+
+
 def mlx_force_oom(size: int = 40000) -> None:
     """
     Force an Out-Of-Memory (OOM) error in MLX by performing large tensor operations.

src/exo/worker/main.py

@@ -7,6 +7,7 @@ from anyio import CancelScope, create_task_group, fail_after
 from anyio.abc import TaskGroup
 from loguru import logger
 
+from exo.routing.connection_message import ConnectionMessage, ConnectionMessageType
 from exo.shared.apply import apply
 from exo.shared.models.model_cards import ModelId
 from exo.shared.types.api import ImageEditsInternalParams
@@ -56,6 +57,7 @@ class Worker:
         node_id: NodeId,
         session_id: SessionId,
         *,
+        connection_message_receiver: Receiver[ConnectionMessage],
         global_event_receiver: Receiver[ForwarderEvent],
         local_event_sender: Sender[ForwarderEvent],
         # This is for requesting updates. It doesn't need to be a general command sender right now,
@@ -72,6 +74,7 @@ class Worker:
         self.event_index_counter = event_index_counter
         self.command_sender = command_sender
         self.download_command_sender = download_command_sender
+        self.connection_message_receiver = connection_message_receiver
         self.event_buffer = OrderedBuffer[Event]()
         self.out_for_delivery: dict[EventId, ForwarderEvent] = {}
 
@@ -102,6 +105,7 @@ class Worker:
             tg.start_soon(info_gatherer.run)
             tg.start_soon(self._forward_info, info_recv)
             tg.start_soon(self.plan_step)
+            tg.start_soon(self._connection_message_event_writer)
             tg.start_soon(self._resend_out_for_delivery)
             tg.start_soon(self._event_applier)
             tg.start_soon(self._forward_events)
@@ -275,6 +279,41 @@ class Worker:
         instance = self.state.instances[task.instance_id]
         return instance.shard_assignments.node_to_runner[self.node_id]
 
+    async def _connection_message_event_writer(self):
+        with self.connection_message_receiver as connection_messages:
+            async for msg in connection_messages:
+                await self.event_sender.send(
+                    self._convert_connection_message_to_event(msg)
+                )
+
+    def _convert_connection_message_to_event(self, msg: ConnectionMessage):
+        match msg.connection_type:
+            case ConnectionMessageType.Connected:
+                return TopologyEdgeCreated(
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
+                        ),
+                    ),
+                )
+
+            case ConnectionMessageType.Disconnected:
+                return TopologyEdgeDeleted(
+                    conn=Connection(
+                        source=self.node_id,
+                        sink=msg.node_id,
+                        edge=SocketConnection(
+                            sink_multiaddr=Multiaddr(
+                                address=f"/ip4/{msg.remote_ipv4}/tcp/{msg.remote_tcp_port}"
+                            ),
+                        ),
+                    ),
+                )
+
     async def _nack_request(self, since_idx: int) -> None:
         # We request all events after (and including) the missing index.
         # This function is started whenever we receive an event that is out of sequence.

src/exo/worker/runner/runner.py

@@ -70,7 +70,6 @@ from exo.worker.engines.image import (
     warmup_image_generator,
 )
 from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.cache import KVPrefixCache
 from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference
 from exo.worker.engines.mlx.utils_mlx import (
     apply_chat_template,
@@ -104,7 +103,6 @@ def main(
     model: Model | DistributedImageModel | None = None
     tokenizer = None
     group = None
-    kv_prefix_cache: KVPrefixCache | None = None
 
     current_status: RunnerStatus = RunnerIdle()
     logger.info("runner created")
@@ -163,8 +161,6 @@ def main(
                         logger.info(
                             f"model has_tool_calling={tokenizer.has_tool_calling}"
                         )
-                        kv_prefix_cache = KVPrefixCache(tokenizer)
-
                     elif (
                         ModelTask.TextToImage in shard_metadata.model_card.tasks
                         or ModelTask.ImageToImage in shard_metadata.model_card.tasks
@@ -174,6 +170,7 @@ def main(
                         raise ValueError(
                             f"Unknown model task(s): {shard_metadata.model_card.tasks}"
                         )
+
                     current_status = RunnerLoaded()
                     logger.info("runner loaded")
                 case StartWarmup() if isinstance(current_status, RunnerLoaded):
@@ -241,7 +238,6 @@ def main(
                             tokenizer=tokenizer,
                             task=task_params,
                             prompt=prompt,
-                            kv_prefix_cache=kv_prefix_cache,
                         )
 
                         # For other thinking models (GLM, etc.), check if we need to

src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py

@@ -1,545 +0,0 @@
-# type: ignore
-import time
-from typing import cast
-from unittest.mock import patch
-
-import mlx.core as mx
-import pytest
-from mlx_lm.models.cache import KVCache
-from mlx_lm.sample_utils import make_sampler
-
-from exo.shared.types.api import ChatCompletionMessage
-from exo.shared.types.common import ModelId
-from exo.shared.types.tasks import ChatCompletionTaskParams
-from exo.worker.engines.mlx import Model
-from exo.worker.engines.mlx.cache import (
-    KVPrefixCache,
-    _cache_length,
-    _get_prefix_length,
-    encode_prompt,
-    make_kv_cache,
-)
-from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill
-from exo.worker.engines.mlx.utils_mlx import apply_chat_template
-from exo.worker.tests.unittests.test_mlx.conftest import (
-    DEFAULT_GPT_OSS_CONFIG,
-    DEFAULT_GPT_OSS_MODEL_ID,
-)
-
-
-def _check_model_exists() -> bool:
-    return DEFAULT_GPT_OSS_CONFIG.model_path.exists()
-
-
-class TestGetPrefixLength:
-    def test_identical_arrays(self):
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3, 4, 5])
-        assert _get_prefix_length(a, b) == 5
-
-    def test_no_common_prefix(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([4, 5, 6])
-        assert _get_prefix_length(a, b) == 0
-
-    def test_partial_prefix(self):
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3, 7, 8])
-        assert _get_prefix_length(a, b) == 3
-
-    def test_prompt_longer_than_cached(self):
-        a = mx.array([1, 2, 3, 4, 5])
-        b = mx.array([1, 2, 3])
-        assert _get_prefix_length(a, b) == 3
-
-    def test_cached_longer_than_prompt(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([1, 2, 3, 4, 5])
-        assert _get_prefix_length(a, b) == 3
-
-    def test_single_token_match(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([1, 5, 6])
-        assert _get_prefix_length(a, b) == 1
-
-    def test_empty_prompt(self):
-        a = mx.array([]).astype(mx.int32)
-        b = mx.array([1, 2, 3])
-        assert _get_prefix_length(a, b) == 0
-
-    def test_empty_cached(self):
-        a = mx.array([1, 2, 3])
-        b = mx.array([]).astype(mx.int32)
-        assert _get_prefix_length(a, b) == 0
-
-    def test_both_empty(self):
-        a = mx.array([]).astype(mx.int32)
-        b = mx.array([]).astype(mx.int32)
-        assert _get_prefix_length(a, b) == 0
-
-
-class TestKVPrefix:
-    @pytest.fixture
-    def mock_tokenizer(self):
-        """Create a minimal mock tokenizer for tests that don't need real tokenization."""
-        from unittest.mock import MagicMock
-
-        tokenizer = MagicMock()
-        tokenizer.encode.return_value = [1, 2, 3]
-        return tokenizer
-
-    def test_starts_empty(self, mock_tokenizer):
-        cache = KVPrefixCache(mock_tokenizer)
-        assert len(cache.prompts) == 0
-        assert len(cache.caches) == 0
-
-    def test_clear_empties_cache(self, mock_tokenizer):
-        cache = KVPrefixCache(mock_tokenizer)
-        cache.prompts.append(mx.array([1, 2, 3]))
-        cache.caches.append([KVCache()])
-        cache.clear()
-        assert len(cache.prompts) == 0
-        assert len(cache.caches) == 0
-
-    def test_clear_on_empty_cache(self, mock_tokenizer):
-        cache = KVPrefixCache(mock_tokenizer)
-        cache.clear()
-        assert len(cache.prompts) == 0
-
-
-def _load_gpt_oss() -> tuple[Model, object]:
-    from mlx_lm.utils import load_model
-
-    from exo.worker.engines.mlx.utils_mlx import load_tokenizer_for_model_id
-
-    model_path = DEFAULT_GPT_OSS_CONFIG.model_path
-    model_id = ModelId(DEFAULT_GPT_OSS_MODEL_ID)
-
-    model, _ = load_model(model_path, lazy=False)
-    tokenizer = load_tokenizer_for_model_id(model_id, model_path)
-    return cast(Model, model), tokenizer
-
-
-@pytest.mark.slow
-@pytest.mark.skipif(
-    not _check_model_exists(),
-    reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}",
-)
-class TestKVPrefixCacheWithModel:
-    @pytest.fixture(scope="class")
-    def model_and_tokenizer(self):
-        model, tokenizer = _load_gpt_oss()
-        return model, tokenizer
-
-    def test_prefill_populates_cache(self, model_and_tokenizer):
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Hello!!")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        # Cache should now hold the prompt tokens
-        assert _cache_length(cache) == len(tokens)
-
-    def test_add_and_get_exact_match(self, model_and_tokenizer):
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Test exact")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        assert len(kv_prefix_cache.prompts) == 1
-        stored_length = _cache_length(kv_prefix_cache.caches[0])
-        assert stored_length > 0
-
-        # Retrieve with same prompt: exact match
-        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, prompt
-        )
-        assert matched_index == 0
-
-        # Exact match returns only last token
-        assert len(remaining_tokens) == 1
-        assert mx.array_equal(remaining_tokens, tokens[-1:])
-
-    def test_add_and_get_prefix_match(self, model_and_tokenizer):
-        """get_kv_cache with a longer prompt sharing prefix should return partial match."""
-        model, tokenizer = model_and_tokenizer
-
-        short_task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Hi")],
-            max_tokens=1,
-        )
-        short_prompt = apply_chat_template(tokenizer, short_task)
-        short_tokens = encode_prompt(tokenizer, short_prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), short_tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(short_prompt, cache)
-
-        # Query with longer prompt that shares the chat template prefix
-        long_task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[
-                ChatCompletionMessage(role="user", content="Hi there, how are you?")
-            ],
-            max_tokens=1,
-        )
-        long_prompt = apply_chat_template(tokenizer, long_task)
-        long_tokens = encode_prompt(tokenizer, long_prompt)
-
-        # The prompts share a prefix (chat template preamble + "Hi")
-        expected_prefix = _get_prefix_length(long_tokens, short_tokens)
-        assert expected_prefix > 0, (
-            "Prompts should share a prefix from the chat template"
-        )
-
-        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, long_prompt
-        )
-        assert matched_index == 0
-
-        # remaining_tokens should be the suffix after the shared prefix
-        assert len(remaining_tokens) == len(long_tokens) - expected_prefix
-        assert mx.array_equal(remaining_tokens, long_tokens[expected_prefix:])
-
-    def test_stored_cache_not_mutated_after_get_and_generation(
-        self, model_and_tokenizer
-    ):
-        """Getting a cache and then mutating it (as generation does) must not corrupt stored cache."""
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Mutation test")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        stored_length = _cache_length(kv_prefix_cache.caches[0])
-
-        # Get cache and mutate it (simulating what generation does)
-        result_cache, _, matched_index = kv_prefix_cache.get_kv_cache(model, prompt)
-        assert matched_index == 0
-
-        # Simulate generation: feed many additional tokens through the cache
-        head_dim = result_cache[0].keys.shape[-1]
-        num_heads = result_cache[0].keys.shape[1]
-        extra_keys = mx.random.normal((1, num_heads, 50, head_dim))
-        extra_values = mx.random.normal((1, num_heads, 50, head_dim))
-        for layer_cache in result_cache:
-            layer_cache.update_and_fetch(extra_keys, extra_values)
-        mx.eval([c.keys for c in result_cache])
-
-        # Stored cache must be unchanged
-        assert _cache_length(kv_prefix_cache.caches[0]) == stored_length
-
-    def test_stored_cache_survives_repeated_get_mutate_cycles(
-        self, model_and_tokenizer
-    ):
-        """Multiple get+mutate cycles (like repeated user requests) must not corrupt cache."""
-        model, tokenizer = model_and_tokenizer
-
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Repeat test")],
-            max_tokens=1,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        tokens = encode_prompt(tokenizer, prompt)
-        cache = make_kv_cache(model)
-
-        prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        stored_length = _cache_length(kv_prefix_cache.caches[0])
-
-        for i in range(3):
-            result_cache, _, _ = kv_prefix_cache.get_kv_cache(model, prompt)
-
-            head_dim = result_cache[0].keys.shape[-1]
-            num_heads = result_cache[0].keys.shape[1]
-            extra = mx.random.normal((1, num_heads, 30, head_dim))
-            for layer_cache in result_cache:
-                layer_cache.update_and_fetch(extra, extra)
-            mx.eval([c.keys for c in result_cache])
-
-            assert _cache_length(kv_prefix_cache.caches[0]) == stored_length, (
-                f"Failed on loop {i}"
-            )
-
-    def test_mlx_generate_populates_cache(self, model_and_tokenizer):
-        """mlx_generate should save the cache after generation completes."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Hello")],
-            max_tokens=5,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        prompt_tokens = encode_prompt(tokenizer, prompt)
-
-        # Consume the entire generator so the cache-saving code after yield runs
-        generated_tokens = 0
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            generated_tokens += 1
-
-        assert len(kv_prefix_cache.prompts) == 1
-        assert len(kv_prefix_cache.caches) == 1
-        # Cache should contain prompt + generated tokens
-        expected_length = len(prompt_tokens) + generated_tokens
-        assert _cache_length(kv_prefix_cache.caches[0]) == expected_length
-
-    def test_mlx_generate_second_call_gets_prefix_hit(self, model_and_tokenizer):
-        """Second mlx_generate call with same prompt should get a prefix hit from stored cache."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Reuse test")],
-            max_tokens=5,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-        prompt_tokens = encode_prompt(tokenizer, prompt)
-
-        # First generation populates cache
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-
-        assert len(kv_prefix_cache.prompts) == 1
-
-        # Second call should find a prefix match (the stored cache contains
-        # prompt + generated tokens, which shares the prompt prefix)
-        result_cache, remaining_tokens, matched_index = kv_prefix_cache.get_kv_cache(
-            model, prompt
-        )
-        # The stored cache is longer than the prompt (it includes generated tokens),
-        # so this is a prefix match where our prompt is fully contained
-        assert matched_index == 0
-        # Exact match: remaining_tokens is just the last token
-        assert len(remaining_tokens) == 1
-        assert mx.array_equal(remaining_tokens, prompt_tokens[-1:])
-
-    def test_mlx_generate_long_prompt_updates_cache_in_place(self, model_and_tokenizer):
-        """With a prompt > 1000 tokens, second generation should update the cache entry in-place."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-
-        # Build a long user message (> 1000 tokens) to exceed _MIN_PREFIX_HIT_TO_UPDATE
-        base_text = "The quick brown fox jumps over the lazy dog. "
-        base_tokens = tokenizer.encode(base_text)
-        repeats = (1200 // len(base_tokens)) + 2
-        long_content = base_text * repeats
-
-        task1 = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content=long_content)],
-            max_tokens=5,
-        )
-        prompt1 = apply_chat_template(tokenizer, task1)
-        prompt1_tokens = encode_prompt(tokenizer, prompt1)
-        assert len(prompt1_tokens) > 1000, (
-            "Prompt must exceed _MIN_PREFIX_HIT_TO_UPDATE"
-        )
-
-        # First generation populates the cache (must prefill all tokens)
-        t0 = time.perf_counter()
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task1,
-            prompt=prompt1,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-        first_gen_time = time.perf_counter() - t0
-
-        assert len(kv_prefix_cache.prompts) == 1
-        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
-
-        # Second generation: same long prompt + extra content (simulating multi-turn)
-        task2 = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[
-                ChatCompletionMessage(role="user", content=long_content),
-                ChatCompletionMessage(role="assistant", content="Sure, I can help."),
-                ChatCompletionMessage(role="user", content="Tell me more."),
-            ],
-            max_tokens=5,
-        )
-        prompt2 = apply_chat_template(tokenizer, task2)
-        prompt2_tokens = encode_prompt(tokenizer, prompt2)
-
-        # Verify the prompts share a long prefix
-        prefix_len = _get_prefix_length(prompt2_tokens, prompt1_tokens)
-        assert prefix_len > 1000, "Prompts must share > 1000 token prefix"
-
-        # Second generation should reuse the cached prefix (only prefill new tokens)
-        t0 = time.perf_counter()
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task2,
-            prompt=prompt2,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-        second_gen_time = time.perf_counter() - t0
-
-        # Second generation should be significantly faster due to prefix cache hit - hopefully not flaky
-        assert second_gen_time < first_gen_time * 0.5, (
-            f"Expected prefix cache speedup: "
-            f"first={first_gen_time:.2f}s, second={second_gen_time:.2f}s"
-        )
-
-        # With prefix_hit > 1000, should update in-place (not add a second entry)
-        assert len(kv_prefix_cache.prompts) == 1
-        # Updated cache should be longer (prompt2 + generated > prompt1 + generated)
-        updated_cache_length = _cache_length(kv_prefix_cache.caches[0])
-        assert updated_cache_length > first_cache_length
-
-    def test_mlx_generate_stored_cache_not_mutated(self, model_and_tokenizer):
-        """After mlx_generate saves a cache, a second generation must not corrupt the stored copy."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-        task = ChatCompletionTaskParams(
-            model=DEFAULT_GPT_OSS_MODEL_ID,
-            messages=[ChatCompletionMessage(role="user", content="Immutable test")],
-            max_tokens=5,
-        )
-        prompt = apply_chat_template(tokenizer, task)
-
-        # First generation populates cache
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-
-        first_cache_length = _cache_length(kv_prefix_cache.caches[0])
-
-        # Second generation gets the cache and mutates it during generation
-        for _response in mlx_generate(
-            model=model,
-            tokenizer=tokenizer,
-            task=task,
-            prompt=prompt,
-            kv_prefix_cache=kv_prefix_cache,
-        ):
-            pass
-
-        # The first stored cache must not have been mutated by the second generation
-        assert _cache_length(kv_prefix_cache.caches[0]) == first_cache_length
-
-    def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer):
-        """Under memory pressure, adding a new cache entry evicts the least recently used one."""
-        model, tokenizer = model_and_tokenizer
-
-        kv_prefix_cache = KVPrefixCache(tokenizer)
-
-        # Add three cache entries with different prompts
-        prompts = ["First entry", "Second entry", "Third entry"]
-        for i, content in enumerate(prompts):
-            task = ChatCompletionTaskParams(
-                model=DEFAULT_GPT_OSS_MODEL_ID,
-                messages=[ChatCompletionMessage(role="user", content=content)],
-                max_tokens=1,
-            )
-            prompt = apply_chat_template(tokenizer, task)
-            tokens = encode_prompt(tokenizer, prompt)
-            cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-            kv_prefix_cache.add_kv_cache(prompt, cache)
-            # Stagger _last_used so LRU order is deterministic
-            kv_prefix_cache._last_used[i] = float(i)
-
-        assert len(kv_prefix_cache.prompts) == 3
-
-        # Access the third entry to make it most recently used
-        kv_prefix_cache._last_used[2] = 100.0
-        # Entry 0 (_last_used=0.0) is LRU, entry 1 (_last_used=1.0) is next
-
-        # Simulate memory pressure: active memory exceeds threshold
-        fake_limit = 1000
-        fake_active = int(fake_limit * 0.90)  # Above _MEMORY_THRESHOLD (0.85)
-
-        with (
-            patch(
-                "exo.worker.engines.mlx.cache.mx.metal.get_active_memory",
-                return_value=fake_active,
-            ),
-            patch(
-                "exo.worker.engines.mlx.cache.mx.metal.device_info",
-                return_value={"max_recommended_working_set_size": fake_limit},
-            ),
-        ):
-            # Trigger eviction by adding a new entry
-            task = ChatCompletionTaskParams(
-                model=DEFAULT_GPT_OSS_MODEL_ID,
-                messages=[ChatCompletionMessage(role="user", content="New entry")],
-                max_tokens=1,
-            )
-            prompt = apply_chat_template(tokenizer, task)
-            tokens = encode_prompt(tokenizer, prompt)
-            cache = make_kv_cache(model)
-            prefill(model, tokenizer, make_sampler(0.0), tokens, cache)
-            kv_prefix_cache.add_kv_cache(prompt, cache)
-
-        # LRU entries should have been evicted (entries 0, 1, 2 in order of _last_used)
-        # Since fake_active stays above threshold after each eviction (we don't change it),
-        # all old entries get evicted, leaving only the newly added one
-        assert len(kv_prefix_cache.prompts) == 1
-        # The surviving entry should be the newly added one
-        new_tokens = encode_prompt(tokenizer, prompt)
-        assert _get_prefix_length(kv_prefix_cache.prompts[0], new_tokens) == len(
-            new_tokens
-        )