diff --git a/python/infinilm/llm/cache_manager.py b/python/infinilm/llm/cache_manager.py index 6c045e4ee..a9551186c 100644 --- a/python/infinilm/llm/cache_manager.py +++ b/python/infinilm/llm/cache_manager.py @@ -336,11 +336,15 @@ def free_blocks(self, block_table: Sequence[int]) -> None: def try_free_blocks(self, num_required: int) -> bool: """Evict unreferenced blocks until the requested capacity is available.""" - to_free = [ - block_id - for block_id in self.used_block_ids - if self.blocks[block_id].ref_count == 0 - ] + # Preserve one eviction, if possible, even when capacity is already enough. + num_to_free = max(1, num_required - len(self.free_block_ids)) + to_free = [] + for block_id in self.used_block_ids: + if self.blocks[block_id].ref_count == 0: + to_free.append(block_id) + if len(to_free) >= num_to_free: + break + # Collect candidates before mutating the set being traversed. for block_id in to_free: self._deallocate_block(block_id) if self.can_allocate(num_required):