@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3eb6cec22a

[Bugfix][HiSparse] Adopt GPU prefix copies after the hit's allocation (#59282)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Matthew Bonanni committed Sep 30, 2026 at 15:43 UTC 3eb6cec22ad9bb098393021b956feaa97081fc85

2 files changed +61 -3

tests/v1/core/test_prefix_caching.py

+46

@@ -1139,6 +1139,52 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():

1139 )

1140

1141

1142 +@pytest.mark.parametrize("free_blocks,adopted_pages", [(7, []), (10, [0, 1])])

1143 +def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies(

1144 + free_blocks, adopted_pages

1145 +):

1146 + """A host prefix hit must not starve its own allocation of free blocks.

1147 +

1148 + Admission counts free GPU copies as free, but adopting a copy pins it. When

1149 + adoption ran before the hit's allocation, it took blocks the allocation was

1150 + promised and the pool ran dry (``Cannot get N free blocks``). The allocation

1151 + now goes first, evicting a prefix's copies tail first; the copies it leaves

1152 + are adopted.

1153 + """

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

1155 + tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

1156 + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

1157 + assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None

1158 + _publish_hisparse_pages(manager)

1159 + copy_ids = [block.block_id for block in manager.get_blocks("original").blocks[2]]

1160 + manager.free(original)

1161 + pool = manager.block_pool

1162 + pool.get_new_blocks(pool.get_num_free_blocks() - free_blocks)

1163 +

1164 + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)

1165 + computed, num_computed, _ = manager.get_computed_blocks(resumed)

1166 + assert num_computed == 3 * HISPARSE_BLOCK_SIZE

1167 + assert manager.allocate_slots(

1168 + resumed,

1169 + num_new_tokens=len(tokens) - num_computed,

1170 + num_new_computed_tokens=num_computed,

1171 + new_computed_blocks=computed,

1172 + )

1173 +

1174 + resumed_blocks = manager.get_blocks("resumed").blocks

1175 + assert [block.block_id for block in resumed_blocks[2][:3]] == [

1176 + copy_ids[page] if page in adopted_pages else pool.null_block.block_id

1177 + for page in range(3)

1178 + ]

1179 + gpu_ids = [

1180 + block.block_id

1181 + for group in resumed_blocks[1:]

1182 + for block in group

1183 + if not block.is_null

1184 + ]

1185 + assert len(gpu_ids) == len(set(gpu_ids))

1186 +

1187 +

1188 def test_hisparse_host_backed_request_accepts_local_prefix_hit():

1189 """Local group-completion hits need no external host import allocation."""

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

vllm/v1/hisparse/coordinator.py

+15 -3

@@ -66,6 +66,8 @@ class _HiSparseRequestState:

66 copies_recorded_blocks: int = 0

67 pinned_clean: set[int] = field(default_factory=set)

68 unpinned_pages: set[int] = field(default_factory=set)

69 + # Prefix pages whose GPU copies are adopted after the admitting allocation.

70 + pages_to_adopt: int = 0

73 @dataclass

@@ -195,7 +197,9 @@ class HiSparseCoordinator:

197 state = self._get_request_state(request_id)

198 state.valid_pages.update(range(num_host_pages))

199 state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages)

198 - self._adopt_copies(request_id, state, host_blocks[:num_host_pages])

200 + # Adopting pins copies out of the free pool, so it must wait until the

201 + # allocation admission counted on has run (see update_residency).

202 + state.pages_to_adopt = max(state.pages_to_adopt, num_host_pages)

203 state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages)

204

205 def take_host_block_copies(self) -> tuple[KVCacheBlockCopy, ...]:

@@ -345,7 +349,8 @@ class HiSparseCoordinator:

349 return True

350

351 def _unpin_clean_pages(self, request_id: str, state: _HiSparseRequestState) -> None:

348 - for page_idx in sorted(state.pinned_clean):

352 + # Tail first, so a prefix loses its tail pages before its head.

353 + for page_idx in sorted(state.pinned_clean, reverse=True):

354 self._unpin_page(request_id, state, page_idx)

355

356 def _page_became_clean(

@@ -391,6 +396,11 @@ class HiSparseCoordinator:

396 if not self.resident_managers:

397 return

398 state = self._get_request_state(request_id)

399 + if state.pages_to_adopt:

400 + assert self.host_manager is not None

401 + host_blocks = self.host_manager.req_to_blocks.get(request_id, ())

402 + self._adopt_copies(request_id, state, host_blocks[: state.pages_to_adopt])

403 + state.pages_to_adopt = 0

404 if not self._can_read_from_host(request_id):

405 assert self.gpu_pool is not None

406 if self.gpu_pool.get_num_free_blocks() >= self.transition_watermark:

@@ -791,7 +801,9 @@ class HiSparseCoordinator:

801 blocks = manager.req_to_blocks.get(request_id)

802 if not blocks:

803 continue

794 - for page_idx, block in enumerate(blocks):

804 + # Tail first, so a prefix loses its tail pages before its head.

805 + for page_idx in reversed(range(len(blocks))):

806 + block = blocks[page_idx]

807 if block.is_null:

808 continue

809 if page_idx in state.unpinned_pages: