From b62dc7b92f965fe74f4c140790a897c96e542fae Mon Sep 17 00:00:00 2001 From: Yueh-Ting Chen Date: Mon, 28 Sep 2026 17:19:39 +0800 Subject: [PATCH] [Bugfix][HiSparse] Preserve host prefix publication after request completion Co-authored-by: OpenAI Codex Signed-off-by: Yueh-Ting Chen --- tests/v1/core/test_prefix_caching.py | 75 ++++++++++++++++++++ tests/v1/e2e/general/test_hisparse.py | 99 +++++++++++++++++++++++++++ vllm/v1/hisparse/coordinator.py | 49 +++++++++++++ 3 files changed, 223 insertions(+) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index c8d42665f08d..fab623eb843d 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1178,6 +1178,81 @@ def test_hisparse_finished_request_leaves_free_copies(): assert not coordinator.spills_to_send +@pytest.mark.parametrize( + "status,expected_cached_pages", + [ + (RequestStatus.FINISHED_STOPPED, 2), + (RequestStatus.FINISHED_ABORTED, 2), + (RequestStatus.RUNNING, 0), + ], + ids=["stopped", "aborted", "running"], +) +def test_hisparse_cleanup_publishes_only_finalized_terminal_pages( + status, expected_cached_pages +): + """Only terminal cleanup may publish completed, finalized pages.""" + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + coordinator = get_hisparse_coordinator(manager) + request = make_request("terminal", list(range(33)), HISPARSE_BLOCK_SIZE, sha256) + request.spec_token_ids = list(range(33, 49)) + assert manager.allocate_slots(request, 49, num_lookahead_tokens=16) is not None + host_blocks = list(manager.get_blocks(request.request_id).blocks[0]) + command = coordinator.build_offload_command() + counts = {transfer.transfer_id: 1 for transfer in command.page_transfers} + assert len(counts) == 2 + request.status = status + manager.free(request) + coordinator.update_spills(counts, counts) + assert all( + block.block_hash is not None for block in host_blocks[:expected_cached_pages] + ) + assert all( + block.block_hash is None for block in host_blocks[expected_cached_pages:] + ) + assert all(block.ref_cnt == 0 for block in host_blocks) + assert not coordinator.has_pending_work() + + +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) + coordinator = get_hisparse_coordinator(manager) + host_pool = coordinator.get_host_block_pool() + tokens = list(range(2 * HISPARSE_BLOCK_SIZE + 1)) + request = make_request("reused", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.allocate_slots(request, len(tokens)) is not None + old_blocks = list(manager.get_blocks(request.request_id).blocks[0]) + command = coordinator.build_offload_command() + counts = {transfer.transfer_id: 2 for transfer in command.page_transfers} + partial = dict.fromkeys(counts, 1) + first_completions = partial.copy() + first_completions[next(iter(counts))] = 2 + coordinator.update_spills(counts, first_completions) + request.status = RequestStatus.FINISHED_LENGTH_CAPPED + manager.free(request) + replacement = make_request( + "reused", list(range(100, 117)), HISPARSE_BLOCK_SIZE, sha256 + ) + assert manager.allocate_slots(replacement, 17) is not None + new_blocks = list(manager.get_blocks(replacement.request_id).blocks[0]) + pressure = host_pool.get_new_blocks(host_pool.get_num_free_blocks()) + 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) + coordinator.update_spills({}, partial) + repeated = make_request("probe", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert manager.get_computed_blocks(repeated)[1] == 2 * HISPARSE_BLOCK_SIZE + assert all(block.block_hash is None for block in new_blocks) + assert all(block.ref_cnt == 0 for block in old_blocks[:2]) + coordinator.update_spills({}, partial) + assert all(block.ref_cnt == 0 for block in old_blocks[:2]) + host_pool.free_blocks(pressure) + manager.free(replacement) + _publish_hisparse_pages(manager) + assert not coordinator.has_pending_work() + + def test_hisparse_reset_prefix_cache_drops_copies(): """Reset must forget GPU copies along with the host hashes they mirror.""" manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index 21f28d0b6b29..573b424ea72e 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time + import prometheus_client import pytest import torch @@ -18,6 +20,7 @@ MultiConnector, ) from vllm.platforms import current_platform +from vllm.sampling_params import SamplingParams from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.v1.hisparse.coordinator import get_hisparse_coordinator @@ -301,3 +304,99 @@ def count(*args, **kwargs): assert checked_steps > 0 assert refusals > 0 assert actual == expected + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" +) +@fork_new_process_for_each_test +def test_hisparse_terminal_prefix_reuse( + monkeypatch: pytest.MonkeyPatch, + vllm_runner: type[VllmRunner], +): + """Finished requests publish mirrored host KV after a late acknowledgement.""" + capability = current_platform.get_device_capability() + if capability is None or capability.major < 9: + pytest.skip("Sparse MLA requires Hopper or newer") + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") + target = [1000 + i % 64 for i in range(257)] + sampling = SamplingParams(temperature=0, max_tokens=1, ignore_eos=True) + with vllm_runner( + MODEL, + load_format="dummy", + hf_overrides=_shrink_config, + attention_config=AttentionConfig( + hisparse_config=HiSparseConfig( + device_buffer_size=512, eager_host_mirror=True + ) + ), + kv_transfer_config=KVTransferConfig( + kv_connector="HiSparseConnector", + kv_role="kv_both", + kv_connector_extra_config={"host_pool_gib": 0.01}, + ), + block_size=64, + max_model_len=320, + max_num_batched_tokens=512, + max_num_seqs=2, + num_gpu_blocks_override=128, + enable_prefix_caching=True, + enable_chunked_prefill=True, + enforce_eager=True, + ) as runner: + engine = runner.llm.llm_engine.engine_core.engine_core + coordinator = get_hisparse_coordinator(engine.scheduler.kv_cache_manager) + worker = _get_hisparse_worker(runner) + original_updates = worker.take_transfer_updates + original_free = coordinator.free + held_completions: list[int] = [] + update_calls = 0 + terminal_pending_pages = 0 + + def take_updates(): + nonlocal update_calls + with gpu_sync_allowed(): + torch.accelerator.synchronize() + enqueued, completed = original_updates() + held_completions.extend(completed) + update_calls += 1 + if update_calls <= 2: + return enqueued, [] + completed = held_completions.copy() + held_completions.clear() + return enqueued, completed + + def free(request_id): + nonlocal terminal_pending_pages + state = coordinator.request_states.get(request_id) + if state is not None: + terminal_pending_pages += len(state.pending_pages) + original_free(request_id) + + def drain_pending_work(): + deadline = time.monotonic() + 10 + while coordinator.has_pending_work(): + assert time.monotonic() < deadline, "HiSparse transfers did not drain" + runner.llm.llm_engine.step() + + monkeypatch.setattr(worker, "take_transfer_updates", take_updates) + monkeypatch.setattr(coordinator, "free", free) + [first] = runner.llm.generate( + [{"prompt_token_ids": target}], sampling, use_tqdm=False + ) + assert terminal_pending_pages == 4 + assert first.num_cached_tokens == 0 + assert coordinator.has_pending_work() + drain_pending_work() + assert not coordinator.has_pending_work() + assert not held_completions + assert update_calls > 2 + + [repeated] = runner.llm.generate( + [{"prompt_token_ids": target}], sampling, use_tqdm=False + ) + assert repeated.num_cached_tokens == 256 + assert repeated.outputs[0].token_ids == first.outputs[0].token_ids + drain_pending_work() + assert not coordinator.has_pending_work() diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 0ca0d40c88b5..26357e0a2acc 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -44,6 +44,8 @@ class _PendingPublication: num_pages: int retention_interval: int | None replay_boundaries: Sequence[int] + detached_blocks: list[KVCacheBlock] | None = None + num_cached_blocks: int = 0 @dataclass @@ -724,6 +726,10 @@ def _complete_host_writes(self) -> None: if pending.restore: self._record_copy(request_id, page_idx, pending.host_block) completed_request_ids.add(request_id) + elif pending.request_state.publication is not None: + if not pending.restore: + pending.request_state.valid_pages.add(page_idx) + self._publish_detached_blocks(pending.request_state) assert self.host_manager is not None self.host_manager.block_pool.free_blocks([pending.host_block]) for request_id in completed_request_ids: @@ -732,12 +738,55 @@ def _complete_host_writes(self) -> None: state.ready_prefix_pages += 1 self._publish_host_blocks_if_ready(request_id) + def _publish_detached_blocks(self, state: _HiSparseRequestState) -> None: + publication = state.publication + if publication is None or state.pending_pages: + return + blocks = publication.detached_blocks + if blocks is None: + return + assert self.host_manager is not None + while state.ready_prefix_pages in state.valid_pages: + state.ready_prefix_pages += 1 + # The host group uses full-attention retention. Only sealed pages with + # completed copies can outlive the request, including on cancellation. + self.host_manager.block_pool.cache_full_blocks( + request=publication.request, + blocks=blocks, + num_cached_blocks=publication.num_cached_blocks, + num_full_blocks=min(state.ready_prefix_pages, publication.num_pages), + block_size=self.host_manager.block_size, + kv_cache_group_id=self.host_manager.kv_cache_group_id, + ) + self.host_manager.block_pool.free_blocks(reversed(blocks)) + state.publication = None + def free(self, request_id: str) -> None: """Detach the request; its clean pages stay readable copies in the pool.""" self._pending_imports.pop(request_id, None) state = self.request_states.pop(request_id, None) if state is None: return + publication = state.publication + if ( + publication is not None + and publication.request.is_finished() + and state.pending_pages + ): + assert self.host_manager is not None + publication.detached_blocks = self.host_manager.req_to_blocks[request_id][ + : publication.num_pages + ] + publication.num_cached_blocks = self.host_manager.num_cached_block.get( + request_id, 0 + ) + # Completed pages need leases too: another allocation may evict + # them before the remaining copies finish and the prefix publishes. + self.host_manager.block_pool.touch(publication.detached_blocks) + else: + # Preemption can cancel a scheduled forward after its spills were + # planned, so its copies do not prove that the KV was computed. + state.publication = None for manager in self.resident_managers: blocks = manager.req_to_blocks.get(request_id) if not blocks: