chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Performance] Optimize radix cache eviction performance (#14339)

Browse Source 
Signed-off-by: Xingrui Yi <yixingrui@linux.alibaba.com>
Co-authored-by: Xuchun Shang <xuchun.shang@gmail.com>
This commit is contained in:
EkiRui
2026-02-03 09:44:20 +08:00
committed by GitHub
co-authored by Xuchun Shang
parent c8da307d7e
commit fd983b09b6
4 changed files with 103 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/hiradix_cache.py
+77 -27
View File
@@ -146,6 +146,8 @@ class HiRadixCache(RadixCache):
# Detach storage backend automatically on process shutdown # Detach storage backend automatically on process shutdown
atexit.register(self.shutdown) atexit.register(self.shutdown)
self.evictable_host_leaves = set()
super().__init__(params=params) super().__init__(params=params)
def shutdown(self): def shutdown(self):
@@ -556,6 +558,7 @@ class HiRadixCache(RadixCache):
TreeNode.counter = 0 TreeNode.counter = 0
self.cache_controller.reset() self.cache_controller.reset()
self.token_to_kv_pool_host.clear() self.token_to_kv_pool_host.clear()
self.evictable_host_leaves.clear()
super().reset() super().reset()
def get_height(self, node: TreeNode): def get_height(self, node: TreeNode):
@@ -695,10 +698,61 @@ class HiRadixCache(RadixCache):
def evictable_size(self): def evictable_size(self):
return self.evictable_size_ 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: def evict(self, params: EvictParams) -> EvictResult:
start_time = time.perf_counter() start_time = time.perf_counter()
num_tokens = params.num_tokens num_tokens = params.num_tokens
leaves = self._collect_leaves_device() leaves = list(self.evictable_leaves)
eviction_heap = [ eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves (self.eviction_strategy.get_priority(node), node) for node in leaves
] ]
@@ -747,6 +801,10 @@ class HiRadixCache(RadixCache):
assert num_evicted > 0 assert num_evicted > 0
self.evictable_size_ -= num_evicted self.evictable_size_ -= num_evicted
node.value = None 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 return num_evicted
def _evict_regular(self, node: TreeNode): def _evict_regular(self, node: TreeNode):
@@ -757,7 +815,7 @@ class HiRadixCache(RadixCache):
return num_evicted return num_evicted
def evict_host(self, num_tokens: int): def evict_host(self, num_tokens: int):
leaves = self._collect_leaves() leaves = list(self.evictable_host_leaves)
eviction_heap = [ eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves (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) key = self.get_child_key_fn(x.key)
v = x.parent.children.pop(key, None) v = x.parent.children.pop(key, None)
assert v == x, f"parent does not have child key, {key}" 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: if len(x.parent.children) == 0 and x.parent.evicted:
new_priority = self.eviction_strategy.get_priority(x.parent) 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.host_value = host_value.clone()
new_node.hash_value = hash_value new_node.hash_value = hash_value
node.children[child_key] = new_node 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 return matched_length
def _match_prefix_helper(self, node: TreeNode, key: RadixKey): 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 # this often happens in the case of KV cache recomputation
node.value = value[:prefix_len] node.value = value[:prefix_len]
self.evictable_size_ += len(node.value) 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: else:
self._inc_hit_count(node, chunked) self._inc_hit_count(node, chunked)
total_prefix_length += prefix_len total_prefix_length += prefix_len
@@ -1240,6 +1309,10 @@ class HiRadixCache(RadixCache):
if new_node.evicted: if new_node.evicted:
new_node.value = value[:prefix_len].clone() new_node.value = value[:prefix_len].clone()
self.evictable_size_ += len(new_node.value) 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: else:
self._inc_hit_count(new_node, chunked) self._inc_hit_count(new_node, chunked)
total_prefix_length += prefix_len total_prefix_length += prefix_len
@@ -1258,6 +1331,8 @@ class HiRadixCache(RadixCache):
new_node.value = value.clone() new_node.value = value.clone()
node.children[child_key] = new_node node.children[child_key] = new_node
self.evictable_size_ += len(value) self.evictable_size_ += len(value)
self._update_leaf_status(node)
self._update_leaf_status(new_node)
# Compute hash_value if storage is enabled # Compute hash_value if storage is enabled
if self.enable_storage: if self.enable_storage:
@@ -1267,31 +1342,6 @@ class HiRadixCache(RadixCache):
self._inc_hit_count(new_node, chunked) self._inc_hit_count(new_node, chunked)
return InsertResult(prefix_len=total_prefix_length) 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): def release_aborted_request(self, rid: str):
if rid not in self.ongoing_prefetch: if rid not in self.ongoing_prefetch:
return return
python/sglang/srt/mem_cache/mamba_radix_cache.py
-13
View File
@@ -1156,19 +1156,6 @@ class MambaRadixCache(BasePrefixCache):
self.full_evictable_size_ -= len(node.key) 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]: def _collect_nontombstone_nodes(self) -> List[TreeNode]:
ret_list = [] ret_list = []
stack = [self.root_node] stack = [self.root_node]
python/sglang/srt/mem_cache/radix_cache.py
+26 -15
View File
@@ -301,6 +301,8 @@ class RadixCache(BasePrefixCache):
raise ValueError( 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'."
) )
self.evictable_leaves = set()
self.reset() self.reset()
@classmethod @classmethod
@@ -333,6 +335,7 @@ class RadixCache(BasePrefixCache):
self.root_node.hash_value = [] self.root_node.hash_value = []
self.evictable_size_ = 0 self.evictable_size_ = 0
self.protected_size_ = 0 self.protected_size_ = 0
self.evictable_leaves.clear()
self._record_all_cleared_event() self._record_all_cleared_event()
def maybe_bigram_convert( def maybe_bigram_convert(
@@ -564,7 +567,7 @@ class RadixCache(BasePrefixCache):
start_time = time.perf_counter() start_time = time.perf_counter()
num_tokens = params.num_tokens num_tokens = params.num_tokens
leaves = self._collect_leaves() leaves = list(self.evictable_leaves)
eviction_heap = [ eviction_heap = [
(self.eviction_strategy.get_priority(node), node) for node in leaves (self.eviction_strategy.get_priority(node), node) for node in leaves
] ]
@@ -598,6 +601,7 @@ class RadixCache(BasePrefixCache):
self.protected_size_ += len(node.key) self.protected_size_ += len(node.key)
delta -= len(node.key) delta -= len(node.key)
node.lock_ref += 1 node.lock_ref += 1
self._update_leaf_status(node)
node = node.parent node = node.parent
return delta return delta
@@ -612,6 +616,7 @@ class RadixCache(BasePrefixCache):
self.protected_size_ -= len(node.key) self.protected_size_ -= len(node.key)
delta += len(node.key) delta += len(node.key)
node.lock_ref -= 1 node.lock_ref -= 1
self._update_leaf_status(node)
if node.parent is None: if node.parent is None:
assert ( assert (
node is self.root_node node is self.root_node
@@ -725,6 +730,8 @@ class RadixCache(BasePrefixCache):
new_node.value = value.clone() new_node.value = value.clone()
node.children[child_key] = new_node node.children[child_key] = new_node
self.evictable_size_ += len(key) self.evictable_size_ += len(key)
self._update_leaf_status(node)
self._update_leaf_status(new_node)
# Hash will be computed lazily during event emission # Hash will be computed lazily during event emission
self._record_store_event(new_node) self._record_store_event(new_node)
return total_prefix_length return total_prefix_length
@@ -753,6 +760,24 @@ class RadixCache(BasePrefixCache):
assert v == node, f"parent does not have child key, {key}" assert v == node, f"parent does not have child key, {key}"
self.evictable_size_ -= len(node.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): def _total_size_helper(self):
total_size = 0 total_size = 0
@@ -766,20 +791,6 @@ class RadixCache(BasePrefixCache):
stack.append(child) stack.append(child)
return total_size 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): def _record_store_event(self, node: TreeNode):
# One BlockStored per ``page_size`` chunk. # One BlockStored per ``page_size`` chunk.
if self.enable_kv_cache_events: if self.enable_kv_cache_events:
python/sglang/srt/mem_cache/swa_radix_cache.py
-13
View File
@@ -1104,19 +1104,6 @@ class SWARadixCache(BasePrefixCache):
self.full_evictable_size_ -= len(node.key) 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]: def _collect_nontombstone_nodes(self) -> List[TreeNode]:
ret_list = [] ret_list = []
stack = [self.root_node] stack = [self.root_node]