Use sliding window

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

@@ -229,6 +229,47 @@ def has_non_kv_caches(cache: KVCacheType) -> bool:
     return any(is_non_trimmable_cache_entry(c) for c in cache)
 
 
+# Max snapshots retained per cache entry. Each CacheSnapshot pins detached GPU
+# copies of every non-trimmable (SSM/ArraysCache, RotatingKVCache) layer, so
+# retaining one per ~4096-token prefill chunk makes snapshot memory grow linearly
+# with context — the dominant residual cost when a single entry is grown to long
+# contexts on hybrid models (~56 MB/snapshot on Qwen3.5-122B, so a full 256K
+# context = 64 snapshots ≈ 3.6 GB). A sliding window of the most-recent N caps
+# this at N×per-snapshot (~0.9 GB here) while preserving the restore points
+# in-place grows actually use (they always extend from the tip).
+_MAX_RETAINED_SNAPSHOTS = 16
+
+
+def _bounded_snapshots(snapshots: list[CacheSnapshot]) -> list[CacheSnapshot]:
+    """Deduplicate snapshots by token position and bound the retained count.
+
+    Returned list is sorted ascending by ``token_count``.
+    """
+    # Deduplicate by position, keeping the most-recently-appended snapshot per
+    # position. Repeated in-place grows re-snapshot positions the kept old
+    # snapshots already cover, which would otherwise grow `_snapshots`
+    # unbounded even at constant context.
+    # TODO: keying on token_count alone is safe only while a position uniquely
+    # identifies the prefix within an entry (grows are strict prefix-extensions).
+    # If edit-and-regenerate, sliding-window/prefix trimming, cross-entry
+    # snapshot sharing, per-request adapter/LoRA swap, or branchy decoding
+    # (beam/parallel/speculative) is added, enrich the key to
+    # (token_count, prefix_hash[, media/adapter id]) — else a stale snapshot
+    # could be restored for a different prefix (silent wrong output).
+    by_position: dict[int, CacheSnapshot] = {}
+    for snapshot in snapshots:
+        by_position[snapshot.token_count] = snapshot
+    deduped = [by_position[pos] for pos in sorted(by_position)]
+
+    # Sliding window: keep only the most-recent N positions. In-place grows
+    # always extend from the tip, so the newest snapshots are the ones future
+    # grows restore from — dropping the oldest is never incorrect: a later hit on
+    # a prefix older than the window finds no snapshot <= target, so get_kv_cache
+    # returns a fresh cache (matched_index=None) and the request takes a full cold
+    # prefill — correct, just slower than a partial-hit reuse for that one request.
+    return deduped[-_MAX_RETAINED_SNAPSHOTS:]
+
+
 class KVPrefixCache:
     def __init__(self, group: mx.distributed.Group | None):
         self.prompts: list[mx.array] = []  # mx array of tokens (ints)
@@ -261,7 +302,9 @@ class KVPrefixCache:
         self._evict_if_needed()
         self.prompts.append(prompt_tokens)
         self.caches.append(deepcopy(cache))
-        self._snapshots.append(ssm_snapshots)
+        self._snapshots.append(
+            _bounded_snapshots(ssm_snapshots) if ssm_snapshots else None
+        )
         self._media_regions.append(media_regions or [])
         self.prefill_tps.append(prefill_tps)
         self._access_counter += 1
@@ -288,7 +331,7 @@ class KVPrefixCache:
 
         self.prompts[index] = prompt_tokens
         self.caches[index] = deepcopy(cache)
-        self._snapshots[index] = merged or None
+        self._snapshots[index] = _bounded_snapshots(merged) or None
         self._media_regions[index] = media_regions or []
         self.prefill_tps[index] = prefill_tps
         self._access_counter += 1

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

@@ -11,6 +11,7 @@ from mlx_lm.sample_utils import make_sampler
 from exo.shared.types.common import ModelId
 from exo.shared.types.text_generation import InputMessage, TextGenerationTaskParams
 from exo.worker.engines.mlx.cache import (
+    CacheSnapshot,
     KVPrefixCache,
     cache_length,
     encode_prompt,
@@ -77,6 +78,74 @@ class TestGetPrefixLength:
         assert get_prefix_length(a, b) == 0
 
 
+class TestSnapshotAccumulation:
+    """Locks in the fix for the actual per-grow Metal leak on hybrid (SSM)
+    models: `update_kv_cache` must not let `_snapshots` grow without bound when
+    the same entry is grown in place many times."""
+
+    def test_repeated_update_does_not_accumulate_snapshots(self):
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            initial = [
+                CacheSnapshot(states=[None], token_count=4096),
+                CacheSnapshot(states=[None], token_count=8192),
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10000), [KVCache()], ssm_snapshots=initial
+            )
+
+            # Each in-place grow re-prefills from restore_pos and produces a
+            # fresh snapshot at a position the retained old snapshots already
+            # cover. Pre-fix this appended one snapshot per grow forever.
+            for _ in range(50):
+                fresh = [CacheSnapshot(states=[None], token_count=8192)]
+                kv_prefix_cache.update_kv_cache(
+                    0, mx.arange(10000), [KVCache()], fresh, restore_pos=8192
+                )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Bounded by the number of distinct snapshot positions (here 2),
+            # not by the 50 grows.
+            assert len(stored) == 2
+            assert sorted(s.token_count for s in stored) == [4096, 8192]
+            # The kept 8192 snapshot must be the most recently supplied one.
+            assert stored[1] is fresh[0]
+
+    def test_extension_caps_snapshots_to_sliding_window(self):
+        """Extending a single entry to a long context (one snapshot per ~4096
+        tokens) must cap retained snapshots to a sliding window of the most-recent
+        N, not keep all of them — that linear-in-context retention was the
+        residual OOM cause."""
+        from exo.worker.engines.mlx.cache import _MAX_RETAINED_SNAPSHOTS
+
+        with patch(
+            "exo.worker.engines.mlx.cache.get_memory_used_percentage",
+            return_value=0.0,
+        ):
+            kv_prefix_cache = KVPrefixCache(None)
+            # 64 distinct positions = a 262144-token context at 4096/chunk.
+            num_positions = 64
+            snaps = [
+                CacheSnapshot(states=[None], token_count=4096 * (i + 1))
+                for i in range(num_positions)
+            ]
+            kv_prefix_cache.add_kv_cache(
+                mx.arange(10), [KVCache()], ssm_snapshots=snaps
+            )
+
+            stored = kv_prefix_cache._snapshots[0]
+            assert stored is not None
+            # Capped at the window; the most-recent N positions are retained
+            # (in-place grows extend from the tip, so these are what get used).
+            assert len(stored) == _MAX_RETAINED_SNAPSHOTS
+            assert stored == snaps[-_MAX_RETAINED_SNAPSHOTS:]
+            assert stored[-1] is snaps[-1]  # tip always kept
+
+
 class TestKVPrefix:
     @pytest.fixture
     def mock_tokenizer(self):