[Bugfix][HiSparse] Adopt GPU prefix copies after the hit's allocation

tests/v1/core/test_prefix_caching.py

@@ -1103,6 +1103,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)

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: