kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-07 15:57:17 +00:00
Code Issues 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
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

104
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -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

13
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -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]

41
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -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:

13
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -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]