# 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]))