diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py index 7cdcc77fbe..159347bc41 100644 --- a/src/exo/worker/engines/mlx/cache.py +++ b/src/exo/worker/engines/mlx/cache.py @@ -476,10 +476,7 @@ def trim_cache( snapshot: CacheSnapshot | None = None, ) -> None: for i, c in enumerate(cache): - non_trimmable = isinstance(c, (ArraysCache, RotatingKVCache)) or ( - isinstance(c, CacheList) and not bool(c.is_trimmable()) # type: ignore[reportUnknownMemberType] - ) - if non_trimmable: + if is_non_trimmable_cache_entry(c): if snapshot is not None and snapshot.states[i] is not None: restored = copy_snapshot_entry(snapshot.states[i]) if restored is not None: @@ -489,7 +486,7 @@ def trim_cache( if isinstance(c, RotatingKVCache): c.offset = 0 c._idx = 0 - else: + elif isinstance(c, CacheList): # CacheList without a snapshot — zero each inner cache's state for inner in c: # type: ignore[reportUnknownVariableType] if isinstance(inner, (ArraysCache, RotatingKVCache)): @@ -497,6 +494,8 @@ def trim_cache( if isinstance(inner, RotatingKVCache): inner.offset = 0 inner._idx = 0 + # DeepseekV4Cache without a snapshot cannot be rolled back in place + # (and is not iterable) — callers must restore via snapshot. else: c.trim(num_tokens) 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..e5d562b72d 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 @@ -6,6 +6,7 @@ import mlx.core as mx import pytest from mlx_lm.models.cache import KVCache +from mlx_lm.models.deepseek_v4 import DeepseekV4Cache from mlx_lm.sample_utils import make_sampler from exo.shared.types.common import ModelId @@ -16,6 +17,8 @@ encode_prompt, get_prefix_length, make_kv_cache, + snapshot_ssm_states, + trim_cache, ) from exo.worker.engines.mlx.generator.generate import mlx_generate, prefill from exo.worker.engines.mlx.types import Model @@ -600,3 +603,56 @@ def test_evicts_lru_entry_under_memory_pressure(self, model_and_tokenizer): assert len(kv_prefix_cache.prompts) == 1 # The surviving entry should be the newly added one assert get_prefix_length(kv_prefix_cache.prompts[0], tokens) == len(tokens) + + +class TestTrimCacheDeepseekV4: + """trim_cache must treat DeepseekV4Cache as non-trimmable. + + DeepseekV4Cache.is_trimmable() is False and .trim() is a no-op, and + is_non_trimmable_cache_entry() already classifies it as non-trimmable — + but trim_cache used a stale inline copy of that predicate which omitted + it, so V4 entries fell through to the no-op .trim() and were silently + reused untrimmed instead of being restored from the snapshot. + """ + + @staticmethod + def _v4_cache(n_tokens: int, sliding_window: int = 64) -> DeepseekV4Cache: + c = DeepseekV4Cache(sliding_window) + c.local.update_and_fetch( + mx.random.normal([1, 1, n_tokens, 8]), + mx.random.normal([1, 1, n_tokens, 8]), + ) + return c + + def test_restores_v4_from_snapshot(self): + v4 = self._v4_cache(200) + kv = KVCache() + kv.update_and_fetch( + mx.random.normal([1, 2, 200, 8]), mx.random.normal([1, 2, 200, 8]) + ) + cache = [kv, v4] + snapshot = snapshot_ssm_states(cache) + assert snapshot.token_count == 200 + + # Advance both entries past the snapshot point. + v4.local.update_and_fetch( + mx.random.normal([1, 1, 100, 8]), mx.random.normal([1, 1, 100, 8]) + ) + kv.update_and_fetch( + mx.random.normal([1, 2, 100, 8]), mx.random.normal([1, 2, 100, 8]) + ) + assert cache_length(cache) == 300 + + trim_cache(cache, 100, snapshot) + + # V4 entry must come back from the snapshot, not stay at 300. + assert cache[1].offset == 200 + # Trimmable entries still trim normally. + assert cache[0].offset == 200 + + def test_v4_without_snapshot_does_not_raise(self): + v4 = self._v4_cache(200) + # No snapshot available: must be a no-op, not a TypeError from + # iterating a non-iterable cache in the CacheList fallback arm. + trim_cache([v4], 50, None) + assert v4.offset == 200