Force a clean prefill after chained KV prefix-cache extensions

src/exo/api/adapters/ollama.py

@@ -195,10 +195,16 @@ async def generate_ollama_chat_stream(
                 continue
 
             case ErrorChunk():
+                # Ollama's response shape has no error field, so failures are
+                # reported through message.content — which many clients render
+                # unconditionally. Never echo chunk.error_message here: it can
+                # carry raw generation text (partial tool-call markup, document
+                # content). The detail is already in the server logs.
                 error_response = OllamaChatResponse(
                     model=str(chunk.model),
                     message=OllamaMessage(
-                        role="assistant", content=chunk.error_message
+                        role="assistant",
+                        content="Internal server error",
                     ),
                     done=True,
                     done_reason="error",

src/exo/api/main.py

@@ -868,6 +868,13 @@ async def _send_text_generation_with_images(
         self, task_params: TextGenerationTaskParams
     ) -> TextGeneration:
         task_params = task_params.with_card_sampling_defaults()
+        if task_params.seed is None and not task_params.bench:
+            # Draw the seed once on the master so every rank samples identically.
+            # Runners fall back to a fixed seed when the task carries none, which
+            # makes a retry of a failed generation replay the exact same failure.
+            task_params = task_params.model_copy(
+                update={"seed": random.randint(0, 2**31 - 1)}
+            )
         images = task_params.images
         if not images:
             command = TextGeneration(task_params=task_params)

src/exo/api/tests/test_seed_default.py

@@ -0,0 +1,42 @@
+"""The master draws a sampling seed when a request doesn't set one.
+
+Runners fall back to a fixed seed when the task carries none, which makes a
+retry of a failed generation replay the exact same failure — so
+_send_text_generation_with_images must draw a seed for normal requests while
+preserving explicit seeds and leaving bench runs deterministic.
+"""
+
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+from exo.api.main import API
+from exo.shared.types.common import ModelId
+from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
+
+
+def _task_params(**overrides: object) -> TextGenerationTaskParams:
+    return TextGenerationTaskParams(
+        model=ModelId("test-org/test-model"),
+        input=[InputMessage(role="user", content="hi")],
+        **overrides,  # pyright: ignore[reportAny]
+    )
+
+
+async def _send_command(task_params: TextGenerationTaskParams):
+    api = SimpleNamespace(_send=AsyncMock())
+    return await API._send_text_generation_with_images(api, task_params)  # pyright: ignore[reportPrivateUsage, reportArgumentType]
+
+
+async def test_missing_seed_gets_drawn():
+    command = await _send_command(_task_params())
+    assert command.task_params.seed is not None
+
+
+async def test_explicit_seed_preserved():
+    command = await _send_command(_task_params(seed=1234))
+    assert command.task_params.seed == 1234
+
+
+async def test_bench_keeps_seed_unset():
+    command = await _send_command(_task_params(bench=True))
+    assert command.task_params.seed is None

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

@@ -47,6 +47,19 @@ def _default_memory_threshold() -> float:
     os.environ.get("EXO_MEMORY_THRESHOLD", _default_memory_threshold())
 )
 
+# An entry whose tail has been built by MAX_CHAIN_DEPTH successive incremental
+# extensions (a tool loop extends the same entry on every iteration) is not
+# offered for further deep reuse — the chained partial prefills accumulate
+# numerical drift in the cached KV that can flip borderline token decisions
+# (observed: the model emitting its end-of-turn token in the middle of a tool
+# call at chain depth >= 2). Forcing a full re-prefill every MAX_CHAIN_DEPTH-th
+# extension resets the entry to a single-pass KV state. Shallow reuse (borrowing
+# only the shared prompt prefix, e.g. the system/tools block) is unaffected:
+# that region was written once by the entry's first prefill and never rewritten
+# by extensions.
+MAX_CHAIN_DEPTH = int(os.environ.get("EXO_KV_MAX_CHAIN_DEPTH", 2))
+_DEEP_REUSE_RATIO = 0.8
+
 
 class CacheSnapshot:
     """Snapshot of states at a known token position."""
@@ -237,6 +250,8 @@ def __init__(self, group: mx.distributed.Group | None):
         self._media_regions: list[list["MediaRegion"]] = []
         self._last_used: list[int] = []  # monotonic counter of last access per entry
         self.prefill_tps: list[float] = []
+        # Successive deep incremental extensions applied to each entry's tail.
+        self.chain_depths: list[int] = []
         self._access_counter: int = 0
         self._group = group
 
@@ -248,6 +263,7 @@ def clear(self):
         self._media_regions.clear()
         self._last_used.clear()
         self.prefill_tps.clear()
+        self.chain_depths.clear()
 
     def add_kv_cache(
         self,
@@ -264,6 +280,7 @@ def add_kv_cache(
         self._snapshots.append(ssm_snapshots)
         self._media_regions.append(media_regions or [])
         self.prefill_tps.append(prefill_tps)
+        self.chain_depths.append(0)
         self._access_counter += 1
         self._last_used.append(self._access_counter)
         logger.info(f"KV cache added: {len(prompt_tokens)} tokens")
@@ -286,6 +303,21 @@ def update_kv_cache(
         if snapshots:
             merged.extend(snapshots)
 
+        # Track how many successive deep extensions built this entry's tail.
+        # Growing the entry while keeping most of its previous content is a
+        # deep extension; any other rewrite (shallow-prefix rebuild, shrink)
+        # replaces the tail with fresh single-pass content and resets the
+        # depth; an exact re-save of the same prompt preserves it.
+        old_len = len(self.prompts[index])
+        deep_growth = (
+            len(prompt_tokens) > old_len > 0
+            and restore_pos >= _DEEP_REUSE_RATIO * old_len
+        )
+        if len(prompt_tokens) != old_len:
+            self.chain_depths[index] = (
+                self.chain_depths[index] + 1 if deep_growth else 0
+            )
+
         self.prompts[index] = prompt_tokens
         self.caches[index] = deepcopy(cache)
         self._snapshots[index] = merged or None
@@ -350,6 +382,19 @@ def get_kv_cache(
                     self._media_regions[i],
                     query_regions,
                 )
+            if (
+                length > 0
+                and self.chain_depths[i] >= MAX_CHAIN_DEPTH
+                and length >= _DEEP_REUSE_RATIO * len(cached_prompt)
+            ):
+                # This entry's tail carries MAX_CHAIN_DEPTH chained partial
+                # prefills — skip it for deep reuse so the caller re-prefills
+                # from scratch and the subsequent save resets the chain.
+                logger.info(
+                    f"KV cache entry {i} at chain depth {self.chain_depths[i]} — "
+                    f"skipping deep reuse ({length}/{len(cached_prompt)} tokens) to force a clean prefill"
+                )
+                continue
             if length >= max_length - 1:
                 best_index, best_length = i, length
                 is_exact = True
@@ -445,6 +490,7 @@ def _evict_if_needed(self):
             self._media_regions.pop(lru_index)
             self._last_used.pop(lru_index)
             self.prefill_tps.pop(lru_index)
+            self.chain_depths.pop(lru_index)
 
             evicted_any = True
             logger.info(

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

@@ -543,7 +543,8 @@ def mlx_generate(
 ) -> Generator[GenerationResponse]:
     # Ensure that generation stats only contains peak memory for this generation
     mx.reset_peak_memory()
-    # TODO: Randomise task seed and set in taskparams, instead of hard coding as 42.
+    # The master draws a random seed for normal requests; bench runs leave it
+    # unset and fall back to a fixed seed so their results stay reproducible.
     seed = task.seed or 42
     mx.random.seed(seed)
 

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

@@ -11,6 +11,7 @@
 from exo.shared.types.common import ModelId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.worker.engines.mlx.cache import (
+    MAX_CHAIN_DEPTH,
     KVPrefixCache,
     cache_length,
     encode_prompt,
@@ -106,6 +107,127 @@ def test_clear_on_empty_cache(self, mock_tokenizer):
         assert len(cache.prompts) == 0
 
 
+def _kv_with_offset(offset: int) -> list[KVCache]:
+    kv = KVCache()
+    kv.offset = offset
+    return [kv]
+
+
+def _fake_model() -> Model:
+    from types import SimpleNamespace
+
+    return cast(Model, SimpleNamespace(layers=[None]))
+
+
+class TestChainDepthCap:
+    """Chained incremental extensions of one entry accumulate numerical drift
+    in the cached KV; the cap forces a clean full prefill at MAX_CHAIN_DEPTH."""
+
+    def _cache_with_entry(self, prompt_len: int = 100, depth: int = 0) -> KVPrefixCache:
+        cache = KVPrefixCache(None)
+        cache.add_kv_cache(
+            mx.array(list(range(prompt_len))), _kv_with_offset(prompt_len - 2)
+        )
+        cache.chain_depths[0] = depth
+        return cache
+
+    def test_add_entry_starts_at_depth_zero(self):
+        cache = self._cache_with_entry()
+        assert cache.chain_depths == [0]
+
+    def test_deep_extension_increments_depth(self):
+        cache = self._cache_with_entry(prompt_len=100)
+        cache.update_kv_cache(
+            0, mx.array(list(range(150))), _kv_with_offset(148), None, restore_pos=98
+        )
+        assert cache.chain_depths == [1]
+        cache.update_kv_cache(
+            0, mx.array(list(range(200))), _kv_with_offset(198), None, restore_pos=148
+        )
+        assert cache.chain_depths == [2]
+
+    def test_shallow_rebuild_resets_depth(self):
+        cache = self._cache_with_entry(prompt_len=200, depth=2)
+        # Rebuild keeping only a shallow shared prefix (< _DEEP_REUSE_RATIO).
+        cache.update_kv_cache(
+            0, mx.array(list(range(300))), _kv_with_offset(298), None, restore_pos=50
+        )
+        assert cache.chain_depths == [0]
+
+    def test_same_length_resave_keeps_depth(self):
+        cache = self._cache_with_entry(prompt_len=100, depth=1)
+        cache.update_kv_cache(
+            0, mx.array(list(range(100))), _kv_with_offset(98), None, restore_pos=98
+        )
+        assert cache.chain_depths == [1]
+
+    def test_shrink_update_resets_depth(self):
+        # A shrinking rewrite replaces the tail with fresh single-pass content;
+        # the previous tail's chain depth must not carry over.
+        cache = self._cache_with_entry(prompt_len=200, depth=MAX_CHAIN_DEPTH)
+        cache.update_kv_cache(
+            0, mx.array(list(range(100))), _kv_with_offset(98), None, restore_pos=98
+        )
+        assert cache.chain_depths == [0]
+
+    def test_deep_extension_at_exact_ratio_boundary(self):
+        # restore_pos == _DEEP_REUSE_RATIO * old_len must count as deep (>=,
+        # not >): 80 == 0.8 * 100.
+        cache = self._cache_with_entry(prompt_len=100)
+        cache.update_kv_cache(
+            0, mx.array(list(range(150))), _kv_with_offset(148), None, restore_pos=80
+        )
+        assert cache.chain_depths == [1]
+
+    def test_deep_reuse_skipped_at_max_chain_depth(self):
+        cache = self._cache_with_entry(prompt_len=100, depth=MAX_CHAIN_DEPTH)
+        query = mx.array(list(range(120)))
+        _, remaining, matched_index, is_exact = cache.get_kv_cache(_fake_model(), query)
+        assert matched_index is None
+        assert not is_exact
+        assert len(remaining) == 120
+
+    def test_deep_reuse_allowed_below_max_chain_depth(self):
+        cache = self._cache_with_entry(prompt_len=100, depth=MAX_CHAIN_DEPTH - 1)
+        query = mx.array(list(range(120)))
+        _, remaining, matched_index, _ = cache.get_kv_cache(_fake_model(), query)
+        assert matched_index == 0
+        # The restore target is clamped to the entry's cached length (98), so
+        # exactly the last 120 - 98 = 22 query tokens remain to prefill.
+        assert len(remaining) == 22
+
+    def test_deep_reuse_skipped_at_exact_ratio_boundary(self):
+        # A shared prefix of exactly _DEEP_REUSE_RATIO * len(cached_prompt)
+        # tokens (80 == 0.8 * 100) counts as deep reuse (>=, not >) and is
+        # skipped at MAX_CHAIN_DEPTH.
+        cache = self._cache_with_entry(prompt_len=100, depth=MAX_CHAIN_DEPTH)
+        query = mx.array(list(range(80)) + [999] * 40)
+        _, remaining, matched_index, _ = cache.get_kv_cache(_fake_model(), query)
+        assert matched_index is None
+        assert len(remaining) == 120
+
+    def test_shallow_borrow_allowed_at_max_chain_depth(self):
+        cache = self._cache_with_entry(prompt_len=100, depth=MAX_CHAIN_DEPTH)
+        # Query shares only the first 50 tokens — well under the deep-reuse
+        # ratio, so borrowing the shared prefix stays allowed.
+        query = mx.array(list(range(50)) + [999] * 70)
+        _, remaining, matched_index, _ = cache.get_kv_cache(_fake_model(), query)
+        assert matched_index == 0
+        assert len(remaining) == 70
+
+    def test_eviction_keeps_depths_aligned(self):
+        cache = KVPrefixCache(None)
+        with patch.object(
+            KVPrefixCache, "get_memory_used_percentage", side_effect=[0.0, 1.0, 0.0]
+        ):
+            cache.add_kv_cache(mx.array([1, 2, 3]), _kv_with_offset(3))
+            cache.add_kv_cache(mx.array([4, 5, 6]), _kv_with_offset(3))
+            cache.chain_depths[0] = 2
+            cache._evict_if_needed()
+        assert len(cache.prompts) == 1
+        assert cache.chain_depths == [0]
+
+
 def _load_gpt_oss() -> tuple[Model, object]:
     from mlx_lm.utils import load_model