[Auto Sync] Update evict_policy.py, radix_cache.py (20251120) (#13669)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: cctry <shiyang@x.ai>

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)

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

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

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]