minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[UnifiedTree]: Support eviction priority (#26549)

Browse Source
This commit is contained in:
Zhangheng
2026-05-30 15:19:43 +08:00
committed by GitHub
parent c048ebd10d
commit 7662210406
5 changed files with 116 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/radix_cache.py
+2 -30
View File
@@ -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 (
EvictionStrategy,
FIFOStrategy,
FILOStrategy,
LFUStrategy,
LRUStrategy,
MRUStrategy,
PriorityStrategy,
SLRUStrategy,
)
from sglang.srt.mem_cache.utils import split_node_hash_value
from sglang.srt.mem_cache.utils import get_eviction_strategy, 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: 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.eviction_strategy = get_eviction_strategy(self.eviction_policy)
self.evictable_leaves = set()
self.reset()
python/sglang/srt/mem_cache/unified_cache_components/full_component.py
+16 -5
View File
@@ -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 = [(n.last_access_time, n) for n in self.cache.evictable_device_leaves]
heap = [
(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
+40 -12
View File
@@ -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
+32 -1
View File
@@ -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