@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 67e5b0acc9

[Bugfix][HiSparse][NIXL] Import full blocks without tail prefill on D (#57049)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>

Nicolò Lucchesi committed Sep 17, 2026 at 21:26 UTC 67e5b0acc9988afc50d019db64ad9da0dadaa15e

8 files changed +277 -53

tests/v1/core/test_prefix_caching.py

+95 -9

@@ -325,7 +325,7 @@ def test_hisparse_reports_when_context_is_fully_resident():

325

326

327 def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():

328 - """Keep indexer-only imports host-backed in the resident block table."""

328 + """Restore the final host page when completing an indexer-only import."""

329 manager = make_hisparse_kv_cache_manager(

330 32,

331 16,

@@ -360,11 +360,15 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():

360 source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks

361 assert len(source) == len(indexer) == len(resident) == 4

362 assert len(hot) == 2

363 - # The local prefix is adopted from shadow pages (GPU-resident), while the

364 - # externally imported page stays host-backed until its tail allocation.

363 assert not any(block.is_null for block in resident[:2])

366 - assert resident[2].is_null

364 + assert not resident[2].is_null

365 assert not resident[3].is_null

366 + coordinator = get_hisparse_coordinator(manager)

367 + assert not coordinator.build_offload_command().page_transfers

368 + coordinator.finish_host_import(resumed.request_id, failed=False)

369 + transfers = coordinator.build_offload_command().page_transfers

370 + assert len(transfers) == 1 and transfers[0].restore

371 + assert transfers[0].resident_block_ids == (resident[2].block_id,)

372

373

374 def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix():

@@ -465,7 +469,8 @@ def test_connector_completes_partial_prefix_without_importing_missing_host_kv(

469 )

470 allocated = manager.get_blocks(resumed.request_id).blocks

471 assert len(allocated[0]) == host_blocks

468 - assert allocated[2][1].is_null == (external > 0)

472 + assert allocated[2][1].is_null == (host_blocks > 2)

473 + assert not allocated[2][-1].is_null

474 manager.free(resumed)

475

476 # A surviving indexer offload must not make absent host KV look computed.

@@ -490,6 +495,85 @@ def allocate_external_prefix(

495 )

496

497

498 +@pytest.mark.parametrize("num_tokens", [1, 15, 16, 17, 31, 32, 33, 127])

499 +def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens):

500 + """Import every real token and restore the page used by last-token replay."""

501 + from tests.v1.kv_connector.unit.utils import create_vllm_config

502 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.pull_scheduler import (

503 + NixlPullConnectorScheduler,

504 + )

505 +

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

507 + config = create_vllm_config(block_size=HISPARSE_BLOCK_SIZE)

508 + connector = NixlPullConnectorScheduler(config, "test", manager.kv_cache_config)

509 + request = make_request(

510 + "import", list(range(num_tokens)), HISPARSE_BLOCK_SIZE, sha256

511 + )

512 + request.kv_transfer_params = {"do_remote_prefill": True}

513 + count, is_async = connector.get_num_new_matched_tokens(request, 0)

514 + assert count == num_tokens and is_async

515 + assert allocate_external_prefix(manager, request, count) is not None

516 + source, _, resident, _ = manager.get_blocks(request.request_id).blocks

517 + assert len(source) == len(resident) == (num_tokens + 15) // 16

518 + assert all(block.is_null for block in resident[:-1])

519 + assert not resident[-1].is_null

520 + coordinator = get_hisparse_coordinator(manager)

521 + assert not coordinator.build_offload_command().page_transfers

522 +

523 + coordinator.finish_host_import(request.request_id, failed=False)

524 + transfers = coordinator.build_offload_command().page_transfers

525 + assert len(transfers) == 1

526 + restore = transfers[0]

527 + assert restore.restore and not restore.after_forward

528 + assert restore.host_block_id == source[-1].block_id

529 + assert restore.resident_block_ids == (resident[-1].block_id,)

530 + request.num_computed_tokens = count

531 + scheduler = SimpleNamespace(

532 + connector=connector,

533 + kv_cache_manager=manager,

534 + failed_recving_kv_req_ids=set(),

535 + finished_recving_kv_req_ids={request.request_id},

536 + )

537 + Scheduler._update_waiting_for_remote_kv(scheduler, request)

538 + assert request.num_tokens - request.num_computed_tokens == 1

539 + assert manager.allocate_slots(request, num_new_tokens=1) is not None

540 + assert coordinator.build_row_mirrors([(request.request_id, num_tokens - 1, 1)])

541 + # A pending restore must not expose uninitialized GPU copies to prefix hits.

542 + assert all(resident[-1] not in copies for copies in coordinator.copies.values())

543 +

544 + coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1})

545 + assert resident[-1].ref_cnt == 2

546 + assert all(resident[-1] not in copies for copies in coordinator.copies.values())

547 + coordinator.update_spills({}, {restore.transfer_id: 1})

548 + assert resident[-1].ref_cnt == 1

549 + assert coordinator.request_states[request.request_id].valid_pages == set(

550 + range(num_tokens // HISPARSE_BLOCK_SIZE)

551 + )

552 + if num_tokens % HISPARSE_BLOCK_SIZE:

553 + assert source[-1].block_hash is None

554 + else:

555 + assert coordinator.copies[source[-1].block_hash] == (resident[-1],)

556 +

557 +

558 +def test_hisparse_aborted_tail_restore_retains_both_endpoints():

559 + """An abort cannot recycle storage still used by a queued tail restore."""

560 + manager = make_hisparse_kv_cache_manager(16, 16)

561 + request = make_request("import", list(range(17)), HISPARSE_BLOCK_SIZE, sha256)

562 + assert allocate_external_prefix(manager, request, 17) is not None

563 + source, _, resident, _ = manager.get_blocks(request.request_id).blocks

564 + host_tail, gpu_tail = source[-1], resident[-1]

565 + coordinator = get_hisparse_coordinator(manager)

566 + coordinator.finish_host_import(request.request_id, failed=False)

567 + restore = coordinator.build_offload_command().page_transfers[0]

568 + manager.free(request)

569 + assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1

570 + coordinator.update_spills({restore.transfer_id: 2}, {restore.transfer_id: 1})

571 + assert host_tail.ref_cnt == gpu_tail.ref_cnt == 1

572 + coordinator.update_spills({}, {restore.transfer_id: 1})

573 + assert host_tail.ref_cnt == gpu_tail.ref_cnt == 0

574 + assert not coordinator.has_pending_work()

575 +

576 +

577 @pytest.mark.parametrize("enable_caching", [False, True])

578 def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_caching):

579 """Later admissions must not remove a scheduled request's only readable tier."""

@@ -669,7 +753,7 @@ def test_hisparse_host_cow_copy_is_drained_without_a_gpu_pool():

753 def test_hisparse_inflight_host_import_reserves_remaining_gpu_pages():

754 """An in-flight host import must reserve its unwritten resident pages."""

755 manager = make_hisparse_kv_cache_manager(

672 - 11,

756 + 12,

757 16,

758 transfer_device_cache=True,

759 )

@@ -1127,9 +1211,10 @@ def test_hisparse_external_import_uses_hard_gpu_footprint():

1211 assert len(source) == num_prompt_blocks

1212 assert len(indexer) == num_prompt_blocks

1213 assert len(resident) == num_prompt_blocks

1130 - assert all(block.is_null for block in resident)

1214 + assert all(block.is_null for block in resident[:-1])

1215 + assert not resident[-1].is_null

1216 assert len(hot) == 2

1132 - assert manager.block_pool.get_num_free_blocks() == 1

1217 + assert manager.block_pool.get_num_free_blocks() == 0

1218

1219

1220 def test_hisparse_external_import_survives_capacity_retry():

@@ -1149,7 +1234,8 @@ def test_hisparse_external_import_survives_capacity_retry():

1234 manager.free(first)

1235 assert allocate_external_prefix(manager, second, len(tokens)) is not None

1236 _, _, resident, hot = manager.get_blocks(second.request_id).blocks

1152 - assert all(block.is_null for block in resident)

1237 + assert all(block.is_null for block in resident[:-1])

1238 + assert not resident[-1].is_null

1239 assert len(hot) == 2

1240

1241

tests/v1/worker/test_utils.py

+65

@@ -531,6 +531,71 @@ def test_hisparse_eager_mirror_records_transfer_without_page_copy():

531 worker._enqueue_transfers.assert_not_called()

532

533

534 +@pytest.mark.parametrize("is_host_writer", [False, True])

535 +@pytest.mark.parametrize("device", ["cpu", "cuda"])

536 +def test_hisparse_tail_restore_preserves_imported_rows(

537 + monkeypatch, request, is_host_writer, device

538 +):

539 + """Restore externally pinned host rows on every rank under sync checking."""

540 + if device == "cuda" and not torch.cuda.is_available():

541 + pytest.skip("CUDA is required for the device copy check")

542 + worker = _make_hisparse_worker()

543 + worker.is_host_writer = is_host_writer

544 + worker.kernel_block_size = 4

545 + worker.host_caches = tuple(

546 + torch.arange(24, dtype=torch.float32).reshape(12, 2) + layer * 100

547 + for layer in range(2)

548 + )

549 + if device == "cuda":

550 + # Match production's cudaHostRegister pool, which torch.is_pinned()

551 + # does not recognize, rather than torch's caching pinned allocator.

552 + host_pool, registered = hisparse_runtime_module.allocate_pinned_host_pool(

553 + sum(cache.nbytes for cache in worker.host_caches)

554 + )

555 + request.addfinalizer(

556 + lambda: hisparse_runtime_module.release_pinned_state([], [registered])

557 + )

558 + host_caches = host_pool.view(torch.float32).view(2, 12, 2).unbind()

559 + for destination, source in zip(host_caches, worker.host_caches):

560 + destination.copy_(source)

561 + worker.host_caches = host_caches

562 + worker.resident_caches = tuple(

563 + torch.full((4, 4, 2), -1.0, device=device) for _ in range(2)

564 + )

565 + worker.cache_handles = [

566 + SimpleNamespace(

567 + runtime=SimpleNamespace(resident_source_index=i, eager_host_mirror=True)

568 + )

569 + for i in range(2)

570 + ]

571 + worker._record_transfer_completion = MagicMock()

572 + stream = torch.cuda.current_stream() if device == "cuda" else MagicMock()

573 + monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)

574 + transfer = SparseKVPageTransfer(7, 2, (1, 3), False, restore=True)

575 +

576 + restore_pages = worker._restore_pages

577 + if device == "cuda":

578 + import vllm.utils.gpu_sync_debug as gpu_sync_debug

579 +

580 + monkeypatch.setattr(gpu_sync_debug, "_SYNC_CHECK_MODE", "error")

581 + monkeypatch.setattr(gpu_sync_debug, "_sync_check_enabled", False)

582 + gpu_sync_debug.enable_gpu_sync_check()

583 + restore_pages = gpu_sync_debug.with_gpu_sync_check(restore_pages)

584 + restore_pages([transfer])

585 +

586 + # Appending into unused capacity must preserve every imported prompt row.

587 + for cache, host, block in zip(

588 + worker.resident_caches, worker.host_caches, transfer.resident_block_ids

589 + ):

590 + assert torch.equal(cache[block].cpu(), host[8:12])

591 + cache[block, 2:] = 999

592 + assert torch.equal(cache[block, :2].cpu(), host[8:10])

593 + assert torch.all(cache[0] == -1)

594 + worker._record_transfer_completion.assert_called_once_with(

595 + [transfer], stream=stream

596 + )

597 +

598 +

599 def test_hisparse_lazy_mirror_copies_transfer_pages():

600 worker = _make_hisparse_worker()

601 worker.cache_handles = [

vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py

+45 -10

@@ -22,6 +22,7 @@ from vllm.distributed.parallel_state import (

22 get_tensor_model_parallel_rank,

23 get_tp_group,

24 )

25 +from vllm.utils.gpu_sync_debug import gpu_sync_allowed

26 from vllm.utils.torch_utils import current_stream

27 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy

28 from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX

@@ -414,11 +415,18 @@ class HiSparseConnectorWorker:

415 transfers = (

416 metadata.command.page_transfers if metadata.command is not None else []

417 )

418 + self._restore_pages([transfer for transfer in transfers if transfer.restore])

419 self._post_forward_transfers = [

418 - transfer for transfer in transfers if transfer.after_forward

420 + transfer

421 + for transfer in transfers

422 + if not transfer.restore and transfer.after_forward

423 ]

424 self._submit_transfers(

421 - [transfer for transfer in transfers if not transfer.after_forward]

425 + [

426 + transfer

427 + for transfer in transfers

428 + if not transfer.restore and not transfer.after_forward

429 + ]

430 )

431 self._pending_invalid_block_ids.extend(metadata.source_block_ids)

432 if request_state_indices is not None:

@@ -708,11 +716,17 @@ class HiSparseConnectorWorker:

716 self._submitted_mirror_layers.update(pending_layers)

717

718 def _record_transfer_completion(

711 - self, transfers: list[SparseKVPageTransfer]

719 + self,

720 + transfers: list[SparseKVPageTransfer],

721 + *,

722 + stream: torch.Stream | None = None,

723 ) -> None:

713 - if not transfers or not self.is_host_writer:

724 + if not transfers:

725 return

715 - stream = self.dma_stream

726 + if stream is None:

727 + if not self.is_host_writer:

728 + return

729 + stream = self.dma_stream

730 assert stream is not None

731 completion_event = torch.Event()

732 completion_event.record(stream)

@@ -720,6 +734,30 @@ class HiSparseConnectorWorker:

734 self._pending_transfer_events.append((completion_event, transfer_ids))

735 self._enqueued_transfer_ids.extend(transfer_ids)

736

737 + def _restore_pages(self, transfers: list[SparseKVPageTransfer]) -> None:

738 + """Restore imported tails on every rank before its forward or graph replay."""

739 + if not transfers:

740 + return

741 + for layer_index, cache in enumerate(self.cache_handles):

742 + source_index = cache.runtime.resident_source_index

743 + source = self.host_caches[layer_index]

744 + destination = self.resident_caches[layer_index]

745 + for transfer in transfers:

746 + host_start = transfer.host_block_id * self.kernel_block_size

747 + resident_block = transfer.resident_block_ids[source_index]

748 + # Each import pays one H2D page copy and enqueue per layer/rank

749 + # instead of tail prefill on D. Restore the prompt rows before

750 + # decode appends to this page; unused rows do not extend the

751 + # logical sequence.

752 + # cudaHostRegister pins this pool outside PyTorch's allocator,

753 + # so the sync checker cannot recognize it via is_pinned().

754 + with gpu_sync_allowed():

755 + destination[resident_block].copy_(

756 + source[host_start : host_start + self.kernel_block_size],

757 + non_blocking=True,

758 + )

759 + self._record_transfer_completion(transfers, stream=current_stream())

760 +

761 def _submit_transfers(self, transfers: list[SparseKVPageTransfer]) -> None:

762 if self.cache_handles[0].runtime.eager_host_mirror:

763 self._record_transfer_completion(transfers)

@@ -733,15 +771,12 @@ class HiSparseConnectorWorker:

771 descriptor_count = len(transfers) * num_layers

772 descriptors = self._acquire_dma_descriptors(descriptor_count)

773 destination_rows = np.fromiter(

736 - (

737 - transfer.destination_block_id * self.kernel_block_size

738 - for transfer in transfers

739 - ),

774 + (transfer.host_block_id * self.kernel_block_size for transfer in transfers),

775 dtype=np.int64,

776 count=len(transfers),

777 )

778 source_blocks_by_transfer = np.asarray(

744 - [transfer.source_block_ids for transfer in transfers], dtype=np.int64

779 + [transfer.resident_block_ids for transfer in transfers], dtype=np.int64

780 )

781 if source_blocks_by_transfer.ndim != 2:

782 raise RuntimeError(

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py

+2 -1

@@ -3278,7 +3278,8 @@ class NixlBaseConnectorWorker:

3278 """Pair an uncached decode suffix with the same prefill regions."""

3279 assert len(decode_block_ids) == len(prefill_block_ids)

3280 if not any(decode_block_ids):

3281 - return [], prefill_block_ids

3281 + empty_regions: list[list[int]] = [[] for _ in decode_block_ids]

3282 + return empty_regions, empty_regions.copy()

3283

3284 if num_computed_blocks is not None:

3285 assert num_remote_blocks is not None

vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py

+8 -2

@@ -168,8 +168,14 @@ class NixlPullConnectorWorker(NixlBaseConnectorWorker):

168 local_block_ids = meta.local_physical_block_ids

169 remote_region_groups = self.dst_region_group_ids[engine_id]

170 local_region_groups = self.region_group_ids or remote_region_groups

171 - groups_differ = local_region_groups != remote_region_groups

172 - if groups_differ:

171 + if not local_block_ids:

172 + # Region expansion cannot index empty groups. Pass empty specs to

173 + # _read_blocks so its existing cache-hit notification path runs.

174 + read_specs = [

175 + ReadSpec(remote_rank=rank, local_block_ids=[], remote_block_ids=[])

176 + for rank in plan.all_source_ranks

177 + ]

178 + elif local_region_groups != remote_region_groups:

179 if not self.use_mla or self._has_mamba:

180 raise NotImplementedError(

181 "Different NIXL cache-group layouts are only supported for "

vllm/v1/core/single_type_kv_cache_manager.py

+12 -8

@@ -2498,12 +2498,14 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):

2498 del num_tokens_main_model

2499 assert not new_computed_blocks

2500 if total_computed_tokens > num_local_computed_tokens:

2501 - if total_computed_tokens % self.block_size != 0:

2502 - raise ValueError(

2503 - "A host-only HiSparse import must end on a cache-block boundary."

2504 - )

2505 - imported_pages = total_computed_tokens // self.block_size

2506 - return max(cdiv(num_tokens, self.block_size) - imported_pages, 0)

2501 + # Keep the last imported page writable, including the last-token

2502 + # replay on a full prompt hit. Only earlier pages stay host-only.

2503 + imported_pages = cdiv(total_computed_tokens, self.block_size)

2504 + existing = len(self.req_to_blocks.get(request_id, ()))

2505 + needs_tail = self.get_resident_page(request_id, imported_pages - 1) is None

2506 + return int(needs_tail) + max(

2507 + cdiv(num_tokens, self.block_size) - max(imported_pages, existing), 0

2508 + )

2509 existing = len(self.req_to_blocks.get(request_id, ()))

2510 host_pages = cdiv(num_local_computed_tokens, self.block_size)

2511 required = cdiv(num_tokens, self.block_size)

@@ -2518,13 +2520,15 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):

2520 num_local_computed_tokens: int,

2521 num_external_computed_tokens: int,

2522 ) -> None:

2521 - """Represent imported host history with null resident pages."""

2523 + """Reserve a writable final page; earlier imported pages stay on host."""

2524 assert num_external_computed_tokens > 0

2525 num_tokens = num_local_computed_tokens + num_external_computed_tokens

2526 blocks = self.req_to_blocks[request_id]

2525 - tail_page = cdiv(num_tokens, self.block_size) - 1

2527 + tail_page = (num_tokens - 1) // self.block_size

2528 if len(blocks) <= tail_page:

2529 blocks.extend([self._null_block] * (tail_page + 1 - len(blocks)))

2530 + if blocks[tail_page].is_null:

2531 + blocks[tail_page] = self.block_pool.get_new_blocks(1)[0]

2532

2533 def add_local_computed_blocks(

2534 self,

vllm/v1/hisparse/coordinator.py

+47 -21

@@ -73,6 +73,7 @@ class _PendingSpill:

73 request_state: _HiSparseRequestState

74 host_block: KVCacheBlock

75 resident_blocks: tuple[KVCacheBlock, ...]

76 + restore: bool = False

77 expected_worker_completions: int = 0

78 worker_completions: int = 0

79 enqueue_applied: bool = False

@@ -271,25 +272,29 @@ class HiSparseCoordinator:

272 num_computed_tokens // self.host_manager.block_size, len(host_blocks)

273 )

274 for host_idx in range(state.copies_recorded_blocks, num_host_blocks):

274 - host_block = host_blocks[host_idx]

275 - if host_block.is_null or host_block.block_hash is None:

276 - continue

277 - blocks: list[KVCacheBlock] = []

278 - for manager in self.resident_managers:

279 - block = manager.get_resident_page(request_id, host_idx)

280 - if block is None:

281 - break

282 - blocks.append(block)

283 - if len(blocks) != len(self.resident_managers):

275 + if host_idx in state.pending_pages:

276 continue

285 - self._drop_copy(host_block.block_hash)

286 - self.copies[host_block.block_hash] = tuple(blocks)

287 - for block in blocks:

288 - self._copy_by_block[block.block_id] = host_block.block_hash

277 + self._record_copy(request_id, host_idx, host_blocks[host_idx])

278 state.copies_recorded_blocks = max(

279 state.copies_recorded_blocks, num_host_blocks

280 )

281

282 + def _record_copy(

283 + self, request_id: str, page_idx: int, host_block: KVCacheBlock

284 + ) -> None:

285 + if host_block.is_null or host_block.block_hash is None:

286 + return

287 + blocks: list[KVCacheBlock] = []

288 + for manager in self.resident_managers:

289 + block = manager.get_resident_page(request_id, page_idx)

290 + if block is None:

291 + return

292 + blocks.append(block)

293 + self._drop_copy(host_block.block_hash)

294 + self.copies[host_block.block_hash] = tuple(blocks)

295 + for block in blocks:

296 + self._copy_by_block[block.block_id] = host_block.block_hash

297 +

298 def _drop_copy(self, host_hash: BlockHashWithGroupId) -> None:

299 blocks = self.copies.pop(host_hash, None)

300 if blocks is None:

@@ -406,9 +411,14 @@ class HiSparseCoordinator:

411 assert self.host_manager is not None

412 host_block_size = self.host_manager.block_size

413 num_pages = num_computed_tokens // host_block_size

414 + importing_pages = cdiv(

415 + self._pending_imports.get(request_id, 0), host_block_size

416 + )

417 state = self._get_request_state(request_id)

418 budget = max(self.max_spill_pages - len(self.spills_to_send), 0)

419 for page_idx in range(num_pages):

420 + if page_idx < importing_pages:

421 + continue

422 if budget == 0:

423 break

424 if page_idx in state.pending_pages:

@@ -420,7 +430,7 @@ class HiSparseCoordinator:

430 for manager in self.resident_managers

431 ):

432 continue

423 - if self._plan_spill(request_id, page_idx, after_forward=True):

433 + if self._plan_page_transfer(request_id, page_idx, after_forward=True):

434 budget -= 1

435

436 def publish_when_ready(

@@ -491,17 +501,27 @@ class HiSparseCoordinator:

501 state = self._get_request_state(request_id)

502 state.valid_pages = set(range(num_pages))

503 state.ready_prefix_pages = num_pages

504 + if num_computed_tokens:

505 + self._plan_page_transfer(

506 + request_id,

507 + (num_computed_tokens - 1) // block_size,

508 + after_forward=False,

509 + restore=True,

510 + )

511 self._publish_host_blocks_if_ready(request_id)

512

496 - def _plan_spill(

513 + def _plan_page_transfer(

514 self,

515 request_id: str,

516 page_idx: int,

517 *,

518 after_forward: bool,

519 + restore: bool = False,

520 ) -> bool:

521 assert self.host_manager is not None

522 state = self._get_request_state(request_id)

523 + if page_idx in state.pending_pages:

524 + return False

525 host_blocks = self.host_manager.req_to_blocks.get(request_id)

526 if host_blocks is None or page_idx >= len(host_blocks):

527 return False

@@ -523,9 +543,10 @@ class HiSparseCoordinator:

543 self.next_spill_id += 1

544 plan = SparseKVPageTransfer(

545 transfer_id=spill_id,

526 - destination_block_id=host_block.block_id,

527 - source_block_ids=tuple(block.block_id for block in blocks),

546 + host_block_id=host_block.block_id,

547 + resident_block_ids=tuple(block.block_id for block in blocks),

548 after_forward=after_forward,

549 + restore=restore,

550 )

551 self.pending_spills[spill_id] = _PendingSpill(

552 transfer_id=spill_id,

@@ -533,6 +554,7 @@ class HiSparseCoordinator:

554 request_state=state,

555 host_block=host_block,

556 resident_blocks=tuple(blocks),

557 + restore=restore,

558 )

559 state.pending_pages[page_idx] = spill_id

560 self.spills_to_send.append(plan)

@@ -666,7 +688,7 @@ class HiSparseCoordinator:

688 self._complete_host_writes()

689

690 def _apply_enqueued_spills(self) -> None:

669 - """Drop the DMA pins once the worker has enqueued a spill."""

691 + """Drop the GPU pins once every worker has completed the transfer."""

692 for pending in self.pending_spills.values():

693 if (

694 pending.enqueue_applied

@@ -692,8 +714,12 @@ class HiSparseCoordinator:

714 request_id, page_idx = pending.page

715 pending.request_state.pending_pages.pop(page_idx, None)

716 if self.request_states.get(request_id) is pending.request_state:

695 - pending.request_state.valid_pages.add(page_idx)

696 - self._page_became_clean(request_id, pending.request_state, page_idx)

717 + if not pending.restore:

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

719 + if page_idx in pending.request_state.valid_pages:

720 + self._page_became_clean(request_id, pending.request_state, page_idx)

721 + if pending.restore:

722 + self._record_copy(request_id, page_idx, pending.host_block)

723 completed_request_ids.add(request_id)

724 assert self.host_manager is not None

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

vllm/v1/hisparse/types.py

+3 -2

@@ -9,9 +9,10 @@ class SparseKVPageTransfer:

9 """Copy one logical KV page between cache-manager and worker-owned tiers."""

11 transfer_id: int

12 - destination_block_id: int

13 - source_block_ids: tuple[int, ...]

12 + host_block_id: int

13 + resident_block_ids: tuple[int, ...]

14 after_forward: bool

15 + restore: bool = False

18 @dataclass(frozen=True)