[Feature] Add SLRU eviction policy & fix RadixCache hit_count bug (#18843)

Co-authored-by: zhangheng <hzh0425@apache.org>

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)

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)

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",

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()