diff --git a/python/sglang/srt/mem_cache/evict_policy.py b/python/sglang/srt/mem_cache/evict_policy.py index 687593418..30ab1983f 100644 --- a/python/sglang/srt/mem_cache/evict_policy.py +++ b/python/sglang/srt/mem_cache/evict_policy.py @@ -44,3 +44,22 @@ class PriorityStrategy(EvictionStrategy): 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) + + +class SLRUStrategy(EvictionStrategy): + def __init__(self, protected_threshold: int = 2): + self.protected_threshold = protected_threshold + + def get_priority(self, node: "TreeNode") -> Tuple[int, float]: + # Priority Logic: + # Smaller value = Evicted earlier. + # + # Segment 0 (Probationary): hit_count < threshold + # Segment 1 (Protected): hit_count >= threshold + # + # Tuple comparison: (segment, last_access_time) + # Nodes in segment 0 will always be evicted before segment 1. + # Inside the same segment, older nodes (smaller time) are evicted first. + + is_protected = 1 if node.hit_count >= self.protected_threshold else 0 + return (is_protected, node.last_access_time) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index c9c32f54d..8e13ff4d1 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -57,6 +57,7 @@ from sglang.srt.mem_cache.evict_policy import ( LRUStrategy, MRUStrategy, PriorityStrategy, + SLRUStrategy, ) from sglang.srt.mem_cache.hicache_storage import get_hash_str, hash_str_to_int64 @@ -322,9 +323,12 @@ class RadixCache(BasePrefixCache): 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'." + f"Unknown eviction policy: {self.eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority', 'slru'." ) self.evictable_leaves = set() @@ -447,13 +451,14 @@ class RadixCache(BasePrefixCache): key = params.key value = params.value priority = params.priority + chunked = params.chunked if value is None: value = torch.tensor(key.token_ids, dtype=torch.int64) key, value = self.maybe_bigram_convert(key, value) - prefix_len = self._insert_helper(self.root_node, key, value, priority) + prefix_len = self._insert_helper(self.root_node, key, value, priority, chunked) return InsertResult(prefix_len=prefix_len) def cache_finished_req(self, req: Req, is_insert: bool = True): @@ -688,6 +693,7 @@ class RadixCache(BasePrefixCache): # new_node -> child # New node inherits child's priority (represents shared prefix) new_node = TreeNode(priority=child.priority) + new_node.hit_count = child.hit_count new_node.children = {self.get_child_key_fn(key[split_len:]): child} new_node.parent = child.parent new_node.lock_ref = child.lock_ref @@ -705,7 +711,22 @@ class RadixCache(BasePrefixCache): return new_node - def _insert_helper(self, node: TreeNode, key: RadixKey, value, priority: int = 0): + def _inc_hit_count(self, node: TreeNode, chunked: bool = False): + # Skip the hit count update for chunked requests to avoid self-referencing + # inflation where a chunked request increments hit_count on nodes it created + # in previous chunks. + if chunked: + return + node.hit_count += 1 + + def _insert_helper( + self, + node: TreeNode, + key: RadixKey, + value, + priority: int = 0, + chunked: bool = False, + ): # Convert None priority to 0 if priority is None: priority = 0 @@ -730,10 +751,11 @@ 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) + self._inc_hit_count(new_node, chunked) node = new_node else: node.priority = max(node.priority, priority) - + self._inc_hit_count(node, chunked) if len(key): child_key = self.get_child_key_fn(key) @@ -742,6 +764,7 @@ class RadixCache(BasePrefixCache): new_node.parent = node new_node.key = key new_node.value = value.clone() + self._inc_hit_count(new_node, chunked) node.children[child_key] = new_node self.evictable_size_ += len(key) self._update_leaf_status(node) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bcaa015f2..f2dfe68d1 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -169,7 +169,7 @@ NSA_CHOICES = [ "trtllm", ] -RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu"] +RADIX_EVICTION_POLICY_CHOICES = ["lru", "lfu", "slru"] RL_ON_POLICY_TARGET_CHOICES = ["fsdp"] @@ -3633,7 +3633,7 @@ class ServerArgs: type=str, choices=RADIX_EVICTION_POLICY_CHOICES, default=ServerArgs.radix_eviction_policy, - help="The eviction policy of radix trees. 'lru' stands for Least Recently Used, 'lfu' stands for Least Frequently Used.", + help="The eviction policy of radix trees. 'lru' stands for Least Recently Used, 'lfu' stands for Least Frequently Used, and 'slru' stands for Segmented Least Recently Used.", ) parser.add_argument( "--enable-prefill-delayer", diff --git a/test/registered/radix_cache/test_radix_cache_slru_accuracy.py b/test/registered/radix_cache/test_radix_cache_slru_accuracy.py new file mode 100644 index 000000000..de3a7cc25 --- /dev/null +++ b/test/registered/radix_cache/test_radix_cache_slru_accuracy.py @@ -0,0 +1,143 @@ +import unittest + +import torch + +from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.mem_cache.base_prefix_cache import ( + EvictParams, + InsertParams, + MatchPrefixParams, +) +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool +from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-test-small-1-gpu") + + +class TestSLRUAccuracy(unittest.TestCase): + + def setUp(self): + """Setup minimal memory pools for testing""" + device = "cpu" # Using CPU for testing simplicity + dtype = torch.float16 + + # Create smaller KV cache to ensure evictions occur + self.kv_cache = MHATokenToKVPool( + size=8, # Very small size to trigger evictions quickly + page_size=1, + dtype=dtype, + head_num=8, + head_dim=64, + layer_num=1, + device=device, + enable_memory_saver=False, + ) + + # Create token-to-KV pool allocator + self.token_to_kv_pool = TokenToKVPoolAllocator( + size=8, dtype=dtype, device=device, kvcache=self.kv_cache, need_sort=False + ) + + # Create req-to-token pool + self.req_to_token_pool = ReqToTokenPool( + size=8, max_context_len=1024, device=device, enable_memory_saver=False + ) + + # Create a cache with the memory pools + params = CacheInitParams( + disable=False, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool, + page_size=1, + eviction_policy="slru", + enable_kv_cache_events=False, + ) + + self.cache = RadixCache(params) + + def test_eviction_mechanism(self): + """Test that SLRU eviction mechanism works correctly""" + + # Insert one key-value three times (high frequency access) + frequent_key = RadixKey( + token_ids=[1, 2], extra_key=None + ) # High hit rate, should be retained + frequent_val = torch.tensor([10, 20], dtype=torch.int64) + + # Insert the frequent key multiple times to increase its hit count + for _ in range(3): + self.cache.insert(InsertParams(key=frequent_key, value=frequent_val)) + + # Insert first low-frequency key-value pair that should be evicted + first_low_freq_key = RadixKey( + token_ids=[5, 6], extra_key=None + ) # Low hit rate, should be evicted + first_low_freq_val = torch.tensor([50, 60], dtype=torch.int64) + + self.cache.insert( + InsertParams(key=first_low_freq_key, value=first_low_freq_val) + ) + + # Insert other key-values once each (low frequency access) - fill up the cache + other_keys = [] + for i in range(4): # Reduce the number to fit in our smaller cache + key = RadixKey( + token_ids=[i + 10], extra_key=None + ) # Unique keys for low-frequency items + val = torch.tensor([i + 100], dtype=torch.int64) + self.cache.insert(InsertParams(key=key, value=val)) + other_keys.append(key) + + # Now insert more items to trigger evictions + for i in range(6, 10): # Add more items to definitely exceed capacity + key = RadixKey( + token_ids=[i * 2], extra_key=None + ) # Different pattern to avoid conflicts + val = torch.tensor([i * 200], dtype=torch.int64) + self.cache.insert(InsertParams(key=key, value=val)) + + # Now trigger eviction explicitly to make space + evict_result = self.cache.evict( + EvictParams(num_tokens=4) + ) # Try to evict 4 tokens worth of space + + # Check if the frequently accessed key-value is still present + # The frequent key should have higher hit count and remain in cache due to SLRU policy + frequent_match_result = self.cache.match_prefix( + MatchPrefixParams(key=frequent_key) + ) + + # Check if the first low-frequency key-value has been evicted + # The first low-freq key should have lower hit count and be evicted due to SLRU policy + first_low_freq_match_result = self.cache.match_prefix( + MatchPrefixParams(key=first_low_freq_key) + ) + + # Verify the frequent key is still present in cache after evictions + self.assertIsNotNone( + frequent_match_result, + "Frequently accessed key should still be in cache after evictions", + ) + + # Check if the tensor is empty, which indicates the key was not found (evicted) + is_frequent_key_present = frequent_match_result.device_indices.numel() > 0 + self.assertTrue( + is_frequent_key_present, + "Frequently accessed key should still be in cache after evictions", + ) + + # Verify the first low-frequency key has been evicted + # The device_indices tensor should be empty when the key is not found + is_first_low_freq_key_present = ( + first_low_freq_match_result.device_indices.numel() > 0 + ) + self.assertFalse( + is_first_low_freq_key_present, + "First inserted low-frequency key should be evicted after evictions", + ) + + +if __name__ == "__main__": + unittest.main()