minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[Fix] Evict only the KV shortfall in evict_from_tree_cache (#32016)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-07-22 12:41:06 -07:00
committed by GitHub
parent 5c6a29b3c3
commit 98cc8d91cf
2 changed files with 6 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+2
View File
@@ -2614,6 +2614,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
return total return total
def check_decode_mem(self, selected_indices: Optional[List[int]] = None): def check_decode_mem(self, selected_indices: Optional[List[int]] = None):
"""Reclaim evictable tree-cache entries (shortfall only), then report
whether the next decode step fits in the KV pool."""
num_tokens = self.new_tokens_required_next_decode(selected_indices) num_tokens = self.new_tokens_required_next_decode(selected_indices)
evict_from_tree_cache(self.tree_cache, num_tokens) evict_from_tree_cache(self.tree_cache, num_tokens)
return self.token_to_kv_pool_allocator.available_size() >= num_tokens return self.token_to_kv_pool_allocator.available_size() >= num_tokens
python/sglang/srt/mem_cache/common.py
+4 -3
View File
@@ -123,9 +123,10 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens) EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
) )
else: else:
# Standard allocator # Standard allocator: evict only the shortfall (mirrors the SWA arm)
if allocator.available_size() < num_tokens: available_size = allocator.available_size()
tree_cache.evict(EvictParams(num_tokens=num_tokens)) if available_size < num_tokens:
tree_cache.evict(EvictParams(num_tokens=num_tokens - available_size))
def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True): def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = True):