From d75c2dd5a24dedf19c707e8b326e26016f524603 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 29 Sep 2026 14:41:09 -0400 Subject: [PATCH 1/2] [Bugfix][HiSparse] Adopt GPU prefix copies after the hit's allocation A HiSparse host-prefix hit re-pins free GPU copies of its pages. Admission counts those copies as free blocks, but adoption pinned them before the hit's own allocation, which could then run the pool dry and crash the engine with "Cannot get N free blocks from the pool". Adopt the copies after the allocation instead, from update_residency, so the allocation always gets the blocks admission counted. To keep the copies the allocation does not need, unpinned blocks now form a true last-resort segment at the tail of the free queue, behind every later-freed block, and HiSparse releases a prefix's pages tail first. 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 | 48 +++++++++++++++++++ vllm/v1/core/block_pool.py | 4 +- vllm/v1/core/kv_cache_utils.py | 69 ++++++++++++++++++---------- vllm/v1/hisparse/coordinator.py | 18 ++++++-- 5 files changed, 134 insertions(+), 30 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c2c86053f0f5..6274d8f73baa 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -690,6 +690,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 99d3235006ad..accca79e61c1 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1103,6 +1103,54 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): ) +@pytest.mark.parametrize( + "free_blocks,adopted_pages", [(7, []), (9, [0, 1]), (10, [0, 1, 2])] +) +def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( + free_blocks, adopted_pages +): + """A host prefix hit must not starve its own allocation of free blocks. + + 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 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)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + _publish_hisparse_pages(manager) + copy_ids = [block.block_id for block in manager.get_blocks("original").blocks[2]] + manager.free(original) + pool = manager.block_pool + pool.get_new_blocks(pool.get_num_free_blocks() - free_blocks) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + assert num_computed == 3 * HISPARSE_BLOCK_SIZE + assert manager.allocate_slots( + resumed, + num_new_tokens=len(tokens) - num_computed, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + ) + + resumed_blocks = manager.get_blocks("resumed").blocks + assert [block.block_id for block in resumed_blocks[2][:3]] == [ + copy_ids[page] if page in adopted_pages else pool.null_block.block_id + for page in range(3) + ] + gpu_ids = [ + block.block_id + for group in resumed_blocks[1:] + for block in group + if not block.is_null + ] + assert len(gpu_ids) == len(set(gpu_ids)) + + def test_hisparse_host_backed_request_accepts_local_prefix_hit(): """Local group-completion hits need no external host import allocation.""" manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index 9289e7c1f81f..d8be6c716af1 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -715,7 +715,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. """ @@ -726,7 +726,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 3f4f4e7bb164..3f5083a5ec0f 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) diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 056e19479a33..f83c9cabd08e 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -64,6 +64,8 @@ class _HiSparseRequestState: copies_recorded_blocks: int = 0 pinned_clean: set[int] = field(default_factory=set) unpinned_pages: set[int] = field(default_factory=set) + # Prefix pages whose GPU copies are adopted after the admitting allocation. + pages_to_adopt: int = 0 @dataclass @@ -193,7 +195,9 @@ def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None: state = self._get_request_state(request_id) state.valid_pages.update(range(num_host_pages)) state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages) - self._adopt_copies(request_id, state, host_blocks[:num_host_pages]) + # Adopting pins copies out of the free pool, so it must wait until the + # allocation admission counted on has run (see update_residency). + state.pages_to_adopt = max(state.pages_to_adopt, num_host_pages) state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages) def take_host_block_copies(self) -> tuple[KVCacheBlockCopy, ...]: @@ -343,7 +347,8 @@ def _unpin_page( return True def _unpin_clean_pages(self, request_id: str, state: _HiSparseRequestState) -> None: - for page_idx in sorted(state.pinned_clean): + # Tail first, so a prefix loses its tail pages before its head. + for page_idx in sorted(state.pinned_clean, reverse=True): self._unpin_page(request_id, state, page_idx) def _page_became_clean( @@ -389,6 +394,11 @@ def update_residency(self, request_id: str) -> None: if not self.resident_managers: return state = self._get_request_state(request_id) + if state.pages_to_adopt: + assert self.host_manager is not None + host_blocks = self.host_manager.req_to_blocks.get(request_id, ()) + self._adopt_copies(request_id, state, host_blocks[: state.pages_to_adopt]) + state.pages_to_adopt = 0 if not self._can_read_from_host(request_id): assert self.gpu_pool is not None if self.gpu_pool.get_num_free_blocks() >= self.transition_watermark: @@ -739,7 +749,9 @@ def free(self, request_id: str) -> None: blocks = manager.req_to_blocks.get(request_id) if not blocks: continue - for page_idx, block in enumerate(blocks): + # Tail first, so a prefix loses its tail pages before its head. + for page_idx in reversed(range(len(blocks))): + block = blocks[page_idx] if block.is_null: continue if page_idx in state.unpinned_pages: From 603c90ad93a2e65b5377fb1c90dfe651dfdccce8 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 29 Sep 2026 15:09:11 -0400 Subject: [PATCH 2/2] [HiSparse] Move the last-resort free-queue tier to a follow-up The crash fix only needs adoption to follow the allocation. Keeping more copies under pressure by making unpinned blocks last-resort in the core free queue is an optional improvement, proposed separately. 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 | 8 ++-- vllm/v1/core/block_pool.py | 4 +- vllm/v1/core/kv_cache_utils.py | 69 ++++++++++------------------ 4 files changed, 30 insertions(+), 76 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 6274d8f73baa..c2c86053f0f5 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -690,31 +690,6 @@ 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 accca79e61c1..4baef7716c52 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1103,9 +1103,7 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): ) -@pytest.mark.parametrize( - "free_blocks,adopted_pages", [(7, []), (9, [0, 1]), (10, [0, 1, 2])] -) +@pytest.mark.parametrize("free_blocks,adopted_pages", [(7, []), (10, [0, 1])]) def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( free_blocks, adopted_pages ): @@ -1114,8 +1112,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 copies only after every other free block and a - prefix's tail before its head; the copies it leaves are adopted. + now goes first, evicting a prefix's copies tail first; 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 d8be6c716af1..9289e7c1f81f 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -715,7 +715,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 stay behind every later-freed block and count as free. + caching: they join the tail of the free queue 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. """ @@ -726,7 +726,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_last_resort_n(released) + self.free_block_queue.append_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 3f5083a5ec0f..3f4f4e7bb164 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -262,10 +262,6 @@ 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. @@ -301,17 +297,6 @@ 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. @@ -338,7 +323,6 @@ 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). @@ -371,7 +355,6 @@ 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 @@ -397,7 +380,6 @@ 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 @@ -409,14 +391,28 @@ 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, ahead of any last-resort - blocks, and increase num_free_blocks by 1. + """Put a block back into the free list and increase + num_free_blocks by 1. Args: block: The block to append. """ - self.append_n([block]) + 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 def prepend_n(self, blocks: list[KVCacheBlock]) -> None: """Put a list of blocks at the front of the free list.""" @@ -440,43 +436,28 @@ 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, ahead of any - last-resort blocks. + """Put a list of blocks back into the free list 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] - 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" + 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" + ) # 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 anchor - last_block.next_free_block = anchor - anchor.prev_free_block = last_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 self.num_free_blocks += len(blocks)