[UnifiedTree]: Support eviction priority (#26549)

python/sglang/srt/mem_cache/radix_cache.py

@@ -47,17 +47,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchResult,
 )
 from sglang.srt.mem_cache.events import KVCacheEventMixin
-from sglang.srt.mem_cache.evict_policy import (
+from sglang.srt.mem_cache.utils import get_eviction_strategy, split_node_hash_value
-    EvictionStrategy,
-    FIFOStrategy,
-    FILOStrategy,
-    LFUStrategy,
-    LRUStrategy,
-    MRUStrategy,
-    PriorityStrategy,
-    SLRUStrategy,
-)
-from sglang.srt.mem_cache.utils import split_node_hash_value
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -292,25 +282,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache):
         else:
             self.device = torch.device("cpu")
 
-        if self.eviction_policy == "lru":
+        self.eviction_strategy = get_eviction_strategy(self.eviction_policy)
-            self.eviction_strategy: EvictionStrategy = LRUStrategy()
-        elif self.eviction_policy == "lfu":
-            self.eviction_strategy: EvictionStrategy = LFUStrategy()
-        elif self.eviction_policy == "fifo":
-            self.eviction_strategy: EvictionStrategy = FIFOStrategy()
-        elif self.eviction_policy == "mru":
-            self.eviction_strategy: EvictionStrategy = MRUStrategy()
-        elif self.eviction_policy == "filo":
-            self.eviction_strategy: EvictionStrategy = FILOStrategy()
-        elif self.eviction_policy == "priority":
-            self.eviction_strategy: EvictionStrategy = PriorityStrategy()
-        elif self.eviction_policy == "slru":
-            self.eviction_strategy: EvictionStrategy = SLRUStrategy()
-
-        else:
-            raise ValueError(
-                f"Unknown eviction policy: {self.eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority', 'slru'."
-            )
 
         self.evictable_leaves = set()
         self.reset()

python/sglang/srt/mem_cache/unified_cache_components/full_component.py

@@ -127,8 +127,10 @@ class FullComponent(TreeComponent):
         self, params: EvictParams, tracker: dict[ComponentType, int]
     ) -> None:
         request = params.num_tokens
-        # Heap-based eviction from evictable_device_leaves, ordered by LRU.
+        heap = [
-        heap = [(n.last_access_time, n) for n in self.cache.evictable_device_leaves]
+            (self.cache.eviction_strategy.get_priority(n), n)
+            for n in self.cache.evictable_device_leaves
+        ]
         heapq.heapify(heap)
         ct = self.component_type
         while tracker[ct] < request and heap:
@@ -137,13 +139,19 @@ class FullComponent(TreeComponent):
                 continue
             self.cache._evict_device_leaf(x, tracker)
             if x.parent is not None and x.parent in self.cache.evictable_device_leaves:
-                heapq.heappush(heap, (x.parent.last_access_time, x.parent))
+                heapq.heappush(
+                    heap,
+                    (self.cache.eviction_strategy.get_priority(x.parent), x.parent),
+                )
 
     def drive_host_eviction(
         self, num_tokens: int, tracker: dict[ComponentType, int]
     ) -> None:
         """Evict host leaves to free KV host pool space."""
-        heap = [(n.last_access_time, n) for n in self.cache.evictable_host_leaves]
+        heap = [
+            (self.cache.eviction_strategy.get_priority(n), n)
+            for n in self.cache.evictable_host_leaves
+        ]
         heapq.heapify(heap)
         ct = self.component_type
         while tracker[ct] < num_tokens and heap:
@@ -152,7 +160,10 @@ class FullComponent(TreeComponent):
                 continue
             self.cache._evict_host_leaf(x, tracker)
             if x.parent is not None and x.parent in self.cache.evictable_host_leaves:
-                heapq.heappush(heap, (x.parent.last_access_time, x.parent))
+                heapq.heappush(
+                    heap,
+                    (self.cache.eviction_strategy.get_priority(x.parent), x.parent),
+                )
 
     def acquire_component_lock(
         self,

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import sys
 import threading
 import time
 from array import array
@@ -48,7 +49,11 @@ from sglang.srt.mem_cache.unified_cache_components import (
     TreeComponent,
     get_and_increase_time_counter,
 )
-from sglang.srt.mem_cache.utils import compute_node_hash_values, split_node_hash_value
+from sglang.srt.mem_cache.utils import (
+    compute_node_hash_values,
+    get_eviction_strategy,
+    split_node_hash_value,
+)
 from sglang.srt.observability.metrics_collector import StorageMetricsCollector
 from sglang.srt.session.streaming_session import StreamingSession
 
@@ -61,7 +66,7 @@ if TYPE_CHECKING:
 class UnifiedTreeNode:
     counter = 0
 
-    def __init__(self, tree_components: tuple[ComponentType, ...]):
+    def __init__(self, tree_components: tuple[ComponentType, ...], priority: int = 0):
         self.children = defaultdict(partial(UnifiedTreeNode, tree_components))
         self.parent: UnifiedTreeNode | None = None
         self.key: Optional[RadixKey] = None
@@ -71,8 +76,10 @@ class UnifiedTreeNode:
             ComponentData() for _ in range(_NUM_COMPONENT_TYPES)
         ]
         self.last_access_time = get_and_increase_time_counter()
+        self.creation_time = get_and_increase_time_counter()
         self.hash_value = None
         self.hit_count = 0
+        self.priority = priority
         self.lru_prev: list[UnifiedTreeNode | None] = [None] * (
             _NUM_COMPONENT_TYPES * 2
         )
@@ -234,6 +241,8 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         self.is_eagle = params.is_eagle
         self.enable_kv_cache_events = params.enable_kv_cache_events
         self.kv_event_queue = []
+        self.eviction_policy = params.eviction_policy.lower()
+        self.eviction_strategy = get_eviction_strategy(self.eviction_policy)
 
         if self.token_to_kv_pool_allocator:
             self.device = self.token_to_kv_pool_allocator.device
@@ -289,6 +298,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
     def _reset_full(self) -> None:
         """Full reset: destroy entire tree and all state."""
         self.root_node = UnifiedTreeNode(self.tree_components)
+        self.root_node.priority = -sys.maxsize
         self.root_node.key = RadixKey(array("q"), None)
         self.root_node.component_data[BASE_COMPONENT_TYPE].value = []
         self.root_node.hash_value = []
@@ -545,7 +555,10 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         insert_params = None
 
         if is_insert:
-            insert_params = InsertParams(prev_prefix_len=req.cache_protected_len)
+            insert_params = InsertParams(
+                prev_prefix_len=req.cache_protected_len,
+                priority=getattr(req, "priority", 0) or 0,
+            )
 
             # components prepare insert data + return effective cache_len
             effective_cache_len = len(token_ids)
@@ -611,7 +624,9 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
 
         # components prepare insert data + return effective cache_len
         insert_params = InsertParams(
-            prev_prefix_len=req.cache_protected_len, chunked=chunked
+            prev_prefix_len=req.cache_protected_len,
+            chunked=chunked,
+            priority=getattr(req, "priority", 0) or 0,
         )
         effective_cache_len = len(token_ids)
         for comp in self._components_tuple:
@@ -823,10 +838,12 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
     def _split_node(
         self, key: RadixKey, child: UnifiedTreeNode, split_len: int
     ) -> UnifiedTreeNode:
-        new_node = UnifiedTreeNode(self.tree_components)
+        new_node = UnifiedTreeNode(self.tree_components, priority=child.priority)
         new_node.children = {key[split_len:].child_key(self.page_size): child}
         new_node.parent = child.parent
         new_node.key = child.key[:split_len]
+        new_node.hit_count = child.hit_count
+        new_node.creation_time = child.creation_time
 
         self._for_each_component_lru(child, UnifiedLRUList.remove_node)
 
@@ -862,8 +879,9 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         parent: UnifiedTreeNode,
         key: RadixKey,
         value: torch.Tensor,
+        priority: int = 0,
     ) -> UnifiedTreeNode:
-        new_node = UnifiedTreeNode(self.tree_components)
+        new_node = UnifiedTreeNode(self.tree_components, priority=priority)
         new_node.parent = parent
         new_node.key = key
         new_node.component_data[BASE_COMPONENT_TYPE].value = value.clone()
@@ -900,7 +918,11 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         value: torch.Tensor,
         params: InsertParams,
     ) -> InsertResult:
+        priority = params.priority
+        if priority is None:
+            priority = 0
         self._touch_node(node)
+        node.priority = max(node.priority, priority)
         if len(key) == 0:
             return InsertResult(prefix_len=0, mamba_exist=True)
 
@@ -912,6 +934,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
             prefix_len = node.key.match(key, page_size=self.page_size)
             if prefix_len < len(node.key):
                 node = self._split_node(node.key, node, prefix_len)
+            node.priority = max(node.priority, priority)
 
             if node.evicted:
                 self._unevict_node_on_insert(node, value[:prefix_len])
@@ -970,7 +993,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
                 # cleanup_after_caching_req can free them properly.
                 self.token_to_kv_pool_allocator.free(value)
                 return InsertResult(prefix_len=total_prefix_length)
-            target_node = self._add_new_node(node, key, value)
+            target_node = self._add_new_node(node, key, value, priority=priority)
             is_new_leaf = True
         else:
             target_node = node
@@ -1030,7 +1053,7 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
                 result.inserted_host_node = node
             return result
 
-        new_node = UnifiedTreeNode(self.tree_components)
+        new_node = UnifiedTreeNode(self.tree_components, priority=node.priority)
         new_node.parent = node
         new_node.key = key
         new_node.hash_value = hash_value
@@ -1514,14 +1537,19 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
 
     def _inc_hit_count(self, node: UnifiedTreeNode, chunked: bool = False) -> None:
         """Increment hit count; trigger write_backup when threshold reached."""
-        if self.cache_controller is None:
-            return
         if node.evicted or chunked:
             return
-        if self.cache_controller.write_policy == "write_back":
+        if (
+            self.cache_controller is not None
+            and self.cache_controller.write_policy == "write_back"
+        ):
             return
         node.hit_count += 1
-        if not node.backuped and node.hit_count >= self.write_through_threshold:
+        if (
+            self.cache_controller is not None
+            and not node.backuped
+            and node.hit_count >= self.write_through_threshold
+        ):
             self.write_backup(node)
 
     def write_backup_storage(self, node: UnifiedTreeNode) -> None:

python/sglang/srt/mem_cache/utils.py

@@ -14,7 +14,7 @@
 """Common utilities."""
 
 import hashlib
-from typing import Any, List, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import torch
 import triton
@@ -22,6 +22,37 @@ import triton.language as tl
 
 from sglang.jit_kernel.utils import is_arch_support_pdl
 from sglang.srt.environ import envs
+from sglang.srt.mem_cache.evict_policy import (
+    EvictionStrategy,
+    FIFOStrategy,
+    FILOStrategy,
+    LFUStrategy,
+    LRUStrategy,
+    MRUStrategy,
+    PriorityStrategy,
+    SLRUStrategy,
+)
+
+_EVICTION_POLICY_FACTORIES: dict[str, Callable[[], EvictionStrategy]] = {
+    "lru": LRUStrategy,
+    "lfu": LFUStrategy,
+    "fifo": FIFOStrategy,
+    "mru": MRUStrategy,
+    "filo": FILOStrategy,
+    "priority": PriorityStrategy,
+    "slru": SLRUStrategy,
+}
+
+
+def get_eviction_strategy(eviction_policy: str) -> EvictionStrategy:
+    policy = eviction_policy.lower()
+    try:
+        return _EVICTION_POLICY_FACTORIES[policy]()
+    except KeyError:
+        supported = "', '".join(_EVICTION_POLICY_FACTORIES)
+        raise ValueError(
+            f"Unknown eviction policy: {policy}. Supported policies: '{supported}'."
+        ) from None
 
 
 @triton.jit

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -2,7 +2,7 @@
 
 import unittest
 from array import array
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Optional
 from unittest import mock
 
@@ -89,6 +89,7 @@ class CacheConfig:
     head_num: int = 2
     head_dim: int = 64
     dtype: torch.dtype = torch.bfloat16
+    eviction_policy: str = "lru"
 
     @property
     def has_mamba(self) -> bool:
@@ -233,6 +234,7 @@ def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False):
         tree_components=cfg.components,
         enable_mamba_extra_buffer=cfg.enable_mamba_extra_buffer,
         enable_kv_cache_events=enable_kv_cache_events,
+        eviction_policy=cfg.eviction_policy,
     )
     tree = UnifiedRadixCache(params=cache_init_params)
     tree.cache_init_params = cache_init_params
@@ -443,11 +445,11 @@ class UnifiedRadixCacheSuite:
         allocator.full_to_swa_index_mapping[full_indices] = swa_indices
         return full_indices[:need_size]
 
-    def _insert(self, tree, allocator, req_to_token_pool, tokens):
+    def _insert(self, tree, allocator, req_to_token_pool, tokens, priority=0):
         """Insert tokens, attaching mamba data when the config has mamba."""
         key = RadixKey(array("q", tokens))
         value = self._alloc(allocator, len(tokens))
-        params = InsertParams(key=key, value=value[: len(key)])
+        params = InsertParams(key=key, value=value[: len(key)], priority=priority)
         if self.cfg.has_mamba:
             req = self._make_req(req_to_token_pool)
             params.mamba_value = req.mamba_pool_idx.unsqueeze(0)
@@ -1282,6 +1284,27 @@ class UnifiedRadixCacheSuite:
         self.assertEqual(len(m_new.device_indices), len(seq_new))
         tree.sanity_check()
 
+    def test_evict_respects_priority_policy(self):
+        if self.cfg.components != (ComponentType.FULL,):
+            self.skipTest("priority policy ordering is covered on Full-only configs")
+        priority_cfg = replace(self.cfg, eviction_policy="priority")
+        tree, allocator, req_to_token_pool = build_fixture(priority_cfg)
+        seq_high = self._make_seq(1, 2)
+        seq_low = self._make_seq(500, 2)
+
+        self._insert(tree, allocator, req_to_token_pool, seq_high, priority=10)
+        self._insert(tree, allocator, req_to_token_pool, seq_low, priority=0)
+
+        tree.evict(EvictParams(num_tokens=len(seq_low)))
+
+        m_high = tree.match_prefix(
+            MatchPrefixParams(key=RadixKey(array("q", seq_high)))
+        )
+        m_low = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_low))))
+        self.assertEqual(len(m_high.device_indices), len(seq_high))
+        self.assertEqual(len(m_low.device_indices), 0)
+        tree.sanity_check()
+
     def test_evict_multiple_independent_leaves(self):
         """Evicting multiple independent leaves works correctly."""
         tree, allocator, req_to_token_pool = build_fixture(self.cfg)