Refactor LRUPromptCache (#1019)

mlx_lm/models/cache.py

@@ -1,6 +1,8 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import copy
+from collections import deque
+from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
@@ -1381,3 +1383,224 @@ class BatchRotatingKVCache(_BaseCache):
         if self.keys is None:
             return 0
         return self.keys.nbytes + self.values.nbytes
+
+
+@dataclass
+class PromptTrieResult:
+    model: Any
+    exact: Optional[List[int]]  # Exact match found
+    shorter: Optional[List[int]]  # Longest prefix with a value
+    longer: Optional[List[int]]  # Shortest value that extends beyond tokens
+    common_prefix: int  # Length of common prefix with any path
+
+
+class PromptTrie:
+    def __init__(self):
+        self._trie = {}
+
+    def add(self, model: Any, tokens: List[int], value: Any):
+        if model not in self._trie:
+            self._trie[model] = {}
+
+        current = self._trie[model]
+        for tok in tokens:
+            if tok not in current:
+                current[tok] = {}
+            current = current[tok]
+        prev = current.get("__value__", None)
+        current["__value__"] = value
+        return prev
+
+    def get(self, model: Any, tokens: List[int]):
+        current = self._trie[model]
+        for tok in tokens:
+            current = current[tok]
+        return current["__value__"]
+
+    def pop(self, model: Any, tokens: List[int]):
+        path = [self._trie[model]]
+        for tok in tokens:
+            path.append(path[-1][tok])
+        value = path[-1].pop("__value__")
+        for i in range(len(tokens), 0, -1):
+            node = path[i]
+            parent = path[i - 1]
+            tok = tokens[i - 1]
+            if len(node) > 0:
+                break
+            del parent[tok]
+        return value
+
+    def pop_prefixes(self, model: Any, tokens: List[int]):
+        values = []
+        current = self._trie[model]
+        for i in range(len(tokens) - 1):
+            if "__value__" in current:
+                values.append((i, current.pop("__value__")))
+            current = current[tokens[i]]
+        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]
+        last_index = -1
+        index = 0
+        while index < len(tokens) and tokens[index] in current:
+            current = current[tokens[index]]
+            if "__value__" in current:
+                last_index = index
+            index += 1
+
+        # Got an exact match
+        if last_index == len(tokens) - 1:
+            return PromptTrieResult(model, tokens, None, None, 0)
+
+        # Check if we found a prefix at any point
+        shorter = None
+        if last_index > 0:
+            shorter = tokens[: last_index + 1]
+
+        # Check for sequences that are longer
+        longer = None
+        common_prefix = index
+        if index > 0:
+            best = None
+            stack = [(current, [])]
+            while stack:
+                current, extra = stack.pop()
+                if "__value__" in current:
+                    if best is None or len(extra) < len(best):
+                        best = extra
+                elif best is None or len(extra) < len(best):
+                    for tok in current:
+                        stack.append((current[tok], extra + [tok]))
+            longer = tokens[:index] + best
+        return PromptTrieResult(model, None, shorter, longer, common_prefix)
+
+
+class LRUPromptCache:
+    @dataclass
+    class CacheEntry:
+        prompt_cache: List[Any]
+        nbytes: int
+
+    class CacheOrder:
+        def __init__(self, ordering: List[str] = ["assistant", "user", "system"]):
+            self._ordering = ordering
+            self._lrus = {k: deque() for k in ordering}
+
+        def __len__(self):
+            return sum(len(lru) for lru in self._lrus.values())
+
+        def push(self, model: Any, tokens: List[Any], cache_type: str = "assistant"):
+            self._lrus[cache_type].append((model, tokens))
+
+        def remove(self, model: Any, tokens: List[Any]):
+            for cache_type in self._ordering:
+                try:
+                    self._lrus[cache_type].remove((model, tokens))
+                    break
+                except ValueError:
+                    pass
+
+        def pop(self):
+            i = 0
+            while i + 1 < len(self._ordering):
+                lru_a = self._lrus[self._ordering[i]]
+                lru_b = self._lrus[self._ordering[i + 1]]
+                if len(lru_a) >= len(lru_b):
+                    return lru_a.popleft()
+                i += 1
+            return lru_b.popleft()
+
+    def __init__(self, max_size: int = 10, max_bytes: int = 1 << 63):
+        self.max_size = max_size
+        self.max_bytes = max_bytes
+        self._trie = PromptTrie()
+        self._lru = LRUPromptCache.CacheOrder()
+        self._n_bytes = 0
+
+    def __len__(self):
+        return len(self._lru)
+
+    @property
+    def nbytes(self):
+        return self._n_bytes
+
+    def fetch_nearest_cache(self, model: Any, tokens: List[int]):
+        result = self._trie.search(model, tokens)
+        if result.exact is not None:
+            cache_entry = self._trie.get(result.model, result.exact)
+            return copy.deepcopy(cache_entry.prompt_cache), []
+
+        short_length = len(result.shorter) if result.shorter is not None else 0
+        if result.longer is not None and result.common_prefix > short_length:
+            cache_entry = self._trie.get(result.model, result.longer)
+            if can_trim_prompt_cache(cache_entry.prompt_cache):
+                cache = copy.deepcopy(cache_entry.prompt_cache)
+                prefix = min(len(tokens) - 1, result.common_prefix)
+                num_to_trim = len(result.longer) - prefix
+                trim_prompt_cache(cache, num_to_trim)
+                return cache, tokens[prefix:]
+
+        if short_length > 0:
+            cache_entry = self._trie.get(result.model, result.shorter)
+            return copy.deepcopy(cache_entry.prompt_cache), tokens[short_length:]
+
+        return None, tokens
+
+    def insert_cache(
+        self,
+        model: Any,
+        tokens: List[int],
+        prompt_cache: List[Any],
+        *,
+        cache_type: str = "assistant",
+    ):
+        # Make the cache entry
+        entry = LRUPromptCache.CacheEntry(
+            prompt_cache, sum(c.nbytes for c in prompt_cache)
+        )
+
+        # Insert into the trie and update the byte counter and lru position
+        self._n_bytes += entry.nbytes
+        prev = self._trie.add(model, tokens, entry)
+        if prev is not None:
+            self._n_bytes -= prev.nbytes
+            self._lru.remove(model, tokens)
+        self._lru.push(model, tokens, cache_type)
+
+        # If it is a trimmable cache remove all prefixes cause they just take
+        # space
+        if can_trim_prompt_cache(prompt_cache):
+            for prefix_len, entry in self._trie.pop_prefixes(model, tokens):
+                self._n_bytes -= entry.nbytes
+                self._lru.remove(model, tokens[:prefix_len])
+
+        # Ensure we match the constraints
+        if len(self._lru) > self.max_size:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes
+        while self._n_bytes > self.max_bytes:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes
+
+    def trim_to(
+        self, *, n_sequences: Optional[int] = None, n_bytes: Optional[int] = None
+    ):
+        n_sequences = max(0, n_sequences) if n_sequences is not None else 1 << 63
+        n_bytes = max(0, n_bytes) if n_bytes is not None else 1 << 63
+
+        while len(self._lru) > n_sequences:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes
+        while self._n_bytes > n_bytes:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes

mlx_lm/server.py

@@ -1,8 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-import copy
-import heapq
 import json
 import logging
 import pickle
@@ -11,7 +9,6 @@ import socket
 import time
 import uuid
 import warnings
-from collections import deque
 from dataclasses import dataclass, field
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from pathlib import Path
@@ -37,6 +34,7 @@ from huggingface_hub import scan_cache_dir
 from ._version import __version__
 from .generate import BatchGenerator, generation_stream, stream_generate
 from .models.cache import (
+    LRUPromptCache,
     can_trim_prompt_cache,
     make_prompt_cache,
     trim_prompt_cache,
@@ -171,195 +169,6 @@ def process_message_content(messages):
                         func["arguments"] = json.loads(args)
 
 
-class LRUPromptCache:
-    @dataclass
-    class CacheEntry:
-        prompt_cache: List[Any]
-        nbytes: int
-
-    class CacheOrder:
-        def __init__(self):
-            self._lru_checkpoints = deque()
-            self._lru = deque()
-
-        def __len__(self):
-            return len(self._lru) + len(self._lru_checkpoints)
-
-        def push(self, model, tokens, checkpoint: bool = False):
-            c = self._lru_checkpoints if checkpoint else self._lru
-            c.append((model, tokens))
-
-        def remove(self, model, tokens):
-            try:
-                self._lru.remove((model, tokens))
-            except ValueError:
-                self._lru_checkpoints.remove((model, tokens))
-
-        def pop(self):
-            if len(self._lru) >= len(self._lru_checkpoints):
-                return self._lru.popleft()
-            else:
-                return self._lru_checkpoints.popleft()
-
-    @dataclass
-    class SearchResult:
-        model: Any
-        exact: List[int]
-        shorter: List[int]
-        longer: List[int]
-        common_prefix: int
-
-    def __init__(self, max_size: int = 10, max_bytes: int = 1 << 63):
-        self.max_size = max_size
-        self.max_bytes = max_bytes
-        self._cache = {}
-        self._lru = self.CacheOrder()
-        self._n_bytes = 0
-
-    def __len__(self):
-        return len(self._lru)
-
-    @property
-    def nbytes(self):
-        return self._n_bytes
-
-    def _search(self, model, tokens):
-        """Search the cache for a prompt cache. Return exact or close match."""
-        if model not in self._cache:
-            return self.SearchResult(model, None, None, None, 0)
-
-        current = self._cache[model]
-        last_cache_index = -1
-        index = 0
-
-        while index < len(tokens) and tokens[index] in current:
-            current = current[tokens[index]]
-            if "cache" in current:
-                last_cache_index = index
-            index += 1
-
-        # Exact match no need to search for longer or shorter caches
-        if last_cache_index == len(tokens) - 1:
-            return self.SearchResult(model, tokens, None, None, 0)
-
-        # Find the shorter cache
-        shorter = None
-        if last_cache_index > 0:
-            shorter = tokens[: last_cache_index + 1]
-
-        # Check for caches that are longer
-        longer = None
-        common_prefix = index
-        if index > 0:
-            best = None
-            stack = [(current, [])]
-            while stack:
-                current, extra = stack.pop()
-                if "cache" in current:
-                    if best is None or len(extra) < len(best):
-                        best = extra
-                else:
-                    for tok in current:
-                        stack.append((current[tok], extra + [tok]))
-            longer = tokens[:index] + best
-        return self.SearchResult(model, None, shorter, longer, common_prefix)
-
-    def _get(self, model, tokens):
-        current = self._cache[model]
-        for tok in tokens:
-            current = current[tok]
-        return current["cache"]
-
-    def _delete(self, model, tokens):
-        path = [self._cache[model]]
-        for tok in tokens:
-            path.append(path[-1][tok])
-        cache_bytes = path[-1]["cache"].nbytes
-        self._n_bytes -= cache_bytes
-        del path[-1]["cache"]
-        for i in reversed(range(len(tokens))):
-            d_prev, d, t = path[i], path[i + 1], tokens[i]
-            if len(d) > 0:
-                break
-            del d_prev[t]
-
-    def fetch_nearest_cache(self, model, tokens):
-        result = self._search(model, tokens)
-        if result.exact is not None:
-            cache_entry = self._get(result.model, result.exact)
-            return copy.deepcopy(cache_entry.prompt_cache), []
-
-        short_length = len(result.shorter) if result.shorter is not None else 0
-        if result.longer is not None and result.common_prefix > short_length:
-            cache_entry = self._get(result.model, result.longer)
-            if can_trim_prompt_cache(cache_entry.prompt_cache):
-                cache = copy.deepcopy(cache_entry.prompt_cache)
-                prefix = min(len(tokens) - 1, result.common_prefix)
-                num_to_trim = len(result.longer) - prefix
-                trim_prompt_cache(cache, num_to_trim)
-                return cache, tokens[prefix:]
-
-        if short_length > 0:
-            cache_entry = self._get(result.model, result.shorter)
-            return copy.deepcopy(cache_entry.prompt_cache), tokens[short_length:]
-
-        return None, tokens
-
-    def insert_cache(self, model, tokens, prompt_cache, checkpoint: bool = False):
-        is_trimmable = can_trim_prompt_cache(prompt_cache)
-
-        if model not in self._cache:
-            self._cache[model] = {}
-        current = self._cache[model]
-        for i, tok in enumerate(tokens):
-            if tok not in current:
-                current[tok] = {}
-            if is_trimmable and "cache" in current:
-                self._n_bytes -= current["cache"].nbytes
-                del current["cache"]
-                self._lru.remove(model, tokens[:i])
-            current = current[tok]
-
-        if "cache" in current:
-            self._lru.remove(model, tokens)
-        else:
-            cache_bytes = sum(c.nbytes for c in prompt_cache)
-            current["cache"] = self.CacheEntry(prompt_cache, cache_bytes)
-            self._n_bytes += cache_bytes
-
-        self._lru.push(model, tokens, checkpoint=checkpoint)
-        if len(self._lru) > self.max_size:
-            model, tokens = self._lru.pop()
-            self._delete(model, tokens)
-        while self._n_bytes > self.max_bytes and len(self._lru) > 1:
-            model, tokens = self._lru.pop()
-            self._delete(model, tokens)
-
-    def trim_to(
-        self, *, n_sequences: Optional[int] = None, n_bytes: Optional[int] = None
-    ):
-        n_sequences = max(0, n_sequences) if n_sequences is not None else 1 << 63
-        n_bytes = max(0, n_bytes) if n_bytes is not None else 1 << 63
-
-        while len(self._lru) > n_sequences:
-            model, tokens = self._lru.pop()
-            self._delete(model, tokens)
-        while self._n_bytes > n_bytes:
-            model, tokens = self._lru.pop()
-            self._delete(model, tokens)
-
-    def log_cache_stats(self):
-        ncaches, nbytes = len(self), self.nbytes
-        ntok = (
-            len(self._lru._lru_checkpoints[-1][1])
-            if len(self._lru._lru_checkpoints) > 0
-            else 0
-        )
-        logging.info(
-            f"KV Caches: {ncaches} seq, {nbytes / 1e9:.2f} GB, latest user cache {ntok} tokens"
-        )
-
-
 @dataclass
 class ModelDescription:
     model: str
@@ -655,6 +464,11 @@ class ResponseGenerator:
     def join(self):
         self._generation_thread.join()
 
+    def _log_cache_stats(self):
+        ncaches = len(self.prompt_cache)
+        nbytes = self.prompt_cache.nbytes
+        logging.info(f"KV Caches: {ncaches} seq, {nbytes / 1e9:.2f} GB")
+
     def _next_request(self, timeout=None):
         request = None
         if not self._is_distributed or self._rank == 0:
@@ -791,7 +605,7 @@ class ResponseGenerator:
                     current_model_key,
                     rs["cache_key"][:-prompt_end],
                     list(cache),
-                    checkpoint=True,
+                    cache_type="user",
                 )
 
         if self._is_distributed:
@@ -842,7 +656,7 @@ class ResponseGenerator:
                     )
                     rqueue.put(ctx)
 
-                    self.prompt_cache.log_cache_stats()
+                    self._log_cache_stats()
                     cache, rest = self.prompt_cache.fetch_nearest_cache(
                         current_model_key, prompt
                     )
@@ -1023,7 +837,7 @@ class ResponseGenerator:
             logits_processors = _make_logits_processors(args)
 
             # Load the KV cache
-            self.prompt_cache.log_cache_stats()
+            self._log_cache_stats()
             cache, rest = self.prompt_cache.fetch_nearest_cache(
                 self.model_provider.model_key, prompt
             )

tests/test_server.py

@@ -514,7 +514,7 @@ class TestLRUPromptCache(unittest.TestCase):
         self.assertEqual(c, [MockCache("test3")])
         self.assertEqual(t, [])
 
-        cache.insert_cache(model, [4, 5], [MockCache("test4")], checkpoint=True)
+        cache.insert_cache(model, [4, 5], [MockCache("test4")], cache_type="user")
         c, t = cache.fetch_nearest_cache(model, [2, 3])
         self.assertEqual(c, None)
         self.assertEqual(t, [2, 3])