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

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).

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)

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

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 <blocks> 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 <blocks> to the anchor
+        last_block.next_free_block = anchor
+        anchor.prev_free_block = last_block
 
         self.num_free_blocks += len(blocks)
 

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: