diff --git a/src/exo/api/adapters/ollama.py b/src/exo/api/adapters/ollama.py index 1f027616ab..041badf99d 100644 --- a/src/exo/api/adapters/ollama.py +++ b/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", diff --git a/src/exo/api/main.py b/src/exo/api/main.py index fcd54c9315..5b7926878b 100644 --- a/src/exo/api/main.py +++ b/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) diff --git a/src/exo/api/tests/test_seed_default.py b/src/exo/api/tests/test_seed_default.py new file mode 100644 index 0000000000..30ed9b268d --- /dev/null +++ b/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 diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index 7cdcc77fbe..2a92560d9d 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/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( diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 2e3d051251..8b97386b2d 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/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) diff --git a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py b/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py index 3d72d47d62..cbd8a03269 100644 --- a/src/exo/worker/tests/unittests/test_mlx/test_kv_prefix_cache.py +++ b/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