perf(cache): index RAM cache longest-prefix lookups

CHANGELOG.md

@@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- perf(cache): index RAM cache longest-prefix lookups by @hakarrr in #2410
+
 - fix(ci): use supported Python versions on Ubuntu by @abetlen in #2407
 - feat: update llama.cpp to ggml-org/llama.cpp@7d7f02875
 

llama_cpp/llama_cache.py

@@ -1,3 +1,4 @@
+from bisect import bisect_left
 import sys
 from abc import ABC, abstractmethod
 from typing import (
@@ -46,6 +47,54 @@ def __setitem__(
         raise NotImplementedError
 
 
+class _LlamaCachePrefixIndex:
+    """Linear-storage sorted-key index with cache-order tie breaking."""
+
+    def __init__(self, cache_state):
+        self.cache_state = cache_state
+        self.order = ()
+        self.keys = []
+        self.ranks = []
+        self.clock = 0
+
+    def sync(self):
+        order = tuple(self.cache_state)
+        if order != self.order:
+            pairs = sorted((key, rank) for rank, key in enumerate(order) if key)
+            self.keys = [key for key, rank in pairs]
+            self.ranks = [rank for key, rank in pairs]
+            self.order = order
+            self.clock = len(order)
+
+    def find(self, query):
+        self.sync()
+        # A lexicographic neighbor attains the maximum shared prefix.
+        index = bisect_left(self.keys, query)
+        best = 0
+        for position in (index - 1, index):
+            if 0 <= position < len(self.keys):
+                length = 0
+                for a, b in zip(self.keys[position], query):
+                    if a != b:
+                        break
+                    length += 1
+                best = max(best, length)
+        if not best:
+            return None
+        prefix = query[:best]
+        start = bisect_left(self.keys, prefix)
+        stop = bisect_left(self.keys, prefix[:-1] + (int(prefix[-1]) + 1,))
+        # All keys in this interval tie; preserve the original LRU winner.
+        winner = min(range(start, stop), key=self.ranks.__getitem__)
+        return self.keys[winner]
+
+    def touch(self, key):
+        position = bisect_left(self.keys, key)
+        self.ranks[position] = self.clock
+        self.clock += 1
+        self.order = tuple(self.cache_state)
+
+
 class LlamaRAMCache(BaseLlamaCache):
     """Cache for a llama.cpp model using RAM."""
 
@@ -55,6 +104,7 @@ def __init__(self, capacity_bytes: int = (2 << 30)):
         self.cache_state: OrderedDict[Tuple[int, ...], "llama_cpp.llama.LlamaState"] = (
             OrderedDict()
         )
+        self._prefix_index = _LlamaCachePrefixIndex(self.cache_state)
 
     @property
     def cache_size(self):
@@ -64,17 +114,19 @@ def _find_longest_prefix_key(
         self,
         key: Tuple[int, ...],
     ) -> Optional[Tuple[int, ...]]:
-        min_len = 0
-        min_key = None
-        keys = (
-            (k, llama_cpp.llama.Llama.longest_token_prefix(k, key))
-            for k in self.cache_state.keys()
-        )
-        for k, prefix_len in keys:
-            if prefix_len > min_len:
-                min_len = prefix_len
-                min_key = k
-        return min_key
+        if len(self.cache_state) <= 1:
+            if (
+                self._prefix_index.order
+                or self._prefix_index.cache_state is not self.cache_state
+            ):
+                self._prefix_index = _LlamaCachePrefixIndex(self.cache_state)
+            for candidate in self.cache_state:
+                if llama_cpp.llama.Llama.longest_token_prefix(candidate, key):
+                    return candidate
+            return None
+        if self._prefix_index.cache_state is not self.cache_state:
+            self._prefix_index = _LlamaCachePrefixIndex(self.cache_state)
+        return self._prefix_index.find(key)
 
     def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState":
         key = tuple(key)
@@ -83,6 +135,8 @@ def __getitem__(self, key: Sequence[int]) -> "llama_cpp.llama.LlamaState":
             raise KeyError("Key not found")
         value = self.cache_state[_key]
         self.cache_state.move_to_end(_key)
+        if len(self.cache_state) > 1:
+            self._prefix_index.touch(_key)
         return value
 
     def __contains__(self, key: Sequence[int]) -> bool:

tests/test_llama_cache.py

@@ -0,0 +1,194 @@
+import random
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+
+from llama_cpp import Llama, LlamaState
+from llama_cpp.llama_cache import LlamaRAMCache
+
+
+def state(tokens, size=1):
+    return LlamaState(
+        input_ids=np.array(tokens, dtype=np.intc),
+        scores=np.empty((0, 0), dtype=np.single),
+        n_tokens=len(tokens),
+        llama_state=bytes(size),
+        llama_state_size=size,
+        seed=0,
+    )
+
+
+def scan(keys, query):
+    best = None
+    longest = 0
+    for key in keys:
+        length = Llama.longest_token_prefix(key, query)
+        if length > longest:
+            best, longest = key, length
+    return best
+
+
+@pytest.mark.parametrize(
+    "keys,query,expected",
+    [
+        ([], (1,), None),
+        ([()], (), None),
+        ([(1, 2)], (), None),
+        ([(1, 2)], (9,), None),
+        ([(1, 2)], (1,), (1, 2)),
+        ([(1, 2)], (1, 2, 3), (1, 2)),
+        ([(1, 2, 3), (1, 2, 4)], (1, 2, 5), (1, 2, 3)),
+        ([(1, 2), (1, 2, 3)], (1, 2, 4), (1, 2)),
+        ([(1, 2, 3), (1, 2)], (1, 2, 4), (1, 2, 3)),
+        ([(1, 2, 3), (1, 2, 4)], (1, 2, 4, 5), (1, 2, 4)),
+        ([(1, 2, 3), (1, 2, 4)], (1, 9), (1, 2, 3)),
+    ],
+)
+def test_ram_cache_longest_shared_prefix(keys, query, expected):
+    cache = LlamaRAMCache()
+    values = {}
+    for key in keys:
+        values[key] = state(key)
+        cache[key] = values[key]
+    assert cache._find_longest_prefix_key(query) == expected
+    assert (list(query) in cache) == (expected is not None)
+    if expected is None:
+        with pytest.raises(KeyError):
+            cache[list(query)]
+    else:
+        assert cache[np.array(query, dtype=np.intc)] is values[expected]
+
+
+def test_ram_cache_ties_follow_lru_order():
+    cache = LlamaRAMCache()
+    a, b, c = (1, 2, 3), (1, 2, 4), (1, 5)
+    for key in (a, b, c):
+        cache[key] = state(key)
+    assert cache._find_longest_prefix_key((1, 2)) == a
+    cache[a]
+    assert cache._find_longest_prefix_key((1, 2)) == b
+    assert cache._find_longest_prefix_key((1,)) == b
+    cache[b] = state(b)
+    assert cache._find_longest_prefix_key((1,)) == c
+    assert cache._find_longest_prefix_key((1, 2)) == a
+
+
+def test_ram_cache_eviction_removes_obsolete_prefixes():
+    cache = LlamaRAMCache(capacity_bytes=3)
+    for key in [(1, 2, 3), (1, 2, 4), (1, 2), (9,)]:
+        cache[key] = state(key)
+    assert (1, 2, 3) not in cache.cache_state
+    assert cache._find_longest_prefix_key((1, 2, 3)) == (1, 2, 4)
+    cache[(8,)] = state((8,), size=2)
+    assert list(cache.cache_state) == [(9,), (8,)]
+    assert (1,) not in cache
+    cache[(7,)] = state((7,), size=4)
+    assert not cache.cache_state
+    assert (7,) not in cache
+    cache[(1,)] = state((1,))
+    assert (1, 2) in cache
+
+
+@pytest.mark.parametrize("seed", [0, 1, 2, 3, 280])
+def test_ram_cache_matches_scan_during_updates_and_evictions(seed):
+    rng = random.Random(seed)
+    cache = LlamaRAMCache(capacity_bytes=16)
+    reference = OrderedDict()
+    for _ in range(3000):
+        key = tuple(rng.randrange(5) for _ in range(rng.randrange(12)))
+        operation = rng.randrange(3)
+        if operation == 0:
+            value = state(key, size=rng.randrange(1, 5))
+            if key in reference:
+                del reference[key]
+            reference[key] = value
+            while sum(v.llama_state_size for v in reference.values()) > 16:
+                reference.popitem(last=False)
+            cache[key] = value
+        else:
+            expected = scan(reference, key)
+            assert cache._find_longest_prefix_key(key) == expected
+            assert (key in cache) == (expected is not None)
+            if operation == 1 and expected is not None:
+                assert cache[key] is reference[expected]
+                reference.move_to_end(expected)
+        assert list(cache.cache_state) == list(reference)
+        assert cache.cache_size == sum(v.llama_state_size for v in reference.values())
+
+
+def test_ram_cache_long_shared_edge_with_bounded_index():
+    cache = LlamaRAMCache()
+    prefix = tuple(range(10000))
+    for suffix in range(32):
+        key = prefix + (suffix,)
+        cache[key] = state(key)
+    assert cache._find_longest_prefix_key(prefix + (31, 99)) == prefix + (31,)
+
+
+def test_ram_cache_nested_prefixes_survive_parent_eviction_and_replacement():
+    cache = LlamaRAMCache(capacity_bytes=64)
+    keys = [tuple(range(128 + n)) for n in range(64)]
+    for key in keys:
+        cache[key] = state(key)
+    for key in keys[::2]:
+        cache[key]
+    for n in range(128):
+        key = (999, n)
+        cache[key] = state(key)
+        for query in (keys[n % 64], (999,), (1234,)):
+            assert cache._find_longest_prefix_key(query) == scan(
+                cache.cache_state, query
+            )
+
+    cache._prefix_index.sync()
+    assert len(cache._prefix_index.keys) == len(cache.cache_state)
+    assert len(cache._prefix_index.ranks) == len(cache.cache_state)
+
+
+def test_ram_cache_public_mapping_mutations_remain_visible():
+    cache = LlamaRAMCache()
+    a, b = (1, 2), (1, 3)
+    cache[a] = state(a)
+    cache[b] = state(b)
+    assert cache._find_longest_prefix_key((1,)) == a
+    cache.cache_state.move_to_end(a)
+    assert cache._find_longest_prefix_key((1,)) == b
+    cache.cache_state.move_to_end(a, last=False)
+    assert cache._find_longest_prefix_key((1,)) == a
+    cache.cache_state.pop(a)
+    assert cache._find_longest_prefix_key((1,)) == b
+    cache.cache_state.clear()
+    assert (1,) not in cache
+    with pytest.raises(KeyError):
+        cache[(1,)]
+    cache.cache_state = OrderedDict([(a, state(a))])
+    assert cache[(1,)] is cache.cache_state[a]
+
+
+def test_ram_cache_replacing_public_mapping_rebuilds_index():
+    cache = LlamaRAMCache()
+    for key in ((1, 2), (1, 3)):
+        cache[key] = state(key)
+    assert (1,) in cache
+    a, b = (9, 3), (9, 2)
+    cache.cache_state = OrderedDict([(a, state(a)), (b, state(b))])
+    assert cache._find_longest_prefix_key((9,)) == a
+    assert (1,) not in cache
+    cache.cache_state.clear()
+    assert (9,) not in cache
+    assert not cache._prefix_index.order
+    assert not cache._prefix_index.keys
+
+
+@pytest.mark.parametrize("dtype", [np.int32, np.int64])
+def test_ram_cache_numpy_token_dtype_boundary(dtype):
+    cache = LlamaRAMCache()
+    boundary = np.iinfo(dtype).max
+    a, b = (int(boundary), 1), (int(boundary), 2)
+    cache[a] = state(())
+    cache[b] = state(())
+    query = np.array([boundary, 3], dtype=dtype)
+    assert cache._find_longest_prefix_key(tuple(query)) == a
+    assert cache[query] is cache.cache_state[a]
+    assert cache._find_longest_prefix_key(tuple(query)) == b