From 3eb6cec22ad9bb098393021b956feaa97081fc85 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Wed, 30 Sep 2026 15:43:09 -0400 Subject: [PATCH] [Bugfix][HiSparse] Adopt GPU prefix copies after the hit's allocation (#59282) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/core/test_prefix_caching.py | 46 ++++++++++++++++++++++++++++ vllm/v1/hisparse/coordinator.py | 18 +++++++++-- 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index fab623eb84..e4c54e880c 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1139,6 +1139,52 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): ) +@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 +): + """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 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)) + 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/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 26357e0a2a..01a6392964 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -66,6 +66,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 @@ -195,7 +197,9 @@ class HiSparseCoordinator: 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, ...]: @@ -345,7 +349,8 @@ class HiSparseCoordinator: 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( @@ -391,6 +396,11 @@ class HiSparseCoordinator: 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: @@ -791,7 +801,9 @@ class HiSparseCoordinator: 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: