diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index e4c54e880c..027bb67b18 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/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 diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index c8b9c54435..a097c9e924 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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( diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 01a6392964..1a783b506b 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/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."""