Class ThreadSafeLRUPromptCache

backend/python/mlx/mlx_cache.py:42–266

Thread-safe LRU cache with prefix matching for prompt KV caches. This cache stores KV caches keyed by token sequences and supports: - Exact match: Return the cache for the exact token sequence - Shorter prefix match: Return a cache for a prefix of the tokens - Longer prefix mat

class ThreadSafeLRUPromptCache:
    """
    Thread-safe LRU cache with prefix matching for prompt KV caches.

    This cache stores KV caches keyed by token sequences and supports:
    - Exact match: Return the cache for the exact token sequence
    - Shorter prefix match: Return a cache for a prefix of the tokens
    - Longer prefix match: If a longer sequence is cached and can be trimmed
    - LRU eviction: When max_size is exceeded, evict least recently used

    Thread safety is provided via a threading.Lock that protects all
    cache operations.

    Args:
        max_size: Maximum number of cache entries (default: 10)
        can_trim_fn: Optional function to check if a cache can be trimmed
        trim_fn: Optional function to trim a cache
    """

    def __init__(
        self,
        max_size: int = 10,
        can_trim_fn: Optional[Any] = None,
        trim_fn: Optional[Any] = None,
    ):
        self.max_size = max_size
        self._cache = {}
        self._lru = deque()
        self._lock = threading.Lock()

        # Optional trim functions (for longer prefix reuse)
        self._can_trim_fn = can_trim_fn
        self._trim_fn = trim_fn

    def _search(self, model, tokens: List[int]) -> SearchResult:
        """
        Search the cache for a prompt cache. Return exact or close match.

        The cache is organized as a trie where each node is keyed by a token.
        This allows efficient prefix matching.
        """
        if model not in self._cache:
            return SearchResult(model, None, None, None, 0)

        current = self._cache[model]
        last_cache_index = -1
        index = 0

        # Traverse the trie following the token sequence
        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 SearchResult(model, tuple(tokens), None, None, 0)