diff --git a/python/sglang/srt/mem_cache/evict_policy.py b/python/sglang/srt/mem_cache/evict_policy.py index dfc7601e7..687593418 100644 --- a/python/sglang/srt/mem_cache/evict_policy.py +++ b/python/sglang/srt/mem_cache/evict_policy.py @@ -36,3 +36,11 @@ class MRUStrategy(EvictionStrategy): class FILOStrategy(EvictionStrategy): def get_priority(self, node: "TreeNode") -> float: return -node.creation_time + + +class PriorityStrategy(EvictionStrategy): + """Priority-aware eviction: lower priority values evicted first, then LRU within same priority.""" + + def get_priority(self, node: "TreeNode") -> Tuple[int, float]: + # Return (priority, last_access_time) so lower priority nodes are evicted first + return (node.priority, node.last_access_time) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 11665d666..1feee6f51 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -786,7 +786,7 @@ class HiRadixCache(RadixCache): child_key = self.get_child_key_fn(key) if len(key): - new_node = TreeNode() + new_node = TreeNode(priority=node.priority) new_node.parent = node new_node.key = key new_node.value = None @@ -823,7 +823,7 @@ class HiRadixCache(RadixCache): def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): # child node split into new_node -> child - new_node = TreeNode() + new_node = TreeNode(priority=child.priority) new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.parent = child.parent new_node.lock_ref = child.lock_ref @@ -848,7 +848,15 @@ class HiRadixCache(RadixCache): new_node.parent.children[self.get_child_key_fn(key)] = new_node return new_node - def insert(self, key: RadixKey, value=None, chunked=False): + def insert( + self, + key: RadixKey, + value=None, + chunked: bool = False, + priority: int | None = None, + ): + if priority is None: + priority = 0 key.token_ids = self.key_convert_fn(key.token_ids) if len(key) == 0: @@ -865,6 +873,7 @@ class HiRadixCache(RadixCache): while len(key) > 0 and child_key in node.children.keys(): node = node.children[child_key] node.last_access_time = time.monotonic() + node.priority = max(node.priority, priority) prefix_len = self.key_match_fn(node.key, key) if prefix_len == len(node.key): @@ -879,6 +888,8 @@ class HiRadixCache(RadixCache): else: # partial match, split the node new_node = self._split_node(node.key, node, prefix_len) + # shared-prefix node should also reflect max priority + new_node.priority = max(new_node.priority, priority) if new_node.evicted: new_node.value = value[:prefix_len] self.evictable_size_ += len(new_node.value) @@ -894,7 +905,7 @@ class HiRadixCache(RadixCache): child_key = self.get_child_key_fn(key) if len(key): - new_node = TreeNode() + new_node = TreeNode(priority=priority) new_node.parent = node new_node.key = key new_node.value = value diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 24a2eaf8b..754a0f004 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -20,6 +20,7 @@ The radix tree data structure for managing the KV cache. """ import heapq +import sys import time from collections import defaultdict from functools import lru_cache, partial @@ -41,6 +42,7 @@ from sglang.srt.mem_cache.evict_policy import ( LFUStrategy, LRUStrategy, MRUStrategy, + PriorityStrategy, ) from sglang.srt.mem_cache.memory_pool import ReqToTokenPool @@ -76,7 +78,7 @@ class TreeNode: counter = 0 - def __init__(self, id: Optional[int] = None): + def __init__(self, id: Optional[int] = None, priority: int = 0): self.children = defaultdict(TreeNode) self.parent: TreeNode = None self.key: RadixKey = None @@ -93,6 +95,8 @@ class TreeNode: self.host_value: Optional[torch.Tensor] = None # store hash values of each pages self.hash_value: Optional[List[str]] = None + # priority for priority-aware eviction + self.priority = priority self.id = TreeNode.counter if id is None else id TreeNode.counter += 1 @@ -195,6 +199,7 @@ class RadixCache(BasePrefixCache): enable_kv_cache_events: bool = False, eviction_policy: str = "lru", is_eagle: bool = False, + disable_finished_insert: bool = False, ): self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator @@ -203,6 +208,7 @@ class RadixCache(BasePrefixCache): self.enable_kv_cache_events = enable_kv_cache_events self.kv_event_queue = [] self.is_eagle = is_eagle + self.disable_finished_insert = disable_finished_insert if enable_metrics: self.init_metrics_collector() @@ -234,16 +240,19 @@ class RadixCache(BasePrefixCache): self.eviction_strategy: EvictionStrategy = MRUStrategy() elif eviction_policy.lower() == "filo": self.eviction_strategy: EvictionStrategy = FILOStrategy() + elif eviction_policy.lower() == "priority": + self.eviction_strategy: EvictionStrategy = PriorityStrategy() else: raise ValueError( - f"Unknown eviction policy: {eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo'." + f"Unknown eviction policy: {eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority'." ) self.reset() ##### Public API ##### def reset(self): - self.root_node = TreeNode() + # Initialize root with minimum priority so any real priority overrides it + self.root_node = TreeNode(priority=-sys.maxsize) self.root_node.key = RadixKey(token_ids=[], extra_key=None) self.root_node.value = [] self.root_node.host_value = [] @@ -324,7 +333,7 @@ class RadixCache(BasePrefixCache): last_host_node=last_node, ) - def insert(self, key: RadixKey, value=None, chunked=False): + def insert(self, key: RadixKey, value=None, chunked=False, priority: int = 0): if self.disable: return 0 @@ -337,10 +346,14 @@ class RadixCache(BasePrefixCache): # Make sure the value len equal to the EAGLE bigram key len value = value[: len(key)] - return self._insert_helper(self.root_node, key, value) + return self._insert_helper(self.root_node, key, value, priority) def cache_finished_req(self, req: Req, is_insert: bool = True): """Cache request when it finishes.""" + # In deterministic mode, disable finished request insertion to radix cache + if self.disable_finished_insert: + is_insert = False + committed_kv_len = req.pop_committed_kv_cache() if self.disable: kv_indices = self.req_to_token_pool.req_to_token[ @@ -379,9 +392,11 @@ class RadixCache(BasePrefixCache): # Radix Cache takes one ref in memory pool if is_insert: + priority = getattr(req, "priority", 0) or 0 new_prefix_len = self.insert( RadixKey(token_ids[:page_aligned_token_len], req.extra_key), page_aligned_kv_indices, + priority=priority, ) # Free the duplicates that were already in the tree self.token_to_kv_pool_allocator.free( @@ -435,10 +450,12 @@ class RadixCache(BasePrefixCache): old_prefix_len -= 1 # Radix Cache takes one ref in memory pool + priority = getattr(req, "priority", 0) or 0 new_prefix_len = self.insert( RadixKey(page_aligned_token_ids, req.extra_key), page_aligned_kv_indices, chunked=chunked, + priority=priority, ) self.token_to_kv_pool_allocator.free(kv_indices[old_prefix_len:new_prefix_len]) @@ -590,8 +607,9 @@ class RadixCache(BasePrefixCache): def _split_node(self, key: RadixKey, child: TreeNode, split_len: int): # new_node -> child + # New node inherits child's priority (represents shared prefix) self._record_remove_event(child) - new_node = TreeNode() + new_node = TreeNode(priority=child.priority) new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.parent = child.parent new_node.lock_ref = child.lock_ref @@ -607,9 +625,14 @@ class RadixCache(BasePrefixCache): return new_node - def _insert_helper(self, node: TreeNode, key: RadixKey, value): + def _insert_helper(self, node: TreeNode, key: RadixKey, value, priority: int = 0): + # Convert None priority to 0 + if priority is None: + priority = 0 access_time = time.monotonic() node.last_access_time = access_time + # Update priority along the path (take max to propagate higher priority) + node.priority = max(node.priority, priority) if len(key) == 0: return 0 @@ -626,13 +649,16 @@ class RadixCache(BasePrefixCache): if prefix_len < len(node.key): new_node = self._split_node(node.key, node, prefix_len) + new_node.priority = max(new_node.priority, priority) node = new_node + else: + node.priority = max(node.priority, priority) if len(key): child_key = self.get_child_key_fn(key) if len(key): - new_node = TreeNode() + new_node = TreeNode(priority=priority) new_node.parent = node new_node.key = key new_node.value = value diff --git a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py index 7b158909a..820653b4b 100644 --- a/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py +++ b/python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py @@ -196,7 +196,7 @@ class LMCRadixCache(RadixCache): if num_retrieved > 0: fetched = num_retrieved - prefix_pad - new_node = TreeNode() + new_node = TreeNode(priority=last_node.priority) start = value.numel() end = start + fetched new_node.key = key[start:end]