[Bugfix][HiSparse] Fix a chunked-prefill preemption livelock (#59494)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

tests/v1/core/test_prefix_caching.py

@@ -1259,6 +1259,63 @@ def test_hisparse_cleanup_publishes_only_finalized_terminal_pages(
     assert not coordinator.has_pending_work()
 
 
+@pytest.mark.parametrize("num_pages,reads_host", [(4, True), (3, False)])
+def test_hisparse_prefill_reads_host_once_it_fills_admission_window(
+    num_pages, reads_host
+):
+    """A prefill switches to host reads once it fills its admission window.
+
+    Admission caps each request's resident pages at the in-flight window,
+    assuming older pages move to host, but pinned pages were released only
+    once free blocks fell below a fixed watermark. A chunked prefill larger
+    than the watermark outgrew the pool first and was preempted.
+    """
+    window_pages = 4
+    manager = make_hisparse_kv_cache_manager(
+        32,
+        16,
+        enable_caching=True,
+        max_in_flight_tokens=window_pages * HISPARSE_BLOCK_SIZE,
+    )
+    coordinator = get_hisparse_coordinator(manager)
+    tokens = list(range(2 * window_pages * HISPARSE_BLOCK_SIZE))
+    request = make_request("prefill", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, num_pages * HISPARSE_BLOCK_SIZE) is not None
+    pool = coordinator.gpu_pool
+    assert pool is not None
+    assert pool.get_num_free_blocks() >= coordinator.transition_watermark
+
+    hot_manager = coordinator.hot_managers[0]
+    assert (request.request_id in hot_manager.hot_required) == reads_host
+
+
+def test_hisparse_preempted_prefill_resumes_from_durable_prefix():
+    """A prefill preempted mid-way must resume from the pages already on host.
+
+    Host writes trail a chunked prefill by a chunk, and host pages were only
+    published once the whole computed prefix was durable. Each chunk moved
+    that target past the finished writes, so nothing was published before the
+    prefill ended: a prefill preempted each time it outgrew the GPU pool
+    recomputed from the start forever.
+    """
+    manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)
+    coordinator = get_hisparse_coordinator(manager)
+    chunk = 2 * HISPARSE_BLOCK_SIZE
+    tokens = list(range(3 * chunk))
+    request = make_request("long", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.allocate_slots(request, chunk) is not None
+    request.num_computed_tokens = chunk
+    first_chunk_writes = coordinator.build_offload_command().page_transfers
+    assert manager.allocate_slots(request, chunk) is not None
+    request.num_computed_tokens = 2 * chunk
+    acks = {transfer.transfer_id: 1 for transfer in first_chunk_writes}
+    coordinator.update_spills(acks, acks)
+    manager.free(request)
+
+    resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert manager.get_computed_blocks(resumed)[1] == chunk
+
+
 def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity():
     """Old completions retain their pages through pressure and request-ID reuse."""
     manager = make_hisparse_kv_cache_manager(32, 8, enable_caching=True)
@@ -1285,7 +1342,9 @@ def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity()
     assert not {block.block_id for block in old_blocks[:2]} & {
         block.block_id for block in pressure + new_blocks
     }
-    assert all(block.block_hash is None for block in old_blocks)
+    # Only the page every worker finished writing is published.
+    assert old_blocks[0].block_hash is not None
+    assert all(block.block_hash is None for block in old_blocks[1:])
     coordinator.update_spills({}, partial)
     repeated = make_request("probe", tokens, HISPARSE_BLOCK_SIZE, sha256)
     assert manager.get_computed_blocks(repeated)[1] == 2 * HISPARSE_BLOCK_SIZE

vllm/v1/core/single_type_kv_cache_manager.py

@@ -112,7 +112,7 @@ class SingleTypeKVCacheManager(ABC):
         self.kv_cache_spec = kv_cache_spec
         self.block_pool = block_pool
         self.enable_caching = enable_caching
-        self._max_admission_blocks_per_request = max_admission_blocks_per_request
+        self.max_admission_blocks_per_request = max_admission_blocks_per_request
         # Record newly allocated block ids only when worker-side zeroing will
         # consume them and this manager holds a spec type that gets zeroed.
         self._record_new_block_ids = (
@@ -202,7 +202,7 @@ class SingleTypeKVCacheManager(ABC):
                 model in spec decode). w/o spec decode, it is num_tokens;
                 with spec decode, it is num_tokens - num_lookahead_tokens.
             apply_admission_cap: If True, clamp by `num_required_blocks` by
-                `_max_admission_blocks_per_request`for recycling-aware specs
+                `max_admission_blocks_per_request`for recycling-aware specs
                 (SWA, chunked-local).
             prefill_end: The token index the request's prefill ends at, the
                 same value the scheduler splits chunks against. Mamba reserves
@@ -214,7 +214,7 @@ class SingleTypeKVCacheManager(ABC):
 
         """
         num_required_blocks = cdiv(num_tokens, self.block_size)
-        if apply_admission_cap and self._max_admission_blocks_per_request is not None:
+        if apply_admission_cap and self.max_admission_blocks_per_request is not None:
             # Recycling-aware specs (SWA, chunked-local) cap the per-request
             # reservation here so admission matches the startup pool sizer
             # (`SlidingWindowSpec.max_admission_blocks_per_request` / its
@@ -225,7 +225,7 @@ class SingleTypeKVCacheManager(ABC):
             # Drift between the two would re-introduce the deadlock from
             # issue #39734 or, worse, mid-prefill OOM.
             num_required_blocks = min(
-                num_required_blocks, self._max_admission_blocks_per_request
+                num_required_blocks, self.max_admission_blocks_per_request
             )
         num_req_blocks = len(self.req_to_blocks.get(request_id, ()))
 
@@ -2488,6 +2488,8 @@ class HiSparseHotManager(_HiSparseAuxiliaryManager):
 class HiSparseResidentManager(_HiSparseAuxiliaryManager):
     """Track GPU-resident pages for otherwise host-backed KV."""
 
+    max_admission_blocks_per_request: int
+
     def get_num_blocks_to_allocate(
         self,
         request_id: str,
@@ -2514,8 +2516,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):
         host_pages = cdiv(num_local_computed_tokens, self.block_size)
         required = cdiv(num_tokens, self.block_size)
         if apply_admission_cap:
-            assert self._max_admission_blocks_per_request is not None
-            required = min(required, self._max_admission_blocks_per_request)
+            required = min(required, self.max_admission_blocks_per_request)
         return max(required - max(existing, host_pages), 0)
 
     def allocate_external_computed_blocks(

vllm/v1/hisparse/coordinator.py

@@ -318,6 +318,19 @@ class HiSparseCoordinator:
             manager.has_hot(request_id) for manager in self.hot_managers
         )
 
+    def _fills_admission_window(self, request_id: str) -> bool:
+        """Whether the request holds the resident pages admission reserved.
+
+        Admission caps a request's resident pages at the in-flight window,
+        assuming older pages move to host; keeping them pinned past it lets a
+        chunked prefill outgrow the pool it was admitted into.
+        """
+        return any(
+            len(manager.req_to_blocks.get(request_id, ()))
+            >= manager.max_admission_blocks_per_request
+            for manager in self.resident_managers
+        )
+
     def _resident_page_blocks(
         self, request_id: str, page_idx: int
     ) -> list[KVCacheBlock] | None:
@@ -390,8 +403,9 @@ class HiSparseCoordinator:
 
         A request that can read from host releases every clean sealed page to
         the pool. One that cannot keeps its pages pinned until the shared pool
-        runs low, then asks for a hot region. Pages remain pinned until that region
-        is allocated on a subsequent scheduling pass.
+        runs low, or until it holds the resident pages admission reserved for
+        it, then asks for a hot region. Pages remain pinned until that region is
+        allocated on a subsequent scheduling pass.
         """
         if not self.resident_managers:
             return
@@ -403,7 +417,10 @@ class HiSparseCoordinator:
             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:
+            if (
+                self.gpu_pool.get_num_free_blocks() >= self.transition_watermark
+                and not self._fills_admission_window(request_id)
+            ):
                 return
             for manager in self.hot_managers:
                 manager.require_hot(request_id)
@@ -453,47 +470,47 @@ class HiSparseCoordinator:
         *,
         replay_boundaries: Sequence[int],
     ) -> None:
-        """Publish host-source hashes only after their pages are durable."""
+        """Publish host-source hashes of the pages that are already durable."""
         manager = self.host_manager
         if manager is None:
             return
-        num_pages = num_computed_tokens // manager.block_size
         request_id = request.request_id
-        state = self._get_request_state(request_id)
-        if state.ready_prefix_pages >= num_pages:
-            manager.publish_blocks(
-                request,
-                num_computed_tokens,
-                retention_interval=retention_interval,
-                replay_boundaries=replay_boundaries,
-            )
-            self._record_copies(request_id, num_computed_tokens)
-            state.publication = None
-            return
-        state.publication = _PendingPublication(
+        self._get_request_state(request_id).publication = _PendingPublication(
             request=request,
             num_computed_tokens=num_computed_tokens,
-            num_pages=num_pages,
+            num_pages=num_computed_tokens // manager.block_size,
             retention_interval=retention_interval,
             replay_boundaries=replay_boundaries,
         )
+        self._publish_host_blocks_if_ready(request_id)
 
     def _publish_host_blocks_if_ready(self, request_id: str) -> None:
+        """Publish the durable prefix of the pending publication.
+
+        Host writes trail a long prefill by about a chunk, so waiting for the
+        whole computed prefix would publish nothing until the prefill ends, and
+        a prefill preempted before then would recompute from the start.
+        """
         state = self.request_states.get(request_id)
         if state is None:
             return
         publication = state.publication
-        if publication is None or state.ready_prefix_pages < publication.num_pages:
+        if publication is None:
             return
         assert self.host_manager is not None
+        num_tokens = min(
+            publication.num_computed_tokens,
+            state.ready_prefix_pages * self.host_manager.block_size,
+        )
         self.host_manager.publish_blocks(
             publication.request,
-            publication.num_computed_tokens,
+            num_tokens,
             retention_interval=publication.retention_interval,
             replay_boundaries=publication.replay_boundaries,
         )
-        self._record_copies(request_id, publication.num_computed_tokens)
-        state.publication = None
+        self._record_copies(request_id, num_tokens)
+        if state.ready_prefix_pages >= publication.num_pages:
+            state.publication = None
 
     def record_pending_host_import(self, request_id: str, num_tokens: int) -> None:
         """Note a prefix an external load is populating in host pages."""