Revert "Fix kv prefix cache (#1262)"

This reverts commit cd8c01b7c8.

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,7 +17,7 @@ dependencies = [
     "loguru>=0.7.3",
     "exo_pyo3_bindings", # rust bindings
     "anyio==4.11.0",
-    "mlx==0.30.3; sys_platform == 'darwin'",
+    "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

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

@@ -622,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
@@ -632,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,
 )
 
@@ -167,10 +170,10 @@ def mlx_distributed_init(
 
                 # TODO: update once upstream fixes
                 logger.info(
-                    f"rank {rank} MLX_JACCL_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
+                    f"rank {rank} MLX_IBV_DEVICES: {coordination_file} with devices: {jaccl_devices_json}"
                 )
                 logger.info(f"rank {rank} MLX_JACCL_COORDINATOR: {jaccl_coordinator}")
-                os.environ["MLX_JACCL_DEVICES"] = coordination_file
+                os.environ["MLX_IBV_DEVICES"] = coordination_file
                 os.environ["MLX_RANK"] = str(rank)
                 os.environ["MLX_JACCL_COORDINATOR"] = jaccl_coordinator
                 group = mx.distributed.init(backend="jaccl", strict=True)
@@ -463,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/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
-        )

uv.lock

@@ -376,8 +376,8 @@ dependencies = [
     { name = "hypercorn", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", extra = ["cpu"], marker = "sys_platform == 'linux'" },
+    { name = "mlx", version = "0.30.3", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" },
+    { name = "mlx", version = "0.30.4.dev20260121+fbe306f9", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git#fbe306f92a47d9b887ee7af2e3af6f1b9e28e663" }, marker = "sys_platform == 'darwin'" },
     { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "openai-harmony", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pillow", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -412,8 +412,8 @@ requires-dist = [
     { name = "huggingface-hub", specifier = ">=0.33.4" },
     { name = "hypercorn", specifier = ">=0.18.0" },
     { name = "loguru", specifier = ">=0.7.3" },
+    { name = "mlx", marker = "sys_platform == 'darwin'", git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git" },
     { name = "mflux", specifier = "==0.15.4" },
-    { name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.3" },
     { name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.3" },
     { name = "mlx-lm", specifier = "==0.30.5" },
     { name = "openai-harmony", specifier = ">=0.0.8" },
@@ -994,8 +994,8 @@ dependencies = [
     { name = "fonttools", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", extra = ["cuda13"], marker = "sys_platform == 'linux'" },
+    { name = "mlx", version = "0.30.3", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" },
+    { name = "mlx", version = "0.30.4.dev20260121+fbe306f9", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git#fbe306f92a47d9b887ee7af2e3af6f1b9e28e663" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1022,18 +1022,12 @@ wheels = [
 name = "mlx"
 version = "0.30.3"
 source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "mlx-metal", marker = "sys_platform == 'darwin'" },
+resolution-markers = [
+    "sys_platform == 'linux'",
 ]
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/d0/22/42935d593fe82d3b98eb9d60e4620ed99703886635106f89d407c68f33bc/mlx-0.30.3-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:743fac1e4f9e8e46c8262943c643a31139c255cdb256c99ad496958215ccac1e", size = 569344, upload-time = "2026-01-14T01:16:54.847Z" },
-    { url = "https://files.pythonhosted.org/packages/7d/27/f2e7a5236289d45315d0215e8553b4dd7e2faaba3bcb5025b34b25d5ab66/mlx-0.30.3-cp313-cp313-macosx_15_0_arm64.whl", hash = "sha256:3b04ae81655aa0e63a6e8f2c749de3bbce64cf5b168ae10f39ed086dfa99e7f8", size = 569345, upload-time = "2026-01-14T01:16:56.564Z" },
-    { url = "https://files.pythonhosted.org/packages/01/41/06b042457f51952456e9bb46b2c6e205ab3a28fc52d6751b5787fdb762b2/mlx-0.30.3-cp313-cp313-macosx_26_0_arm64.whl", hash = "sha256:ba9b5bdb1e929cc130af72efd7f73508c0f4e526d224489af7ec1c6419564659", size = 569213, upload-time = "2026-01-14T05:52:10.86Z" },
     { url = "https://files.pythonhosted.org/packages/ec/1e/f62c98fc0d2d878ee4235671f9d406b13cc9240493ba6fcfde2f72c2ff83/mlx-0.30.3-cp313-cp313-manylinux_2_35_aarch64.whl", hash = "sha256:dfe5c5b64e55398a22100804abbf9681996b03129e720e36b1727ed704db12b5", size = 617309, upload-time = "2026-01-14T01:16:57.58Z" },
     { url = "https://files.pythonhosted.org/packages/e9/62/811f064693449de740350d27793ce39343a460305ec8d878c318b80921d0/mlx-0.30.3-cp313-cp313-manylinux_2_35_x86_64.whl", hash = "sha256:a3364924610929936e6aaf13c71106161258e5a5d3f7813a64c07cc2435f9f55", size = 659521, upload-time = "2026-01-14T01:16:58.719Z" },
-    { url = "https://files.pythonhosted.org/packages/82/e2/6e551bd48fb350fbf0ee4cc5cd09485437d260b8f4937f22d8623e14687a/mlx-0.30.3-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:2c27fd8daaae14ca6cf407fcd236006a6e968f7708c8f61a2709116f2e754852", size = 571920, upload-time = "2026-01-14T01:16:59.683Z" },
-    { url = "https://files.pythonhosted.org/packages/82/c0/561d1c9d3d12830b0e7fdcbd807585ef20909e398d4bcdbf25e4367543eb/mlx-0.30.3-cp314-cp314-macosx_15_0_arm64.whl", hash = "sha256:b755fd4ed4b6a2ae4dee3766b5a2ea52fcbe83ebd1cf018458e18b74139409f3", size = 571921, upload-time = "2026-01-14T01:17:00.868Z" },
-    { url = "https://files.pythonhosted.org/packages/42/1a/fb573fc2edc22a777fa254ff5c0c886ffd2c88aeb1f21c45778ef170f990/mlx-0.30.3-cp314-cp314-macosx_26_0_arm64.whl", hash = "sha256:7e352c0369a2f7e54d4f317b434eab3333918ea9edde1c43c61d36386b6f76bf", size = 571732, upload-time = "2026-01-14T05:52:11.893Z" },
     { url = "https://files.pythonhosted.org/packages/9e/db/d0083e8f2205b3b2dcd9670eb6f0d6c1b7cbfea6b01a1f8bff39142edf44/mlx-0.30.3-cp314-cp314-manylinux_2_35_aarch64.whl", hash = "sha256:00ac867f3d003c1477a66a579442c2040ba7ea43ce3c174490d1f8bf379606bd", size = 619635, upload-time = "2026-01-14T01:17:01.812Z" },
     { url = "https://files.pythonhosted.org/packages/ab/90/ab0b93ff0e76da4fe0e878722c76a308cfb950b044a4676e9617276d8ccd/mlx-0.30.3-cp314-cp314-manylinux_2_35_x86_64.whl", hash = "sha256:5be7d0329036f09c6ed003ea3e307e97e3144f20a3e4711b01810d7d5013cf2c", size = 659652, upload-time = "2026-01-14T01:17:02.915Z" },
 ]
@@ -1046,6 +1040,14 @@ cuda13 = [
     { name = "mlx-cuda-13", marker = "sys_platform == 'linux'" },
 ]
 
+[[package]]
+name = "mlx"
+version = "0.30.4.dev20260121+fbe306f9"
+source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git#fbe306f92a47d9b887ee7af2e3af6f1b9e28e663" }
+resolution-markers = [
+    "sys_platform == 'darwin'",
+]
+
 [[package]]
 name = "mlx-cpu"
 version = "0.30.3"
@@ -1076,7 +1078,7 @@ version = "0.30.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
-    { name = "mlx", marker = "sys_platform == 'darwin'" },
+    { name = "mlx", version = "0.30.4.dev20260121+fbe306f9", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git#fbe306f92a47d9b887ee7af2e3af6f1b9e28e663" }, marker = "sys_platform == 'darwin'" },
     { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
     { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -1088,16 +1090,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/89/ba/66db6e1e5f1ef506655b562932f6bd8f72600116d5f31f92d71c1f200b3f/mlx_lm-0.30.5-py3-none-any.whl", hash = "sha256:a80bc8e3efdebe81813b0f6eb403fb66a7a15071e256f4e7102ada986acb75bb", size = 366716, upload-time = "2026-01-25T15:29:28.29Z" },
 ]
 
-[[package]]
-name = "mlx-metal"
-version = "0.30.3"
-source = { registry = "https://pypi.org/simple" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f6/63/4d8f6fefb507c028df4454dabfe8d8e0ad2961bb06510b6aca23d2d5b2be/mlx_metal-0.30.3-py3-none-macosx_14_0_arm64.whl", hash = "sha256:6276312b02353714c7c6515169569fe1c4bebe3229c8ecf1fdb375a13e78c966", size = 37716245, upload-time = "2026-01-14T01:16:34.838Z" },
-    { url = "https://files.pythonhosted.org/packages/35/91/1d452e48a4bb4958844fd3bb28ae31b8de110549c009ebec5024ce27ebf3/mlx_metal-0.30.3-py3-none-macosx_15_0_arm64.whl", hash = "sha256:c096c0a3428f3f96a06220f97a36f9528b18bc05173f821eb05bc8458e723fa8", size = 37712125, upload-time = "2026-01-14T01:16:38.619Z" },
-    { url = "https://files.pythonhosted.org/packages/fe/36/7a3cbca85542b5ca4faf871e35927f43aa0e3fc830ae5b699780fe723677/mlx_metal-0.30.3-py3-none-macosx_26_0_arm64.whl", hash = "sha256:69068533bd1ee8b0379ce5de57ed5fd313577a10ecab58e1332fd1ff7248a75e", size = 46488962, upload-time = "2026-01-14T05:52:04.523Z" },
-]
-
 [[package]]
 name = "more-itertools"
 version = "10.8.0"