[hakarrr]

Related to #280. This is a dependency-free sorted-key alternative for multi-state RAM caches with long shared or nested prompt prefixes, not a pygtrie implementation.

Workload limits and tradeoffs

This is not a general faster-cache claim. In the synthetic 64-state, 2048-token benchmark below, disjoint-key get regresses from 339 to 570 microseconds (about 68% slower), and alternating overwrite/get regresses from 826 to 1063 microseconds (about 29% slower). Shared/nested prefixes benefit, but this tradeoff may not suit every workload. No model-level inference speedup was measured.

Every lookup takes an O(N) snapshot of public mapping order. Tie selection can scan O(N) ranks. Key/order changes trigger a lazy O(N log N) comparison sort, with token-length-dependent comparisons. Auxiliary O(N) key references/ranks are outside capacity_bytes, which still measures serialized state bytes. Removed states are released normally; indexed key tuples may remain referenced until a subsequent lookup.

Changes

• Use lexicographic neighbors to determine the maximum common-prefix length, then select the oldest matching rank from its contiguous key interval.
• Preserve first-in-LRU-order tie selection, access reordering, overwrites and capacity eviction.
• Detect direct cache_state key/order changes and mapping replacement, including clear/pop/move_to_end. Keep a simple path for zero/one state.
• Add scan-equivalence, long/nested prefix, LRU/eviction, public mapping and NumPy integer-boundary regression tests. Disk cache is unchanged; no dependency added.

Validation

Python 3.11.17, Linux x86_64 CPU, GCC 11.4, CMake 3.22.1; exact pinned llama.cpp 7d7f0287530bc8981690f361d8c1627dbae1c99a, built with GGML_NATIVE=OFF, GGML_OPENMP=OFF, GGML_BLAS=OFF.

• Full python -m pytest tests -q: 50 passed, no skips, including the existing native/model-backed tests. Unchanged original suite: 26 passed.
• ruff check llama_cpp tests and ruff format --check llama_cpp tests: passed.
• git diff --check: passed.
• Additional private scan-oracle probe: 100,000 mixed operations including public mapping clear/pop/reorder passed. This is separate from the checked-in tests.
• Fetched the published branch, verified all three changed files byte-for-byte, then reran the full 50-test suite and lint/format there.

Synthetic cache microbenchmark

Microseconds per operation, median of 7 batches, CPU0 affinity. Real LlamaRAMCache and LlamaState; 1-byte serialized states, 2048-token keys, 64 entries. This deliberately isolates Python cache overhead, not inference or large-state copying. The alternating case rotates overwrites through the stored keys before each get, so it includes lazy rebuild cost.

Key shape	get: scan	get: index	overwrite/get: scan	overwrite/get: index
shared	2604.3	586.4	3104.3	1260.6
disjoint	339.4	569.8	825.6	1062.7
nested	4892.9	602.3	5412.9	1433.5

Index-only nested-key allocation tracing at 64/128/256/512/1024 keys was 1576/3112/6376/19624/47336 bytes. These are Python traced allocations after forcing lazy index creation, excluding the preallocated token/state payloads, not process RSS.

Reproduce the same-workload timing comparison

Run from this branch with its native library built and test dependencies installed. The baseline module is extracted from the parent before this PR; only this script's baseline import is redirected. CPU0 was used for the reported run; choose an available CPU on another machine.

git show 05c124e33f7f33b7e66ef887274ad9ee6775099f:llama_cpp/llama_cache.py > /tmp/llama_cache_baseline.py
OPENBLAS_NUM_THREADS=1 taskset -c 0 python /tmp/cache_benchmark.py

Save this as /tmp/cache_benchmark.py:

import importlib.util
import json
import os
import platform
import random
import statistics
import timeit
import tracemalloc

import numpy as np
from llama_cpp import LlamaState
from llama_cpp.llama_cache import LlamaRAMCache

spec = importlib.util.spec_from_file_location('llama_cpp.baseline_cache', '/tmp/llama_cache_baseline.py')
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
Baseline = module.LlamaRAMCache

def value(key):
    return LlamaState(np.array(key, dtype=np.intc), np.empty((0,0), dtype=np.single),len(key),b'x',1,0)

def make(cls, pairs):
    cache=cls()
    for key,state in pairs:cache[key]=state
    return cache

def measure(fn, number=100):
    return statistics.median(timeit.repeat(fn, number=number, repeat=7))*1e6/number

def mixed(cache, keys, pairs, query):
    counter = [0]
    def run():
        index = counter[0] % len(keys)
        counter[0] += 1
        cache[keys[index]] = pairs[index][1]
        cache[query]
    return run

rows=[]
for shape in ['shared', 'disjoint', 'nested']:
    for count in [1, 16, 64]:
        rng=random.Random(280)
        common=tuple(rng.randrange(32000) for _ in range(2048))
        if shape=='shared':keys=[common[:1024]+(i,)+common[1025:] for i in range(count)]
        elif shape=='disjoint':keys=[(i,)+common[1:] for i in range(count)]
        else:keys=[common[:2048-count+i+1] for i in range(count)]
        pairs=[(k,value(k)) for k in keys]
        query=keys[-1]+(99999,)
        base=make(Baseline,pairs); head=make(LlamaRAMCache,pairs)
        assert base._find_longest_prefix_key(query)==head._find_longest_prefix_key(query)
        tracemalloc.start(); tree=make(LlamaRAMCache,pairs); tree._find_longest_prefix_key(query); tree_mem=tracemalloc.get_traced_memory()[0];tracemalloc.stop()
        tracemalloc.start(); old=make(Baseline,pairs); old_mem=tracemalloc.get_traced_memory()[0];tracemalloc.stop()
        rows.append(dict(shape=shape,count=count,lookup_base_us=measure(lambda:base._find_longest_prefix_key(query)),lookup_head_us=measure(lambda:head._find_longest_prefix_key(query)),get_base_us=measure(lambda:base[query]),get_head_us=measure(lambda:head[query]),set_base_us=measure(lambda:base.__setitem__(keys[-1],pairs[-1][1]),20),set_head_us=measure(lambda:head.__setitem__(keys[-1],pairs[-1][1]),20),set_get_base_us=measure(mixed(base,keys,pairs,query),20),set_get_head_us=measure(mixed(head,keys,pairs,query),20),extra_traced_bytes=tree_mem-old_mem))
print(json.dumps(dict(python=platform.python_version(),platform=platform.platform(),affinity=sorted(os.sched_getaffinity(0)),rows=rows),indent=2))