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

2026-07-22 12:41:06 -07:00
parent 5c6a29b3c3
commit 98cc8d91cf
2 changed files with 6 additions and 3 deletions
@@ -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