diff --git a/mlx_lm/models/cache.py b/mlx_lm/models/cache.py index e699324..756ce4e 100644 --- a/mlx_lm/models/cache.py +++ b/mlx_lm/models/cache.py @@ -1434,18 +1434,22 @@ class PromptTrie: def pop_prefixes(self, model: Any, tokens: List[int]): values = [] current = self._trie[model] - for i in range(len(tokens) - 1): + for i, tok in enumerate(tokens): if "__value__" in current: values.append((i, current.pop("__value__"))) - current = current[tokens[i]] + current = current[tok] return values def search(self, model: Any, tokens: List[int]) -> PromptTrieResult: if model not in self._trie: return PromptTrieResult(model, None, None, None, 0) - # Walk the tokens as far as we can current = self._trie[model] + + if not tokens and "__value__" in current: + return PromptTrieResult(model, [], None, None, 0) + + # Walk the tokens as far as we can last_index = -1 index = 0 while index < len(tokens) and tokens[index] in current: @@ -1455,7 +1459,7 @@ class PromptTrie: index += 1 # Got an exact match - if last_index == len(tokens) - 1: + if last_index == len(tokens) - 1 >= 0: return PromptTrieResult(model, tokens, None, None, 0) # Check if we found a prefix at any point diff --git a/tests/test_server.py b/tests/test_server.py index 9fa2fbf..b6907ac 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -537,6 +537,41 @@ class TestLRUPromptCache(unittest.TestCase): self.assertEqual(c, [MockCache("test4")]) self.assertEqual(t, []) + def test_insert_trimmable_cache_removes_immediate_prefix(self): + cache = LRUPromptCache(max_size=10) + model = ("test", None, None) + + cache.insert_cache(model, [1, 2], [MockCache("ab")]) + self.assertEqual(len(cache), 1) + self.assertEqual(cache.nbytes, 2) + + cache.insert_cache(model, [1, 2, 3], [MockCache("abc")]) + self.assertEqual(len(cache), 1) + self.assertEqual(cache.nbytes, 3) + + def test_insert_empty_tokens_does_not_self_destruct(self): + cache = LRUPromptCache(max_size=10) + model = ("test", None, None) + + cache.insert_cache(model, [], [MockCache("root")]) + self.assertEqual(len(cache), 1) + self.assertEqual(cache.nbytes, 4) + + c, t = cache.fetch_nearest_cache(model, []) + self.assertIsNotNone(c) + self.assertEqual(t, []) + + def test_fetch_empty_tokens_after_root_eviction(self): + cache = LRUPromptCache(max_size=10) + model = ("test", None, None) + + cache.insert_cache(model, [], [MockCache("root")]) + cache.insert_cache(model, [1], [MockCache("a")]) + + c, t = cache.fetch_nearest_cache(model, []) + self.assertIsNone(c) + self.assertEqual(t, []) + def test_lru_bytes(self): cache = LRUPromptCache(max_size=100, max_bytes=10) model = ("test", None, None)