feat : add radix prefix cache

- replace hash-only lookup with page-granular radix matching
- keep partial pages private and cache only materialized KV prefixes
- integrate completed-request caching and add radix behavior tests

astrai/inference/__init__.py

@@ -35,7 +35,7 @@ from astrai.inference.core import (
     KVCache,
     KVStorage,
     PagePool,
-    PrefixCache,
+    RadixCache,
     ReqToTokenPool,
     Task,
     TaskManager,
@@ -66,7 +66,7 @@ __all__ = [
     "KVCache",
     "KVStorage",
     "PagePool",
-    "PrefixCache",
+    "RadixCache",
     "ReqToTokenPool",
     "page_hash",
     "sample",

astrai/inference/core/__init__.py

@@ -5,7 +5,7 @@ from astrai.inference.core.cache import (
     KVCache,
     KVStorage,
     PagePool,
-    PrefixCache,
+    RadixCache,
     ReqToTokenPool,
     page_hash,
 )
@@ -18,7 +18,7 @@ __all__ = [
     "KVCache",
     "KVStorage",
     "PagePool",
-    "PrefixCache",
+    "RadixCache",
     "ReqToTokenPool",
     "page_hash",
     "Executor",

astrai/inference/core/cache.py

@@ -4,7 +4,7 @@ Layer 1 — KVStorage:      flat token-level K/V buffers [n_layers, size, H, D]
 Layer 2 — ReqToTokenPool:  index table [req_idx, pos] → physical token slot
 Layer 3 — Allocator:       slot/page allocation with ref-counting and LRU
 
-PagePool orchestrates all three plus PrefixCache (content addressing).
+PagePool orchestrates all three plus RadixCache (prefix addressing).
 KVCache is a pure dataclass passed to the model for direct buffer access.
 
 Two modes:
@@ -23,10 +23,12 @@ from torch import Tensor
 from astrai.inference.core.workspace import InferenceWorkspace
 
 
-def page_hash(token_ids: List[int], page_idx: int, page_size: int) -> int:
+def page_hash(
+    token_ids: List[int], page_idx: int, page_size: int, parent_hash: int = 0
+) -> int:
     start = page_idx * page_size
     end = min(start + page_size, len(token_ids))
-    h = 0
+    h = parent_hash
     for i in range(start, end):
         h = (h * 31 + token_ids[i]) & 0xFFFFFFFFFFFFFFFF
     return h
@@ -83,45 +85,96 @@ class Allocator:
                 self._lru.move_to_end(idx)
 
 
-class PrefixCache:
+class RadixNode:
-    """Hash-based prefix matching: maps page hashes to physical page indices."""
+    """A page-aligned edge in the CPU-side prefix radix."""
+
+    __slots__ = ("parent", "children", "page_idx", "tokens", "lock_ref")
+
+    def __init__(self, parent=None, tokens=(), page_idx=None):
+        self.parent = parent
+        self.children: Dict[tuple, "RadixNode"] = {}
+        self.page_idx = page_idx
+        self.tokens = tuple(tokens)
+        self.lock_ref = 0
+
+
+class RadixCache:
+    """Page-granular radix prefix index with exact token matching."""
 
     def __init__(self, page_size: int):
         self._page_size = page_size
+        self._root = RadixNode()
+        self._page_to_node: Dict[int, RadixNode] = {}
+        # Retained as an introspection-compatible map; matching never relies on
+        # this lossy value.
         self._page_to_hash: Dict[int, int] = {}
-        self._hash_to_page: Dict[int, int] = {}
         self._lock = threading.Lock()
 
     def evict(self, idx: int):
         with self._lock:
-            h = self._page_to_hash.pop(idx, None)
+            node = self._page_to_node.pop(idx, None)
-            if h is not None:
+            self._page_to_hash.pop(idx, None)
-                self._hash_to_page.pop(h, None)
+            if node is None:
+                return
+            node.page_idx = None
+            parent = node.parent
+            if parent is not None:
+                parent.children.pop(node.tokens, None)
 
     def has_page(self, idx: int) -> bool:
         with self._lock:
-            return idx in self._page_to_hash
+            return idx in self._page_to_node
 
     def lookup(self, token_ids: List[int]) -> List[int]:
         with self._lock:
             full_pages = len(token_ids) // self._page_size
             hits: List[int] = []
+            node = self._root
             for i in range(full_pages):
-                h = page_hash(token_ids, i, self._page_size)
+                start = i * self._page_size
-                p = self._hash_to_page.get(h)
+                page_tokens = tuple(token_ids[start : start + self._page_size])
-                if p is None:
+                child = node.children.get(page_tokens)
+                if child is None or child.page_idx is None:
                     break
-                hits.append(p)
+                hits.append(child.page_idx)
+                node = child
             return hits
 
     def record(self, page_idx: int, token_ids: List[int], logical_page_idx: int):
         with self._lock:
-            h = page_hash(token_ids, logical_page_idx, self._page_size)
+            full_pages = len(token_ids) // self._page_size
-            old_h = self._page_to_hash.pop(page_idx, None)
+            if logical_page_idx >= full_pages:
-            if old_h is not None:
+                return
-                self._hash_to_page.pop(old_h, None)
+            old = self._page_to_node.pop(page_idx, None)
-            self._page_to_hash[page_idx] = h
+            self._page_to_hash.pop(page_idx, None)
-            self._hash_to_page[h] = page_idx
+            if old is not None and old.parent is not None:
+                old.parent.children.pop(old.tokens, None)
+
+            node = self._root
+            for i in range(logical_page_idx + 1):
+                start = i * self._page_size
+                page_tokens = tuple(token_ids[start : start + self._page_size])
+                child = node.children.get(page_tokens)
+                if child is None:
+                    child = RadixNode(node, page_tokens)
+                    node.children[page_tokens] = child
+                node = child
+            if node.page_idx is not None and node.page_idx != page_idx:
+                replaced = node.page_idx
+                self._page_to_node.pop(replaced, None)
+                self._page_to_hash.pop(replaced, None)
+            node.page_idx = page_idx
+            self._page_to_node[page_idx] = node
+            self._page_to_hash[page_idx] = page_hash(
+                token_ids, logical_page_idx, self._page_size
+            )
+
+    def release(self, pages: List[int]) -> None:
+        with self._lock:
+            for page_idx in pages:
+                node = self._page_to_node.get(page_idx)
+                if node is not None and node.lock_ref:
+                    node.lock_ref -= 1
 
 
 class ReqToTokenPool:
@@ -223,7 +276,7 @@ class KVCache:
 class PagePool:
     """Top-level KV cache manager.
 
-    Combines KVStorage + ReqToTokenPool + Allocator + PrefixCache.
+    Combines KVStorage + ReqToTokenPool + Allocator + RadixCache.
 
     Args:
         n_layers: Number of transformer layers.
@@ -275,11 +328,11 @@ class PagePool:
                     i * max_seq_len, (i + 1) * max_seq_len, device=device
                 )
             self._alloc: Optional[Allocator] = None
-            self._prefix: Optional[PrefixCache] = None
+            self._prefix: Optional[RadixCache] = None
         else:
             n_pages = self.n_tokens // page_size
             self._alloc = Allocator(n_pages)
-            self._prefix = PrefixCache(page_size) if page_size > 1 else None
+            self._prefix = RadixCache(page_size) if page_size > 1 else None
             if self._prefix is not None:
                 self._alloc.on_evict = self._prefix.evict
 
@@ -430,6 +483,17 @@ class PagePool:
         for i in range(start_logical_page, min(full_pages, len(pages))):
             self._prefix.record(pages[i], prompt_ids, i)
 
+    def task_cacheable_ids(
+        self, task_id: str, prompt_ids: List[int], output_ids: List[int]
+    ):
+        """Return the sequence whose KV entries are already materialized.
+
+        The first sampled output is produced by prompt prefill, and the last
+        sampled output has not been decoded into KV yet.  Therefore the cache
+        can safely retain the prompt plus every output except the last one.
+        """
+        return list(prompt_ids) + list(output_ids[:-1])
+
     # ---- bind for forward ----
 
     def bind_tasks(

astrai/inference/core/scheduler.py

@@ -164,6 +164,13 @@ class InferenceScheduler:
             while not self._stop_event.is_set():
                 finished = self._task_mgr.remove_finished_tasks(stop_ids)
                 for task in finished:
+                    if task.status == TaskStatus.FINISHED:
+                        cache.task_record_hashes(
+                            task.task_id,
+                            cache.task_cacheable_ids(
+                                task.task_id, task.prompt_ids, task.output_ids
+                            ),
+                        )
                     cache.task_free(task.task_id)
 
                 active = self._task_mgr.get_active_tasks()

tests/inference/test_cache.py

@@ -6,7 +6,7 @@ from astrai.inference import (
     Allocator,
     KVStorage,
     PagePool,
-    PrefixCache,
+    RadixCache,
     ReqToTokenPool,
     page_hash,
 )
@@ -77,12 +77,12 @@ def test_allocator_inc_ref_and_free():
     assert alloc._refs[p] == 0
 
 
-# ---- PrefixCache ----
+# ---- RadixCache ----
 
 
 def test_prefix_cache_lookup_returns_hits():
     token_ids = list(range(256))
-    prefix = PrefixCache(64)
+    prefix = RadixCache(64)
     pages = [0, 1, 2, 3]
     for i, p in enumerate(pages):
         prefix.record(p, token_ids, i)
@@ -92,7 +92,7 @@ def test_prefix_cache_lookup_returns_hits():
 
 def test_prefix_cache_lookup_stops_at_first_miss():
     token_ids = list(range(256))
-    prefix = PrefixCache(64)
+    prefix = RadixCache(64)
     prefix.record(0, token_ids, 0)
     prefix.record(1, [99] * 64, 1)
     hits = prefix.lookup(token_ids)
@@ -102,14 +102,14 @@ def test_prefix_cache_lookup_stops_at_first_miss():
 
 def test_prefix_cache_ignores_partial_last_page():
     token_ids = list(range(100))
-    prefix = PrefixCache(64)
+    prefix = RadixCache(64)
     prefix.record(0, token_ids, 0)
     hits = prefix.lookup(token_ids)
     assert len(hits) == 1
 
 
 def test_prefix_cache_on_evict_clears_mappings():
-    prefix = PrefixCache(64)
+    prefix = RadixCache(64)
     prefix.record(0, list(range(64)), 0)
     assert 0 in prefix._page_to_hash
     prefix.evict(0)
@@ -117,12 +117,49 @@ def test_prefix_cache_on_evict_clears_mappings():
 
 
 def test_prefix_cache_has_page():
-    prefix = PrefixCache(64)
+    prefix = RadixCache(64)
     assert not prefix.has_page(0)
     prefix.record(0, list(range(64)), 0)
     assert prefix.has_page(0)
 
 
+def test_prefix_cache_does_not_reuse_page_without_parent_prefix():
+    prefix = RadixCache(2)
+    prefix.record(0, [1, 2, 3, 4], 0)
+    prefix.record(1, [1, 2, 3, 4, 5, 6], 1)
+    prefix.record(2, [9, 10, 5, 6], 0)
+    prefix.record(3, [9, 10, 5, 6, 7, 8], 1)
+    assert prefix.lookup([1, 2, 3, 4, 5, 6]) == [0, 1]
+    assert prefix.lookup([9, 10, 5, 6, 7, 8]) == [2, 3]
+
+
+def test_prefix_cache_shares_branch_prefix():
+    prefix = RadixCache(2)
+    prefix.record(0, [1, 2, 3, 4], 0)
+    prefix.record(1, [1, 2, 3, 4], 1)
+    prefix.record(2, [1, 2, 7, 8], 1)
+    assert prefix.lookup([1, 2, 3, 4]) == [0, 1]
+    assert prefix.lookup([1, 2, 7, 8]) == [0, 2]
+    prefix.evict(1)
+    assert prefix.lookup([1, 2, 3, 4]) == [0]
+    assert prefix.lookup([1, 2, 7, 8]) == [0, 2]
+
+
+def test_prefix_cache_does_not_record_partial_page():
+    prefix = RadixCache(4)
+    prefix.record(0, [1, 2, 3, 4, 5, 6], 0)
+    prefix.record(1, [1, 2, 3, 4, 5, 6], 1)
+    assert prefix.lookup([1, 2, 3, 4, 5, 6]) == [0]
+
+    prefix.record(1, [1, 2, 3, 4, 5, 6, 7, 8], 1)
+    assert prefix.lookup([1, 2, 3, 4, 5, 6, 7, 8]) == [0, 1]
+
+
+def test_page_pool_task_cacheable_ids_excludes_unmaterialized_tail():
+    pool = _make_paged_pool_ps64()
+    assert pool.task_cacheable_ids("missing", [1, 2], [3, 4]) == [1, 2, 3]
+
+
 # ---- ReqToTokenPool ----