diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 6e35d1a31..41c392fec 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/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 ( - 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() diff --git a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py b/python/sglang/srt/mem_cache/unified_cache_components/full_component.py index 5ed89e3b4..ea35536eb 100644 --- a/python/sglang/srt/mem_cache/unified_cache_components/full_component.py +++ b/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 = [(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, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index bc47648ad..22f418541 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/utils.py b/python/sglang/srt/mem_cache/utils.py index 6d654b84c..a55b63b46 100644 --- a/python/sglang/srt/mem_cache/utils.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index a59f6ffd7..f3a3d5f2e 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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)