@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ff1b87cca2

[Bugfix][HiSparse] Preserve host prefix publication after request completion (#59007)

Signed-off-by: Yueh-Ting Chen <yueh.ting.chen@gmail.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Yueh-Ting (eop) Chen committed Oct 1, 2026 at 00:14 UTC ff1b87cca25690fef6bd12667fd0d26690d949af

3 files changed +223

tests/v1/core/test_prefix_caching.py

+75

@@ -1178,6 +1178,81 @@ def test_hisparse_finished_request_leaves_free_copies():

1178 assert not coordinator.spills_to_send

1179

1180

1181 +@pytest.mark.parametrize(

1182 + "status,expected_cached_pages",

1183 + [

1184 + (RequestStatus.FINISHED_STOPPED, 2),

1185 + (RequestStatus.FINISHED_ABORTED, 2),

1186 + (RequestStatus.RUNNING, 0),

1187 + ],

1188 + ids=["stopped", "aborted", "running"],

1189 +)

1190 +def test_hisparse_cleanup_publishes_only_finalized_terminal_pages(

1191 + status, expected_cached_pages

1192 +):

1193 + """Only terminal cleanup may publish completed, finalized pages."""

1194 + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)

1195 + coordinator = get_hisparse_coordinator(manager)

1196 + request = make_request("terminal", list(range(33)), HISPARSE_BLOCK_SIZE, sha256)

1197 + request.spec_token_ids = list(range(33, 49))

1198 + assert manager.allocate_slots(request, 49, num_lookahead_tokens=16) is not None

1199 + host_blocks = list(manager.get_blocks(request.request_id).blocks[0])

1200 + command = coordinator.build_offload_command()

1201 + counts = {transfer.transfer_id: 1 for transfer in command.page_transfers}

1202 + assert len(counts) == 2

1203 + request.status = status

1204 + manager.free(request)

1205 + coordinator.update_spills(counts, counts)

1206 + assert all(

1207 + block.block_hash is not None for block in host_blocks[:expected_cached_pages]

1208 + )

1209 + assert all(

1210 + block.block_hash is None for block in host_blocks[expected_cached_pages:]

1211 + )

1212 + assert all(block.ref_cnt == 0 for block in host_blocks)

1213 + assert not coordinator.has_pending_work()

1214 +

1215 +

1216 +def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity():

1217 + """Old completions retain their pages through pressure and request-ID reuse."""

1218 + manager = make_hisparse_kv_cache_manager(32, 8, enable_caching=True)

1219 + coordinator = get_hisparse_coordinator(manager)

1220 + host_pool = coordinator.get_host_block_pool()

1221 + tokens = list(range(2 * HISPARSE_BLOCK_SIZE + 1))

1222 + request = make_request("reused", tokens, HISPARSE_BLOCK_SIZE, sha256)

1223 + assert manager.allocate_slots(request, len(tokens)) is not None

1224 + old_blocks = list(manager.get_blocks(request.request_id).blocks[0])

1225 + command = coordinator.build_offload_command()

1226 + counts = {transfer.transfer_id: 2 for transfer in command.page_transfers}

1227 + partial = dict.fromkeys(counts, 1)

1228 + first_completions = partial.copy()

1229 + first_completions[next(iter(counts))] = 2

1230 + coordinator.update_spills(counts, first_completions)

1231 + request.status = RequestStatus.FINISHED_LENGTH_CAPPED

1232 + manager.free(request)

1233 + replacement = make_request(

1234 + "reused", list(range(100, 117)), HISPARSE_BLOCK_SIZE, sha256

1235 + )

1236 + assert manager.allocate_slots(replacement, 17) is not None

1237 + new_blocks = list(manager.get_blocks(replacement.request_id).blocks[0])

1238 + pressure = host_pool.get_new_blocks(host_pool.get_num_free_blocks())

1239 + assert not {block.block_id for block in old_blocks[:2]} & {

1240 + block.block_id for block in pressure + new_blocks

1241 + }

1242 + assert all(block.block_hash is None for block in old_blocks)

1243 + coordinator.update_spills({}, partial)

1244 + repeated = make_request("probe", tokens, HISPARSE_BLOCK_SIZE, sha256)

1245 + assert manager.get_computed_blocks(repeated)[1] == 2 * HISPARSE_BLOCK_SIZE

1246 + assert all(block.block_hash is None for block in new_blocks)

1247 + assert all(block.ref_cnt == 0 for block in old_blocks[:2])

1248 + coordinator.update_spills({}, partial)

1249 + assert all(block.ref_cnt == 0 for block in old_blocks[:2])

1250 + host_pool.free_blocks(pressure)

1251 + manager.free(replacement)

1252 + _publish_hisparse_pages(manager)

1253 + assert not coordinator.has_pending_work()

1254 +

1255 +

1256 def test_hisparse_reset_prefix_cache_drops_copies():

1257 """Reset must forget GPU copies along with the host hashes they mirror."""

1258 manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True)

tests/v1/e2e/general/test_hisparse.py

+99

@@ -1,6 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import time

5 +

6 import prometheus_client

7 import pytest

8 import torch

@@ -18,6 +20,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (

20 MultiConnector,

21 )

22 from vllm.platforms import current_platform

23 +from vllm.sampling_params import SamplingParams

24 from vllm.utils.gpu_sync_debug import gpu_sync_allowed

25 from vllm.v1.hisparse.coordinator import get_hisparse_coordinator

@@ -301,3 +304,99 @@ def test_hisparse_host_exhaustion_defers_requests(

304 assert checked_steps > 0

305 assert refusals > 0

306 assert actual == expected

307 +

308 +

309 +@pytest.mark.skipif(

310 + not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA"

311 +)

312 +@fork_new_process_for_each_test

313 +def test_hisparse_terminal_prefix_reuse(

314 + monkeypatch: pytest.MonkeyPatch,

315 + vllm_runner: type[VllmRunner],

316 +):

317 + """Finished requests publish mirrored host KV after a late acknowledgement."""

318 + capability = current_platform.get_device_capability()

319 + if capability is None or capability.major < 9:

320 + pytest.skip("Sparse MLA requires Hopper or newer")

321 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

322 + monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")

323 + target = [1000 + i % 64 for i in range(257)]

324 + sampling = SamplingParams(temperature=0, max_tokens=1, ignore_eos=True)

325 + with vllm_runner(

326 + MODEL,

327 + load_format="dummy",

328 + hf_overrides=_shrink_config,

329 + attention_config=AttentionConfig(

330 + hisparse_config=HiSparseConfig(

331 + device_buffer_size=512, eager_host_mirror=True

332 + )

333 + ),

334 + kv_transfer_config=KVTransferConfig(

335 + kv_connector="HiSparseConnector",

336 + kv_role="kv_both",

337 + kv_connector_extra_config={"host_pool_gib": 0.01},

338 + ),

339 + block_size=64,

340 + max_model_len=320,

341 + max_num_batched_tokens=512,

342 + max_num_seqs=2,

343 + num_gpu_blocks_override=128,

344 + enable_prefix_caching=True,

345 + enable_chunked_prefill=True,

346 + enforce_eager=True,

347 + ) as runner:

348 + engine = runner.llm.llm_engine.engine_core.engine_core

349 + coordinator = get_hisparse_coordinator(engine.scheduler.kv_cache_manager)

350 + worker = _get_hisparse_worker(runner)

351 + original_updates = worker.take_transfer_updates

352 + original_free = coordinator.free

353 + held_completions: list[int] = []

354 + update_calls = 0

355 + terminal_pending_pages = 0

356 +

357 + def take_updates():

358 + nonlocal update_calls

359 + with gpu_sync_allowed():

360 + torch.accelerator.synchronize()

361 + enqueued, completed = original_updates()

362 + held_completions.extend(completed)

363 + update_calls += 1

364 + if update_calls <= 2:

365 + return enqueued, []

366 + completed = held_completions.copy()

367 + held_completions.clear()

368 + return enqueued, completed

369 +

370 + def free(request_id):

371 + nonlocal terminal_pending_pages

372 + state = coordinator.request_states.get(request_id)

373 + if state is not None:

374 + terminal_pending_pages += len(state.pending_pages)

375 + original_free(request_id)

376 +

377 + def drain_pending_work():

378 + deadline = time.monotonic() + 10

379 + while coordinator.has_pending_work():

380 + assert time.monotonic() < deadline, "HiSparse transfers did not drain"

381 + runner.llm.llm_engine.step()

382 +

383 + monkeypatch.setattr(worker, "take_transfer_updates", take_updates)

384 + monkeypatch.setattr(coordinator, "free", free)

385 + [first] = runner.llm.generate(

386 + [{"prompt_token_ids": target}], sampling, use_tqdm=False

387 + )

388 + assert terminal_pending_pages == 4

389 + assert first.num_cached_tokens == 0

390 + assert coordinator.has_pending_work()

391 + drain_pending_work()

392 + assert not coordinator.has_pending_work()

393 + assert not held_completions

394 + assert update_calls > 2

395 +

396 + [repeated] = runner.llm.generate(

397 + [{"prompt_token_ids": target}], sampling, use_tqdm=False

398 + )

399 + assert repeated.num_cached_tokens == 256

400 + assert repeated.outputs[0].token_ids == first.outputs[0].token_ids

401 + drain_pending_work()

402 + assert not coordinator.has_pending_work()

vllm/v1/hisparse/coordinator.py

+49

@@ -44,6 +44,8 @@ class _PendingPublication:

44 num_pages: int

45 retention_interval: int | None

46 replay_boundaries: Sequence[int]

47 + detached_blocks: list[KVCacheBlock] | None = None

48 + num_cached_blocks: int = 0

51 @dataclass

@@ -724,6 +726,10 @@ class HiSparseCoordinator:

726 if pending.restore:

727 self._record_copy(request_id, page_idx, pending.host_block)

728 completed_request_ids.add(request_id)

729 + elif pending.request_state.publication is not None:

730 + if not pending.restore:

731 + pending.request_state.valid_pages.add(page_idx)

732 + self._publish_detached_blocks(pending.request_state)

733 assert self.host_manager is not None

734 self.host_manager.block_pool.free_blocks([pending.host_block])

735 for request_id in completed_request_ids:

@@ -732,12 +738,55 @@ class HiSparseCoordinator:

738 state.ready_prefix_pages += 1

739 self._publish_host_blocks_if_ready(request_id)

740

741 + def _publish_detached_blocks(self, state: _HiSparseRequestState) -> None:

742 + publication = state.publication

743 + if publication is None or state.pending_pages:

744 + return

745 + blocks = publication.detached_blocks

746 + if blocks is None:

747 + return

748 + assert self.host_manager is not None

749 + while state.ready_prefix_pages in state.valid_pages:

750 + state.ready_prefix_pages += 1

751 + # The host group uses full-attention retention. Only sealed pages with

752 + # completed copies can outlive the request, including on cancellation.

753 + self.host_manager.block_pool.cache_full_blocks(

754 + request=publication.request,

755 + blocks=blocks,

756 + num_cached_blocks=publication.num_cached_blocks,

757 + num_full_blocks=min(state.ready_prefix_pages, publication.num_pages),

758 + block_size=self.host_manager.block_size,

759 + kv_cache_group_id=self.host_manager.kv_cache_group_id,

760 + )

761 + self.host_manager.block_pool.free_blocks(reversed(blocks))

762 + state.publication = None

763 +

764 def free(self, request_id: str) -> None:

765 """Detach the request; its clean pages stay readable copies in the pool."""

766 self._pending_imports.pop(request_id, None)

767 state = self.request_states.pop(request_id, None)

768 if state is None:

769 return

770 + publication = state.publication

771 + if (

772 + publication is not None

773 + and publication.request.is_finished()

774 + and state.pending_pages

775 + ):

776 + assert self.host_manager is not None

777 + publication.detached_blocks = self.host_manager.req_to_blocks[request_id][

778 + : publication.num_pages

779 + ]

780 + publication.num_cached_blocks = self.host_manager.num_cached_block.get(

781 + request_id, 0

782 + )

783 + # Completed pages need leases too: another allocation may evict

784 + # them before the remaining copies finish and the prefix publishes.

785 + self.host_manager.block_pool.touch(publication.detached_blocks)

786 + else:

787 + # Preemption can cancel a scheduled forward after its spills were

788 + # planned, so its copies do not prove that the KV was computed.

789 + state.publication = None

790 for manager in self.resident_managers:

791 blocks = manager.req_to_blocks.get(request_id)

792 if not blocks: