chore: import upstream snapshot with attribution

python/tokenspeed/runtime/cache/prefix_cache.py

@@ -0,0 +1,647 @@
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""
+The prefix tree data structure for managing the KV cache.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import heapq
+import time
+from collections import defaultdict
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
+import torch
+
+from tokenspeed.runtime.cache.allocator import KVAllocator
+from tokenspeed.runtime.cache.base_prefix_cache import BasePrefixCache, MatchResult
+from tokenspeed.runtime.cache.evict_policy import (
+    EvictionStrategy,
+    FIFOStrategy,
+    FILOStrategy,
+    LFUStrategy,
+    LRUStrategy,
+    MRUStrategy,
+    PriorityStrategy,
+)
+from tokenspeed.runtime.cache.req_to_token_pool import ReqToTokenPool
+from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool
+from tokenspeed.runtime.utils import get_colorful_logger
+
+logger = get_colorful_logger(__name__)
+
+if TYPE_CHECKING:
+    from tokenspeed.runtime.engine.request import Req
+
+# Import KV events classes if needed
+try:
+    from tokenspeed.runtime.pd.kv_events import (
+        AllBlocksCleared,
+        BlockRemoved,
+        BlockStored,
+    )
+
+    KV_EVENTS_AVAILABLE = True
+except ImportError:
+    BlockStored = None
+    BlockRemoved = None
+    AllBlocksCleared = None
+    KV_EVENTS_AVAILABLE = False
+    logger.warning(
+        "KV Events module not available. "
+        "KV cache events will not be emitted even if enable_kv_cache_events=True."
+    )
+
+
+@dataclasses.dataclass
+class CacheInitParams:
+    disable: bool
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool_allocator: KVAllocator
+    page_size: int
+
+    token_to_kv_pool: BaseTokenToKVPool = None
+    tp_cache_group: torch.distributed.ProcessGroup | None = None
+    eviction_policy: str = "lru"
+    disable_finished_insert: bool = False
+
+    enable_metrics: bool = False
+    enable_kv_cache_events: bool = False
+
+
+class TreeNode:
+    counter = 0
+
+    def __init__(self, id: int | None = None, priority: int = 0):
+        self.children = defaultdict(TreeNode)
+        self.parent = None
+        self.key = None
+        self.value = None
+        self.lock_ref = 0
+        self.last_access_time = time.monotonic()
+        self.creation_time = time.monotonic()
+
+        self.id = TreeNode.counter if id is None else id
+        TreeNode.counter += 1
+
+        ##### KVStore #####
+        # host_value is token level, others are page level
+        self.hit_count = 0
+        # indicating the node is locked to protect from eviction
+        # incremented when the node is referenced by a storage operation
+        self.host_ref_counter = 0
+        # store the host indices of KV cache
+        self.host_value: torch.Tensor | None = None
+        # store hash values of each pages
+        self.hash_value: list[str] | None = None
+        # priority for priority-aware eviction
+        self.priority = priority
+
+    @property
+    def evicted(self):
+        return self.value is None
+
+    def __lt__(self, other: "TreeNode"):
+        return self.last_access_time < other.last_access_time
+
+
+def _key_match(key0: list, key1: list):
+    i = 0
+    for k0, k1 in zip(key0, key1):
+        if k0 != k1:
+            break
+        i += 1
+    return i
+
+
+class PrefixCache(BasePrefixCache):
+    def __init__(self, params: CacheInitParams):
+        self.disable = params.disable
+        self.req_to_token_pool = params.req_to_token_pool
+        self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
+        self.enable_kv_cache_events = params.enable_kv_cache_events
+        self.page_size = (
+            1
+            if params.token_to_kv_pool_allocator is None
+            else params.token_to_kv_pool_allocator.page_size
+        )
+        self.eviction_policy = params.eviction_policy.lower()
+        self.key_match_fn = _key_match
+
+        strategy_types: dict[str, type[EvictionStrategy]] = {
+            "lru": LRUStrategy,
+            "lfu": LFUStrategy,
+            "fifo": FIFOStrategy,
+            "mru": MRUStrategy,
+            "filo": FILOStrategy,
+            "priority": PriorityStrategy,
+        }
+        strategy_type = strategy_types.get(self.eviction_policy)
+        if strategy_type is None:
+            raise ValueError(
+                f"Unknown eviction policy: {self.eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority'."
+            )
+        self.eviction_strategy = strategy_type()
+        # Initialize KV event queue
+        self.kv_event_queue = []
+
+        # Warn if KV events are enabled but module is not available
+        if self.enable_kv_cache_events and not KV_EVENTS_AVAILABLE:
+            logger.warning(
+                "KV cache events are enabled (enable_kv_cache_events=True) "
+                "but KV Events module is not available. Events will not be emitted."
+            )
+
+        self.reset()
+
+    ##### Public API #####
+
+    def reset(self):
+        self.root_node = TreeNode()
+        self.root_node.key = []
+        self.root_node.value = []
+        self.root_node.host_value = []
+        self.root_node.lock_ref = 1
+        self.root_node.hash_value = []
+        self.evictable_size_ = 0
+        self.protected_size_ = 0
+        self.evictable_leaves: set[TreeNode] = set()
+
+        if self.enable_kv_cache_events and KV_EVENTS_AVAILABLE:
+            self.kv_event_queue.append(AllBlocksCleared())
+
+    def match_prefix(self, key: list, **kwargs) -> MatchResult:
+        """Find the matching prefix from the prefix tree.
+        Args:
+            key: A list of token IDs to find a matching prefix.
+        Returns:
+            A tuple of a tensor of matching prefix token IDs and
+            the last node that contains the prefix values. Note that
+            this API can modify the internal state of the prefix tree.
+            The last node creates a new child if the prefix is shorter
+            than the last node's value.
+        """
+
+        if self.disable or len(key) == 0:
+            return self._empty_match_result()
+
+        page_size = self.token_to_kv_pool_allocator.page_size
+        # Compatible with whether the incoming key is paged
+        if not isinstance(key[0], tuple):
+            full_page_num = len(key) // page_size
+            paged_token_ids = [
+                tuple(key[i * page_size : (i + 1) * page_size])
+                for i in range(0, full_page_num)
+            ]
+        else:
+            paged_token_ids = key
+        if len(paged_token_ids) == 0:
+            return self._empty_match_result()
+        value = []
+        last_node = [self.root_node]
+        self._match_prefix_helper(self.root_node, paged_token_ids, value, last_node)
+        if value and isinstance(value[0], list):
+            flat_value = [e for arr in value for e in arr]
+            value = torch.concat(flat_value)
+        elif value:
+            value = torch.concat(value)
+        else:
+            value = torch.tensor([], dtype=torch.int32)
+        return MatchResult(
+            device_indices=value,
+            last_device_node=last_node[0],
+            last_host_node=last_node[0],
+            device_prefix_length=len(value) * page_size,
+        )
+
+    def insert(self, key: list, value=None):
+        if self.disable:
+            return 0
+
+        if value is None:
+            value = [x for x in key]
+        return self._insert_helper(self.root_node, key, value)
+
+    def cache_finished_req(
+        self,
+        req: Req,
+        token_ids: list[int] | None = None,
+        delay_req_pool_release=False,
+    ):
+        """Cache request when it finishes."""
+        if self.disable:
+            alloced_len = self.req_to_token_pool.alloced_lens[req.req_pool_idx].item()
+            self.token_to_kv_pool_allocator.free_req_cache(
+                req.req_pool_idx, alloced_len
+            )
+            if not delay_req_pool_release:
+                self.req_to_token_pool.free(req.req_pool_idx)
+                req.req_pool_idx = None
+            return
+
+        if token_ids is None:
+            token_ids = (req.origin_input_ids + req.output_ids)[:-1]
+        # Prefix Cache takes one ref in memory pool
+        req_pool_idx = req.req_pool_idx
+        page_size = self.token_to_kv_pool_allocator.page_size
+        seq_len = len(token_ids)
+        # without last not full page
+        full_page_num = seq_len // page_size
+        paged_token_ids = [
+            tuple(token_ids[i * page_size : (i + 1) * page_size])
+            for i in range(0, full_page_num)
+        ]
+        page_ids = self.token_to_kv_pool_allocator.req_to_page[
+            req_pool_idx, :full_page_num
+        ].clone()
+        logger.debug("insert_page_ids=%s", page_ids)
+        _ = self.insert(paged_token_ids, page_ids)
+
+        # After insert, the tree changed so the match prefix ids changed
+        new_prefix_page_ids = self.match_prefix(paged_token_ids).device_indices
+        # new_prefix_page_ids is cached in tree, free the diff part in page_ids
+
+        if new_prefix_page_ids.numel() > 0:
+            self.token_to_kv_pool_allocator.free_with_diff(
+                new_prefix_page_ids, page_ids
+            )
+        self.token_to_kv_pool_allocator.free_extra_pages_not_cached(
+            req_pool_idx,
+            seq_len,
+            self.req_to_token_pool.alloced_lens[req_pool_idx].item(),
+        )
+        if not delay_req_pool_release:
+            self.req_to_token_pool.free(req.req_pool_idx)
+            req.req_pool_idx = None
+        # Remove req slot release the cache lock
+        if req.last_node is not None:
+            self.dec_lock_ref(req.last_node)
+            req.last_node = None
+        logger.debug(
+            "[cache_finished_req]\nold_prefix_page_ids=%s\nnew_prefix_page_ids=%s\nlen(token_ids)=%s \nalloced_lens=%s, self.evictable_size_=%r, self.protected_size_=%r, req.last_node=%r",
+            self.token_to_kv_pool_allocator.req_to_page[
+                req_pool_idx, : (seq_len + page_size - 1) // page_size
+            ].tolist(),
+            new_prefix_page_ids.tolist(),
+            len(token_ids),
+            self.req_to_token_pool.alloced_lens[req_pool_idx],
+            self.evictable_size_,
+            self.protected_size_,
+            req.last_node,
+        )
+
+    def cache_unfinished_req(self, req: Req, token_ids: list[int] | None = None):
+        """Cache request when it is unfinished."""
+
+        if self.disable:
+            return
+
+        if token_ids is None:
+            token_ids = req.fill_ids
+
+        page_size = self.token_to_kv_pool_allocator.page_size
+        req_pool_idx = req.req_pool_idx
+        seq_len = len(token_ids)
+        # without last not full page
+        full_page_num = seq_len // page_size
+        paged_token_ids = [
+            tuple(token_ids[i * page_size : (i + 1) * page_size])
+            for i in range(0, full_page_num)
+        ]
+
+        page_ids = self.token_to_kv_pool_allocator.req_to_page[
+            req_pool_idx, :full_page_num
+        ].clone()
+        _ = self.insert(paged_token_ids, page_ids)
+        # After insert, perform matching, use page_id in prefix tree to replace the allocated
+        # page before insert, release diff part, and write to req_to_token_pool
+        match_result = self.match_prefix(paged_token_ids)
+        new_prefix_page_ids, new_last_node = (
+            match_result.device_indices,
+            match_result.last_device_node,
+        )
+        if new_prefix_page_ids.numel() > 0:
+            diff = self.token_to_kv_pool_allocator.free_with_diff(
+                new_prefix_page_ids, page_ids
+            )
+            diff_idxs = torch.nonzero(diff)
+            self.token_to_kv_pool_allocator.req_to_page[
+                req.req_pool_idx, diff_idxs.squeeze()
+            ] = new_prefix_page_ids[diff]
+            # Fix requires prefix cache to store CPU page IDs alongside GPU ones.
+            self.token_to_kv_pool_allocator.req_to_page_cpu[
+                req.req_pool_idx, diff_idxs.squeeze().cpu()
+            ] = new_prefix_page_ids[diff].cpu()
+            token_level_offsets = torch.arange(
+                self.page_size, device=self.req_to_token_pool.device
+            )
+            indices_start_locs = diff_idxs * self.page_size
+            diff_slots_indices = (
+                indices_start_locs[:, None] + token_level_offsets
+            ).flatten()
+            new_slots_start_locs = new_prefix_page_ids[diff] * self.page_size
+            new_slots = (new_slots_start_locs[:, None] + token_level_offsets).flatten()
+            self.req_to_token_pool.req_to_token[
+                req.req_pool_idx, diff_slots_indices
+            ] = new_slots.to(torch.int32)
+        self.dec_lock_ref(req.last_node)
+        self.inc_lock_ref(new_last_node)
+        req.last_node = new_last_node
+
+    def pretty_print(self, start_str: str = ""):
+        logger.debug(
+            self._print_helper(
+                self.root_node, 0, start_str + f" #tokens: {self.total_size()} "
+            )
+        )
+
+    def total_size(self):
+        return self._total_size_helper(self.root_node)
+
+    def evict(self, num_tokens: int, evict_callback: Callable = None):
+        if self.disable:
+            return
+
+        heap = [
+            (self.eviction_strategy.get_priority(node), node)
+            for node in self.evictable_leaves
+        ]
+        heapq.heapify(heap)
+
+        num_evicted = 0
+        while num_evicted < num_tokens and heap:
+            _, x = heapq.heappop(heap)
+
+            # evictable_leaves only contains unlocked leaves so lock_ref > 0 can
+            # only happen for cascade parents pushed mid-loop; guard defensively.
+            if x.lock_ref > 0:
+                continue
+
+            self.token_to_kv_pool_allocator.append_to_later_free(x.value)
+            num_evicted += len(x.value)
+            self._delete_leaf(x)  # removes x from evictable_leaves, may add parent
+
+            # Push cascade parent onto the working heap so it can be evicted in
+            # this same call. _delete_leaf already added it to evictable_leaves.
+            parent = x.parent
+            if (
+                not parent.children
+                and parent.lock_ref == 0
+                and parent != self.root_node
+            ):
+                heapq.heappush(
+                    heap, (self.eviction_strategy.get_priority(parent), parent)
+                )
+
+        self.token_to_kv_pool_allocator.free_group_end()
+        return num_evicted
+
+    def inc_lock_ref(self, node: TreeNode):
+        if self.disable:
+            return 0
+
+        delta = 0
+        while node != self.root_node:
+            if node.lock_ref == 0:
+                self.evictable_size_ -= len(node.value)
+                self.protected_size_ += len(node.value)
+                delta -= len(node.value)
+                self.evictable_leaves.discard(node)
+            node.lock_ref += 1
+            node = node.parent
+        return delta
+
+    def dec_lock_ref(self, node: TreeNode):
+        if self.disable:
+            return 0
+
+        if node is None:
+            return 0
+
+        delta = 0
+        while node != self.root_node:
+            if node.lock_ref == 1:
+                self.evictable_size_ += len(node.value)
+                self.protected_size_ -= len(node.value)
+                delta += len(node.value)
+                if not node.children:
+                    self.evictable_leaves.add(node)
+            node.lock_ref -= 1
+            node = node.parent
+        return delta
+
+    def evictable_size(self):
+        return self.evictable_size_
+
+    def protected_size(self):
+        # protected size refers to the size of the cache that is locked
+        return self.protected_size_
+
+    ##### Internal Helper Functions #####
+
+    def _update_leaf_status(self, node: TreeNode) -> None:
+        if node == self.root_node:
+            return
+        if not node.children and node.lock_ref == 0:
+            self.evictable_leaves.add(node)
+        else:
+            self.evictable_leaves.discard(node)
+
+    def _empty_match_result(self):
+        return MatchResult(
+            device_indices=torch.tensor([], dtype=torch.int32),
+            device_prefix_length=0,
+            host_hit_length=0,
+            last_device_node=self.root_node,
+            last_host_node=self.root_node,
+        )
+
+    def _match_prefix_helper(
+        self, node: TreeNode, key: list, value, last_node: TreeNode
+    ):
+        node.last_access_time = time.monotonic()
+        if not key:
+            return
+
+        if key[0] in node.children:
+            child = node.children[key[0]]
+            prefix_len = _key_match(child.key, key)
+            if prefix_len < len(child.key):
+                new_node = self._split_node(child.key, child, prefix_len)
+                value.append(new_node.value)
+                last_node[0] = new_node
+            else:
+                value.append(child.value)
+                last_node[0] = child
+                self._match_prefix_helper(child, key[prefix_len:], value, last_node)
+
+    def _split_node(self, key, child: TreeNode, split_len: int):
+        # new_node -> child
+        new_node = TreeNode()
+        new_node.children = {key[split_len]: child}
+        new_node.parent = child.parent
+        new_node.lock_ref = child.lock_ref
+        new_node.key = child.key[:split_len]
+        new_node.value = child.value[:split_len]
+        child.parent = new_node
+        child.key = child.key[split_len:]
+        child.value = child.value[split_len:]
+        new_node.parent.children[key[0]] = new_node
+        return new_node
+
+    def _insert_helper(self, node: TreeNode, key: list, value):
+        node.last_access_time = time.monotonic()
+        if not key:
+            return 0
+
+        if key[0] in node.children:
+            child = node.children[key[0]]
+            prefix_len = _key_match(child.key, key)
+
+            if prefix_len == len(child.key):
+                if prefix_len == len(key):
+                    return prefix_len
+                else:
+                    key = key[prefix_len:]
+                    value = value[prefix_len:]
+                    return prefix_len + self._insert_helper(child, key, value)
+
+            new_node = self._split_node(child.key, child, prefix_len)
+            return prefix_len + self._insert_helper(
+                new_node, key[prefix_len:], value[prefix_len:]
+            )
+
+        if key:
+            new_node = TreeNode()
+            new_node.parent = node
+            new_node.key = key
+            new_node.value = value
+            node.children[key[0]] = new_node
+            self.evictable_size_ += len(value)
+            # New node is a leaf; parent may have transitioned from leaf to internal.
+            self._update_leaf_status(new_node)
+            self._update_leaf_status(node)
+
+            # Emit BlockStored event when new KV blocks are inserted
+            if self.enable_kv_cache_events:
+                self._record_store_event(new_node)
+        return 0
+
+    def _print_helper(self, node: TreeNode, indent: int, print_str: str) -> str:
+        for _, child in node.children.items():
+            print_str += "\n"
+            print_str += " " * indent
+            print_str += f"{child=} key_len={len(child.key)} "
+            print_str += f"value={child.value} {child.lock_ref=}"
+            print_str = self._print_helper(
+                child, indent=indent + 2, print_str=print_str
+            )
+        return print_str
+
+    def _delete_leaf(self, node):
+        del node.parent.children[node.key[0]]
+        self.evictable_size_ -= len(node.key)
+        self.evictable_leaves.discard(node)
+        # Parent may have become a childless leaf.
+        self._update_leaf_status(node.parent)
+
+        # Emit BlockRemoved event when KV blocks are removed
+        if self.enable_kv_cache_events:
+            self._record_remove_event(node)
+
+    def _total_size_helper(self, node: TreeNode):
+        if node.evicted:
+            return 0
+        x = len(node.value)
+        for child in node.children.values():
+            x += self._total_size_helper(child)
+        return x
+
+    def _collect_leaves(self):
+        ret_list = []
+        stack = [self.root_node]
+
+        while stack:
+            cur_node = stack.pop()
+            if not cur_node.children:
+                ret_list.append(cur_node)
+            else:
+                stack.extend(cur_node.children.values())
+
+        return ret_list
+
+    def _record_store_event(self, node: TreeNode):
+        """Record BlockStored event for a node."""
+        if (
+            not self.enable_kv_cache_events
+            or not KV_EVENTS_AVAILABLE
+            or node.key is None
+        ):
+            return
+
+        # TokenSpeed uses tuples as keys, where each tuple represents a page
+        # node.key is a list of tuples: [(token1, token2, ...), (token3, token4, ...), ...]
+        # Each tuple is already a page, so we iterate over them directly
+
+        parent_block_hash = None
+        if node.parent and node.parent != self.root_node and node.parent.key:
+            # Use the last page (tuple) of the parent
+            parent_block_hash = hash(node.parent.key[-1])
+
+        # Each element in node.key is already a page (tuple of tokens)
+        for page_tuple in node.key:
+            if not page_tuple:
+                continue
+
+            # Convert tuple to list for token_ids
+            token_ids = list(page_tuple)
+            block_hash = hash(page_tuple)
+
+            self.kv_event_queue.append(
+                BlockStored(
+                    block_hashes=[block_hash],
+                    parent_block_hash=parent_block_hash,
+                    token_ids=token_ids,
+                    block_size=len(token_ids),
+                    lora_id=None,
+                )
+            )
+            # Chain next chunk to this one
+            parent_block_hash = block_hash
+
+    def _record_remove_event(self, node: TreeNode):
+        """Record BlockRemoved event for a node."""
+        if (
+            not self.enable_kv_cache_events
+            or not KV_EVENTS_AVAILABLE
+            or node.key is None
+        ):
+            return
+
+        # Create BlockRemoved event for each page
+        for start in range(0, len(node.key), self.page_size):
+            page_tokens = node.key[start : start + self.page_size]
+            if not page_tokens:
+                continue
+            block_hash = hash(tuple(page_tokens))
+            self.kv_event_queue.append(BlockRemoved(block_hashes=[block_hash]))