From fd983b09b68c076d448a55266a29ffdfdff2d06e Mon Sep 17 00:00:00 2001 From: EkiRui <374389385@qq.com> Date: Tue, 3 Feb 2026 09:44:20 +0800 Subject: [PATCH] [Performance] Optimize radix cache eviction performance (#14339) Signed-off-by: Xingrui Yi Co-authored-by: Xuchun Shang --- python/sglang/srt/mem_cache/hiradix_cache.py | 104 +++++++++++++----- .../sglang/srt/mem_cache/mamba_radix_cache.py | 13 --- python/sglang/srt/mem_cache/radix_cache.py | 41 ++++--- .../sglang/srt/mem_cache/swa_radix_cache.py | 13 --- 4 files changed, 103 insertions(+), 68 deletions(-) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index e8fa02cfd..7b016a144 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -146,6 +146,8 @@ class HiRadixCache(RadixCache): # Detach storage backend automatically on process shutdown atexit.register(self.shutdown) + self.evictable_host_leaves = set() + super().__init__(params=params) def shutdown(self): @@ -556,6 +558,7 @@ class HiRadixCache(RadixCache): TreeNode.counter = 0 self.cache_controller.reset() self.token_to_kv_pool_host.clear() + self.evictable_host_leaves.clear() super().reset() def get_height(self, node: TreeNode): @@ -695,10 +698,61 @@ class HiRadixCache(RadixCache): def evictable_size(self): return self.evictable_size_ + def inc_lock_ref(self, node: TreeNode): + if self.disable: + return 0 + + delta = 0 + while node != self.root_node: + if node.lock_ref == 0: + self.evictable_size_ -= len(node.key) + self.protected_size_ += len(node.key) + delta -= len(node.key) + node.lock_ref += 1 + self._update_leaf_status(node) + self._update_host_leaf_status(node) + node = node.parent + return delta + + def dec_lock_ref(self, node: TreeNode): + if self.disable: + return 0 + + delta = 0 + while node != self.root_node: + if node.lock_ref == 1: + self.evictable_size_ += len(node.key) + self.protected_size_ -= len(node.key) + delta += len(node.key) + node.lock_ref -= 1 + self._update_leaf_status(node) + self._update_host_leaf_status(node) + if node.parent is None: + assert ( + node is self.root_node + ), f"This request holds the node from another tree" + node = node.parent + return delta + + def _update_host_leaf_status(self, node: TreeNode): + if not node.evicted or node.lock_ref > 0: + if node in self.evictable_host_leaves: + self.evictable_host_leaves.remove(node) + return + + for child in node.children.values(): + if child.evicted: + if node in self.evictable_host_leaves: + self.evictable_host_leaves.remove(node) + return + + if node not in self.evictable_host_leaves: + self.evictable_host_leaves.add(node) + def evict(self, params: EvictParams) -> EvictResult: start_time = time.perf_counter() num_tokens = params.num_tokens - leaves = self._collect_leaves_device() + leaves = list(self.evictable_leaves) eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves ] @@ -747,6 +801,10 @@ class HiRadixCache(RadixCache): assert num_evicted > 0 self.evictable_size_ -= num_evicted node.value = None + self._update_leaf_status(node) + self._update_host_leaf_status(node) + # update leaf status for the parent because the node is evicted + self._update_leaf_status(node.parent) return num_evicted def _evict_regular(self, node: TreeNode): @@ -757,7 +815,7 @@ class HiRadixCache(RadixCache): return num_evicted def evict_host(self, num_tokens: int): - leaves = self._collect_leaves() + leaves = list(self.evictable_host_leaves) eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves ] @@ -781,6 +839,9 @@ class HiRadixCache(RadixCache): key = self.get_child_key_fn(x.key) v = x.parent.children.pop(key, None) assert v == x, f"parent does not have child key, {key}" + if x in self.evictable_host_leaves: + self.evictable_host_leaves.remove(x) + self._update_host_leaf_status(x.parent) if len(x.parent.children) == 0 and x.parent.evicted: new_priority = self.eviction_strategy.get_priority(x.parent) @@ -1141,6 +1202,10 @@ class HiRadixCache(RadixCache): new_node.host_value = host_value.clone() new_node.hash_value = hash_value node.children[child_key] = new_node + self._update_host_leaf_status(new_node) + self._update_leaf_status(node) + self._update_host_leaf_status(node) + return matched_length def _match_prefix_helper(self, node: TreeNode, key: RadixKey): @@ -1229,6 +1294,10 @@ class HiRadixCache(RadixCache): # this often happens in the case of KV cache recomputation node.value = value[:prefix_len] self.evictable_size_ += len(node.value) + self._update_leaf_status(node) + self._update_host_leaf_status(node) + # update parent status as a new leaf is added into device + self._update_leaf_status(node.parent) else: self._inc_hit_count(node, chunked) total_prefix_length += prefix_len @@ -1240,6 +1309,10 @@ class HiRadixCache(RadixCache): if new_node.evicted: new_node.value = value[:prefix_len].clone() self.evictable_size_ += len(new_node.value) + self._update_leaf_status(new_node) + self._update_host_leaf_status(new_node) + # update parent status as a new leaf is added into device + self._update_leaf_status(new_node.parent) else: self._inc_hit_count(new_node, chunked) total_prefix_length += prefix_len @@ -1258,6 +1331,8 @@ class HiRadixCache(RadixCache): new_node.value = value.clone() node.children[child_key] = new_node self.evictable_size_ += len(value) + self._update_leaf_status(node) + self._update_leaf_status(new_node) # Compute hash_value if storage is enabled if self.enable_storage: @@ -1267,31 +1342,6 @@ class HiRadixCache(RadixCache): self._inc_hit_count(new_node, chunked) return InsertResult(prefix_len=total_prefix_length) - def _collect_leaves_device(self): - def is_leaf(node): - if node.evicted: - return False - if node == self.root_node: - return False - if len(node.children) == 0: - return True - for child in node.children.values(): - if not child.evicted: - return False - return True - - ret_list = [] - stack = [self.root_node] - while stack: - cur_node = stack.pop() - if is_leaf(cur_node): - ret_list.append(cur_node) - else: - for cur_child in cur_node.children.values(): - if not cur_child.evicted: - stack.append(cur_child) - return ret_list - def release_aborted_request(self, rid: str): if rid not in self.ongoing_prefetch: return diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index ad4ff29e8..b89c3b15c 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -1156,19 +1156,6 @@ class MambaRadixCache(BasePrefixCache): self.full_evictable_size_ -= len(node.key) - def _collect_leaves(self) -> List[TreeNode]: - ret_list = [] - stack = [self.root_node] - - while stack: - cur_node = stack.pop() - if len(cur_node.children) == 0: - ret_list.append(cur_node) - else: - stack.extend(cur_node.children.values()) - - return ret_list - def _collect_nontombstone_nodes(self) -> List[TreeNode]: ret_list = [] stack = [self.root_node] diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 4ff789431..fcec2c9ac 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -301,6 +301,8 @@ class RadixCache(BasePrefixCache): raise ValueError( f"Unknown eviction policy: {self.eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority'." ) + + self.evictable_leaves = set() self.reset() @classmethod @@ -333,6 +335,7 @@ class RadixCache(BasePrefixCache): self.root_node.hash_value = [] self.evictable_size_ = 0 self.protected_size_ = 0 + self.evictable_leaves.clear() self._record_all_cleared_event() def maybe_bigram_convert( @@ -564,7 +567,7 @@ class RadixCache(BasePrefixCache): start_time = time.perf_counter() num_tokens = params.num_tokens - leaves = self._collect_leaves() + leaves = list(self.evictable_leaves) eviction_heap = [ (self.eviction_strategy.get_priority(node), node) for node in leaves ] @@ -598,6 +601,7 @@ class RadixCache(BasePrefixCache): self.protected_size_ += len(node.key) delta -= len(node.key) node.lock_ref += 1 + self._update_leaf_status(node) node = node.parent return delta @@ -612,6 +616,7 @@ class RadixCache(BasePrefixCache): self.protected_size_ -= len(node.key) delta += len(node.key) node.lock_ref -= 1 + self._update_leaf_status(node) if node.parent is None: assert ( node is self.root_node @@ -725,6 +730,8 @@ class RadixCache(BasePrefixCache): new_node.value = value.clone() node.children[child_key] = new_node self.evictable_size_ += len(key) + self._update_leaf_status(node) + self._update_leaf_status(new_node) # Hash will be computed lazily during event emission self._record_store_event(new_node) return total_prefix_length @@ -753,6 +760,24 @@ class RadixCache(BasePrefixCache): assert v == node, f"parent does not have child key, {key}" self.evictable_size_ -= len(node.key) + if node in self.evictable_leaves: + self.evictable_leaves.remove(node) + self._update_leaf_status(node.parent) + + def _update_leaf_status(self, node: TreeNode): + if node.evicted or node.lock_ref > 0: + if node in self.evictable_leaves: + self.evictable_leaves.remove(node) + return + + for child in node.children.values(): + if not child.evicted: + if node in self.evictable_leaves: + self.evictable_leaves.remove(node) + return + + if node not in self.evictable_leaves: + self.evictable_leaves.add(node) def _total_size_helper(self): total_size = 0 @@ -766,20 +791,6 @@ class RadixCache(BasePrefixCache): stack.append(child) return total_size - def _collect_leaves(self): - ret_list = [] - stack = list(self.root_node.children.values()) - - while stack: - cur_node = stack.pop() - if len(cur_node.children) == 0: - if cur_node.lock_ref == 0: - ret_list.append(cur_node) - else: - stack.extend(cur_node.children.values()) - - return ret_list - def _record_store_event(self, node: TreeNode): # One BlockStored per ``page_size`` chunk. if self.enable_kv_cache_events: diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index b992b4727..1feaca670 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -1104,19 +1104,6 @@ class SWARadixCache(BasePrefixCache): self.full_evictable_size_ -= len(node.key) - def _collect_leaves(self) -> List[TreeNode]: - ret_list = [] - stack = [self.root_node] - - while stack: - cur_node = stack.pop() - if len(cur_node.children) == 0: - ret_list.append(cur_node) - else: - stack.extend(cur_node.children.values()) - - return ret_list - def _collect_nontombstone_nodes(self) -> List[TreeNode]: ret_list = [] stack = [self.root_node]