diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 97a3dceb6..08a837833 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2614,6 +2614,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return total 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) evict_from_tree_cache(self.tree_cache, num_tokens) return self.token_to_kv_pool_allocator.available_size() >= num_tokens diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 7a3e1c39d..25d3f96e8 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -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) ) else: - # Standard allocator - if allocator.available_size() < num_tokens: - tree_cache.evict(EvictParams(num_tokens=num_tokens)) + # Standard allocator: evict only the shortfall (mirrors the SWA arm) + available_size = allocator.available_size() + 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):