From 25916608693e669e91472c7096750ea8fcae5238 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 29 Sep 2026 15:09:56 -0400 Subject: [PATCH] [Core][HiSparse] Keep unpinned blocks last in the free queue BlockPool.unpin_blocks documents its blocks as last-resort eviction candidates, but blocks freed afterwards were appended behind them, so an allocation could evict HiSparse's GPU prefix copies before blocks nobody would reuse. FreeKVCacheBlockQueue now keeps a last-resort segment at its tail that later appends stay in front of, and unpin_blocks uses it. Only HiSparse unpins blocks; without them the queue behaves as before. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Matthew Bonanni --- tests/v1/core/test_kv_cache_utils.py | 25 ++++++++++ tests/v1/core/test_prefix_caching.py | 6 +-- vllm/v1/core/block_pool.py | 4 +- vllm/v1/core/kv_cache_utils.py | 69 ++++++++++++++++++---------- 4 files changed, 74 insertions(+), 30 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 0c56e89d597c..f9a80f0a540d 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -736,6 +736,31 @@ def test_free_kv_cache_block_queue_append_n(): ) +def test_free_kv_cache_block_queue_last_resort_stays_behind(): + """Later appends stay ahead of the last-resort segment even after the + segment's first block leaves the queue by removal or popping.""" + blocks = [KVCacheBlock(block_id=i) for i in range(7)] + queue = FreeKVCacheBlockQueue(blocks[0:1]) + + def order() -> list[int]: + return [block.block_id for block in queue.get_all_free_blocks()] + + queue.append_last_resort_n(blocks[1:3]) + queue.append(blocks[3]) + assert order() == [0, 3, 1, 2] + + queue.remove(blocks[1]) + queue.append_n(blocks[4:5]) + assert order() == [0, 3, 4, 2] + + assert [block.block_id for block in queue.popleft_n(4)] == [0, 3, 4, 2] + queue.append_n(blocks[5:6]) + queue.append_last_resort_n(blocks[6:7]) + queue.append(blocks[1]) + assert order() == [5, 1, 6] + assert queue.num_free_blocks == 3 + + def test_free_kv_cache_block_queue_prepend_n(): # Seed the queue with one block so prepend has an existing head to splice # in front of (fake_head->b0->fake_tail). diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index e4c54e880ccc..6ec6fc58b760 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1139,7 +1139,7 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): ) -@pytest.mark.parametrize("free_blocks,adopted_pages", [(7, []), (10, [0, 1])]) +@pytest.mark.parametrize("free_blocks,adopted_pages", [(7, []), (10, [0, 1, 2])]) def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( free_blocks, adopted_pages ): @@ -1148,8 +1148,8 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( Admission counts free GPU copies as free, but adopting a copy pins it. When adoption ran before the hit's allocation, it took blocks the allocation was promised and the pool ran dry (``Cannot get N free blocks``). The allocation - now goes first, evicting a prefix's copies tail first; the copies it leaves - are adopted. + now goes first, evicting copies only after every other free block and a + prefix's tail before its head; the copies it leaves are adopted. """ manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 664804d50de4..b7514b572466 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -716,7 +716,7 @@ def unpin_blocks( """Release references to blocks that stay readable until reused. The blocks become last-resort eviction candidates regardless of prefix - caching: they join the tail of the free queue and count as free. + caching: they stay behind every later-freed block and count as free. ``on_reuse`` fires when ``get_new_blocks`` hands a block out (or the cache is reset), so the caller can stop reading it. """ @@ -727,7 +727,7 @@ def unpin_blocks( self._reuse_watchers[block.block_id] = on_reuse if block.ref_cnt == 0: released.append(block) - self.free_block_queue.append_n(released) + self.free_block_queue.append_last_resort_n(released) def _maybe_evict_cached_block(self, block: KVCacheBlock) -> bool: """If a block is cached in `cached_block_hash_to_block`, we reset its hash diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index a82d3e811358..33a833559292 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -262,6 +262,10 @@ class FreeKVCacheBlockQueue: Note that we maintain this order by reversing the block order when free blocks of a request. This operation is outside of this class. + Blocks added with ``append_last_resort_n`` form a segment at the tail that + later ``append``/``append_n`` calls stay in front of, so they are evicted + only after every other free block. + Args: blocks: A list of KVCacheBlock objects. @@ -297,6 +301,17 @@ def __init__(self, blocks: list[KVCacheBlock]) -> None: self.fake_free_list_head.next_free_block = self.fake_free_list_tail self.fake_free_list_tail.prev_free_block = self.fake_free_list_head + # First block of the last-resort tail segment, if any. + self.last_resort_head: KVCacheBlock | None = None + + def _unlink_last_resort_head(self, block: KVCacheBlock) -> None: + """Advance the last-resort segment past a block leaving the list.""" + if block is self.last_resort_head: + next_block = block.next_free_block + self.last_resort_head = ( + None if next_block is self.fake_free_list_tail else next_block + ) + def popleft(self) -> KVCacheBlock: """Pop the first free block and reduce num_free_blocks by 1. @@ -323,6 +338,7 @@ def popleft(self) -> KVCacheBlock: "Invalid block found in popleft() " "which doesn't have a valid next_free_block" ) + self._unlink_last_resort_head(first_block) # Connect fake_head and the next block of first_block (i.e. second block # or fake tail). @@ -355,6 +371,7 @@ def popleft_n(self, n: int) -> list[KVCacheBlock]: ret = [] for _ in range(n): assert curr_block is not None + self._unlink_last_resort_head(curr_block) ret.append(curr_block) last_block = curr_block curr_block = curr_block.next_free_block @@ -380,6 +397,7 @@ def remove(self, block: KVCacheBlock) -> None: # This should not happen if the block is from the free list. # It indicates a bug in the caller's logic. raise RuntimeError(f"remove() called on an invalid block: {block}") + self._unlink_last_resort_head(block) # Link the previous block to the next block. block.prev_free_block.next_free_block = block.next_free_block @@ -391,28 +409,14 @@ def remove(self, block: KVCacheBlock) -> None: self.num_free_blocks -= 1 def append(self, block: KVCacheBlock) -> None: - """Put a block back into the free list and increase - num_free_blocks by 1. + """Put a block back into the free list, ahead of any last-resort + blocks, and increase num_free_blocks by 1. Args: block: The block to append. """ - if self.fake_free_list_tail.prev_free_block is None: - raise RuntimeError( - "prev_free_block of fake_free_list_tail should always exist" - ) - last_block: KVCacheBlock = self.fake_free_list_tail.prev_free_block - - # Connect the new block after the last block. - last_block.next_free_block = block - block.prev_free_block = last_block - - # Connect the fake tail after the new block. - block.next_free_block = self.fake_free_list_tail - self.fake_free_list_tail.prev_free_block = block - - self.num_free_blocks += 1 + self.append_n([block]) def prepend_n(self, blocks: list[KVCacheBlock]) -> None: """Put a list of blocks at the front of the free list.""" @@ -436,28 +440,43 @@ def prepend_n(self, blocks: list[KVCacheBlock]) -> None: self.num_free_blocks += len(blocks) def append_n(self, blocks: list[KVCacheBlock]) -> None: - """Put a list of blocks back into the free list + """Put a list of blocks back into the free list, ahead of any + last-resort blocks. Args: blocks: The blocks to append. + """ + self._insert_before(self.last_resort_head or self.fake_free_list_tail, blocks) + + def append_last_resort_n(self, blocks: list[KVCacheBlock]) -> None: + """Put a list of blocks at the tail, behind every block appended later. + + Args: + blocks: The blocks to append, in eviction order. + """ if len(blocks) == 0: return + self._insert_before(self.fake_free_list_tail, blocks) + if self.last_resort_head is None: + self.last_resort_head = blocks[0] - last_block = self.fake_free_list_tail.prev_free_block - assert last_block is not None, ( - "prev_free_block of fake_free_list_tail should always exist" - ) + def _insert_before(self, anchor: KVCacheBlock, blocks: list[KVCacheBlock]) -> None: + if len(blocks) == 0: + return + + last_block = anchor.prev_free_block + assert last_block is not None, "free list anchors always have a predecessor" # Add inter-connections between consecutive blocks for block in blocks: block.prev_free_block = last_block last_block.next_free_block = block last_block = block - # Connect the last block of to the fake tail - last_block.next_free_block = self.fake_free_list_tail - self.fake_free_list_tail.prev_free_block = last_block + # Connect the last block of to the anchor + last_block.next_free_block = anchor + anchor.prev_free_block = last_block self.num_free_blocks += len(blocks)