[Bugfix][HiSparse] Preserve host prefix publication after request completion

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)

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

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: