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

[Perf][HiSparse] Avoid repeated prefix scans and residency updates (#57930)

Signed-off-by: S1ro1 <matej.sirovatka@gmail.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Matej Sirovatka committed Oct 1, 2026 at 14:48 UTC ca65eb67d995f5401c575ec90cba523757127ec7

5 files changed +236 -83

tests/v1/core/test_prefix_caching.py

+134 -46

@@ -168,6 +168,7 @@ def make_hisparse_kv_cache_config(

168 host_num_blocks: int,

169 *,

170 transfer_device_cache: bool = False,

171 + num_resident_groups: int = 1,

172 ) -> KVCacheConfig:

173 source_spec = FullAttentionSpec(

174 block_size=HISPARSE_BLOCK_SIZE,

@@ -210,6 +211,13 @@ def make_hisparse_kv_cache_config(

211 enable_kv_transfer=False,

212 )

213 )

214 + for group_idx in range(1, num_resident_groups):

215 + groups.extend(

216 + replace(

217 + group, layer_names=[f"{name}_{group_idx}" for name in group.layer_names]

218 + )

219 + for group in groups[2:4]

220 + )

221 return KVCacheConfig(

222 num_blocks=num_blocks,

223 hisparse_host_num_blocks=host_num_blocks,

@@ -225,6 +233,7 @@ def make_hisparse_kv_cache_manager(

233 max_model_len: int = 128,

234 max_in_flight_tokens: int | None = None,

235 enable_caching: bool = False,

236 + num_prefill_lookahead: int = 0,

237 **config_kwargs,

238 ) -> KVCacheManager:

239 config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs)

@@ -234,9 +243,53 @@ def make_hisparse_kv_cache_manager(

243 max_in_flight_tokens=max_in_flight_tokens,

244 enable_caching=enable_caching,

245 hash_block_size=HISPARSE_BLOCK_SIZE,

246 + num_prefill_lookahead=num_prefill_lookahead,

247 )

248

249

250 +def _allocate_scheduled(

251 + manager: KVCacheManager,

252 + request: Request,

253 + num_new_tokens: int,

254 + *,

255 + num_new_computed_tokens: int = 0,

256 + num_external_computed_tokens: int = 0,

257 + **kwargs,

258 +) -> KVCacheBlocks | None:

259 + """``allocate_slots`` plus the per-step residency work the connector runs."""

260 + blocks = manager.allocate_slots(

261 + request,

262 + num_new_tokens,

263 + num_new_computed_tokens=num_new_computed_tokens,

264 + num_external_computed_tokens=num_external_computed_tokens,

265 + **kwargs,

266 + )

267 + if blocks is None or kwargs.get("delay_cache_blocks"):

268 + return blocks

269 + num_computed_tokens = min(

270 + request.num_computed_tokens

271 + + num_new_computed_tokens

272 + + num_external_computed_tokens,

273 + manager.max_model_len,

274 + )

275 + num_tokens = min(num_computed_tokens + num_new_tokens, request.num_tokens)

276 + get_hisparse_coordinator(manager).advance_scheduled(

277 + [(request.request_id, num_tokens)]

278 + )

279 + return blocks

280 +

281 +

282 +def test_hisparse_does_not_write_back_reprefillable_tokens():

283 + """Multi-module MTP re-prefills the last tokens, so they are not final."""

284 + manager = make_hisparse_kv_cache_manager(32, 16, num_prefill_lookahead=3)

285 + request = make_request(

286 + "request", list(range(2 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256

287 + )

288 + assert _allocate_scheduled(manager, request, 2 * HISPARSE_BLOCK_SIZE)

289 + transfers = get_hisparse_coordinator(manager).build_offload_command().page_transfers

290 + assert len(transfers) == 1

291 +

292 +

293 def test_hisparse_builds_dma_row_mirrors_across_pages():

294 manager = make_hisparse_kv_cache_manager(32, 16)

295 request = make_request(

@@ -245,7 +298,7 @@ def test_hisparse_builds_dma_row_mirrors_across_pages():

298 HISPARSE_BLOCK_SIZE,

299 sha256,

300 )

248 - assert manager.allocate_slots(request, num_new_tokens=32) is not None

301 + assert _allocate_scheduled(manager, request, num_new_tokens=32) is not None

302 coordinator = get_hisparse_coordinator(manager)

303 resident_blocks = coordinator.resident_managers[0].req_to_blocks[request.request_id]

304 host_blocks = coordinator.host_manager.req_to_blocks[request.request_id]

@@ -308,14 +361,14 @@ def test_hisparse_reports_when_context_is_fully_resident():

361 manager = make_hisparse_kv_cache_manager(32, 16)

362 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

363 request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)

311 - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None

364 + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None

365 coordinator = get_hisparse_coordinator(manager)

366 scheduled = ((request.request_id, len(tokens) - 1, 1),)

367 assert coordinator.all_context_pages_resident(scheduled)

368

369 for hot_manager in coordinator.hot_managers:

370 hot_manager.require_hot(request.request_id)

318 - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None

371 + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None

372 _publish_hisparse_pages(manager)

373 pool = manager.block_pool

374 pool.get_new_blocks(pool.get_num_free_blocks())

@@ -333,7 +386,9 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():

386 )

387 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

388 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

336 - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None

389 + assert (

390 + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None

391 + )

392 spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers

393 spill_counts = {spill.transfer_id: 1 for spill in spills}

394 get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)

@@ -348,7 +403,8 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():

403 assert num_local == 2 * HISPARSE_BLOCK_SIZE

404 assert [len(group_blocks) for group_blocks in blocks.blocks] == [3, 2, 0, 0]

405

351 - allocated = manager.allocate_slots(

406 + allocated = _allocate_scheduled(

407 + manager,

408 resumed,

409 num_new_tokens=1,

410 num_new_computed_tokens=num_local,

@@ -379,7 +435,9 @@ def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix():

435 )

436 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

437 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

382 - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None

438 + assert (

439 + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None

440 + )

441 spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers

442 spill_counts = {spill.transfer_id: 1 for spill in spills}

443 get_hisparse_coordinator(manager).update_spills(spill_counts, spill_counts)

@@ -410,7 +468,7 @@ def test_connector_completes_partial_prefix_without_importing_missing_host_kv(

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

469 tokens = list(range(64))

470 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

413 - assert manager.allocate_slots(original, num_new_tokens=64) is not None

471 + assert _allocate_scheduled(manager, original, num_new_tokens=64) is not None

472 _publish_hisparse_pages(manager)

473 host, indexer, _, _ = manager.get_blocks(original.request_id).blocks

474 missing_indexer = indexer[1].block_id

@@ -457,7 +515,8 @@ def test_connector_completes_partial_prefix_without_importing_missing_host_kv(

515 assert asynchronous == (external > 0)

516 assert [len(group) for group in blocks.blocks] == [3, 1, 0, 0]

517 assert (

460 - manager.allocate_slots(

518 + _allocate_scheduled(

519 + manager,

520 resumed,

521 int(external == 0),

522 delay_cache_blocks=True,

@@ -486,7 +545,8 @@ def test_connector_completes_partial_prefix_without_importing_missing_host_kv(

545 def allocate_external_prefix(

546 manager: KVCacheManager, request: Request, num_tokens: int

547 ) -> KVCacheBlocks | None:

489 - return manager.allocate_slots(

548 + return _allocate_scheduled(

549 + manager,

550 request,

551 num_new_tokens=0,

552 num_external_computed_tokens=num_tokens,

@@ -538,7 +598,7 @@ def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens):

598 scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler)

599 Scheduler._update_waiting_for_remote_kv(scheduler, request)

600 assert request.num_tokens - request.num_computed_tokens == 1

541 - assert manager.allocate_slots(request, num_new_tokens=1) is not None

601 + assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None

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

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

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

@@ -583,15 +643,15 @@ def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_cach

643 32, 32, max_model_len=256, enable_caching=enable_caching

644 )

645 first = make_request("first", list(range(160)), HISPARSE_BLOCK_SIZE, sha256)

586 - assert manager.allocate_slots(first, num_new_tokens=128) is not None

646 + assert _allocate_scheduled(manager, first, num_new_tokens=128) is not None

647 _publish_hisparse_pages(manager)

648 first.num_computed_tokens = 128

649 pool = manager.block_pool

650 held = pool.get_new_blocks(pool.get_num_free_blocks() - 3)

651

592 - assert manager.allocate_slots(first, num_new_tokens=1) is not None

652 + assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None

653 second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)

594 - manager.allocate_slots(second, num_new_tokens=16)

654 + _allocate_scheduled(manager, second, num_new_tokens=16)

655 coordinator = get_hisparse_coordinator(manager)

656 scheduled = (("first", 128, 1),)

657 assert coordinator.all_context_pages_resident(scheduled)

@@ -600,7 +660,7 @@ def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_cach

660 # Once another allocation finishes, A can acquire its hot buffer and

661 # safely release its clean pages for other requests to reuse.

662 pool.free_blocks(held)

603 - assert manager.allocate_slots(first, num_new_tokens=1) is not None

663 + assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None

664 assert all(m.has_hot("first") for m in coordinator.hot_managers)

665 pool.get_new_blocks(pool.get_num_free_blocks())

666 assert not coordinator.all_context_pages_resident(scheduled)

@@ -611,21 +671,24 @@ def test_hisparse_full_pool_keeps_pages_pinned_until_preemption():

671 """Without room for a hot buffer, admission must defer instead of losing KV."""

672 manager = make_hisparse_kv_cache_manager(18, 18, max_model_len=160)

673 first = make_request("first", list(range(128)), HISPARSE_BLOCK_SIZE, sha256)

614 - assert manager.allocate_slots(first, num_new_tokens=128) is not None

674 + assert _allocate_scheduled(manager, first, num_new_tokens=128) is not None

675 _publish_hisparse_pages(manager)

676 assert manager.block_pool.get_num_free_blocks() == 1

677 first.num_computed_tokens = 128

618 - assert manager.allocate_slots(first, num_new_tokens=1) is None

678 + assert _allocate_scheduled(manager, first, num_new_tokens=1) is None

679 assert get_hisparse_coordinator(manager).all_context_pages_resident(

680 (("first", 127, 1),)

681 )

682

683 manager.free(first)

684 second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256)

625 - assert manager.allocate_slots(second, num_new_tokens=16) is not None

685 + assert _allocate_scheduled(manager, second, num_new_tokens=16) is not None

686

687

628 -def test_hisparse_materializes_prefix_without_allocating_hot_blocks():

688 +@pytest.mark.parametrize("num_resident_groups", [1, 3])

689 +def test_hisparse_materializes_prefix_without_allocating_hot_blocks(

690 + num_resident_groups,

691 +):

692 """A host prefix becomes visible only when every page is durable.

693

694 Completing a later page first must not expose a prefix with a hole.

@@ -634,10 +697,11 @@ def test_hisparse_materializes_prefix_without_allocating_hot_blocks():

697 32,

698 16,

699 enable_caching=True,

700 + num_resident_groups=num_resident_groups,

701 )

702 tokens = list(range(2 * HISPARSE_BLOCK_SIZE))

703 request = make_request("resident", tokens, HISPARSE_BLOCK_SIZE, sha256)

640 - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None

704 + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None

705 blocks = manager.get_block_ids(request.request_id)

706 assert len(blocks[2]) == 2

707 assert blocks[3] == []

@@ -675,19 +739,21 @@ def test_hisparse_materializes_prefix_without_allocating_hot_blocks():

739 assert blocks[3] == []

740

741

678 -def test_hisparse_materialization_respects_per_step_spill_budget():

742 +@pytest.mark.parametrize("num_resident_groups", [1, 3])

743 +def test_hisparse_materialization_respects_per_step_spill_budget(num_resident_groups):

744 """Prefix publication must not bypass the configured spill batch limit."""

745 manager = make_hisparse_kv_cache_manager(

746 32,

747 16,

748 enable_caching=True,

749 + num_resident_groups=num_resident_groups,

750 )

751 coordinator = get_hisparse_coordinator(manager)

752 coordinator.max_spill_pages = 1

753 tokens = list(range(2 * HISPARSE_BLOCK_SIZE))

754 request = make_request("bounded", tokens, HISPARSE_BLOCK_SIZE, sha256)

755

690 - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None

756 + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None

757 first = coordinator.build_offload_command().page_transfers

758 assert len(first) == 1

759

@@ -705,7 +771,9 @@ def test_hisparse_deferred_free_retains_host_blocks_until_fence():

771 host_free = host_pool.get_num_free_blocks()

772 device_free = manager.block_pool.get_num_free_blocks()

773 request = make_request("free", list(range(16)), HISPARSE_BLOCK_SIZE, sha256)

708 - assert manager.allocate_slots(request, 16, delay_cache_blocks=True) is not None

774 + assert (

775 + _allocate_scheduled(manager, request, 16, delay_cache_blocks=True) is not None

776 + )

777 blocks = manager.pop_blocks_for_free(request)

778 assert host_pool.get_num_free_blocks() < host_free

779 assert blocks and any(host_pool.blocks[b.block_id] is b for b in blocks)

@@ -903,7 +971,7 @@ def test_hisparse_host_exhaustion_defers_allocation(enable_caching):

971 """Pages are only allocated with host backing, so every one can be reclaimed."""

972 manager = make_hisparse_kv_cache_manager(32, 5, enable_caching=enable_caching)

973 donor = make_request("donor", [99] * 16, HISPARSE_BLOCK_SIZE, sha256)

906 - assert manager.allocate_slots(donor, 16) is not None

974 + assert _allocate_scheduled(manager, donor, 16) is not None

975 request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)

976 free_gpu = manager.block_pool.get_num_free_blocks()

977 assert manager.allocate_slots(request, 64) is None

@@ -916,7 +984,7 @@ def test_hisparse_host_exhaustion_defers_allocation(enable_caching):

984 assert manager.allocate_slots(request, 64) is None

985 _publish_hisparse_pages(manager)

986 assert not coordinator.has_pending_reclamation()

919 - assert manager.allocate_slots(request, 64) is not None

987 + assert _allocate_scheduled(manager, request, 64) is not None

988 host = coordinator.host_manager

989 assert host is not None

990 assert not any(b.is_null for b in host.req_to_blocks[request.request_id])

@@ -974,7 +1042,7 @@ def test_hisparse_recomputation_does_not_complete_failed_host_import():

1042 # allocation reaches the original import boundary before the worker runs.

1043 request.num_computed_tokens = 16

1044 manager.cache_blocks(request, 16)

977 - assert manager.allocate_slots(request, 16) is not None

1045 + assert _allocate_scheduled(manager, request, 16) is not None

1046 state = get_hisparse_coordinator(manager).request_states.get(request.request_id)

1047 assert state is None or 1 not in state.valid_pages

1048

@@ -995,12 +1063,12 @@ def test_hisparse_resident_request_can_grow_without_hot_capacity():

1063 """Running resident history must not be mistaken for a new host-prefix hit."""

1064 manager = make_hisparse_kv_cache_manager(8, 16)

1065 request = make_request("resident", list(range(48)), HISPARSE_BLOCK_SIZE, sha256)

998 - assert manager.allocate_slots(request, num_new_tokens=32) is not None

1066 + assert _allocate_scheduled(manager, request, num_new_tokens=32) is not None

1067 request.num_computed_tokens = 32

1068

1069 # Three free blocks fit the next indexer/resident pages, but not a hot region.

1070 assert manager.block_pool.get_num_free_blocks() == 3

1003 - assert manager.allocate_slots(request, num_new_tokens=1) is not None

1071 + assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None

1072 assert manager.get_block_ids(request.request_id)[3] == []

1073

1074

@@ -1050,7 +1118,7 @@ def test_hisparse_events_report_host_and_device_placement():

1118 enable_kv_cache_events=True,

1119 )

1120 request = make_request("request", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)

1053 - assert manager.allocate_slots(request, 64) is not None

1121 + assert _allocate_scheduled(manager, request, 64) is not None

1122 _publish_hisparse_pages(manager)

1123

1124 stored = [

@@ -1113,7 +1181,9 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():

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

1182 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

1183 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

1116 - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None

1184 + assert (

1185 + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None

1186 + )

1187 _publish_hisparse_pages(manager)

1188 original_resident_ids = [

1189 block.block_id for block in manager.get_blocks("original").blocks[2]

@@ -1123,7 +1193,8 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():

1193 resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)

1194 computed, num_computed, _ = manager.get_computed_blocks(resumed)

1195 assert num_computed == 3 * HISPARSE_BLOCK_SIZE

1126 - assert manager.allocate_slots(

1196 + assert _allocate_scheduled(

1197 + manager,

1198 resumed,

1199 num_new_tokens=len(tokens) - num_computed,

1200 num_new_computed_tokens=num_computed,

@@ -1154,7 +1225,9 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies(

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

1226 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

1227 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

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

1228 + assert (

1229 + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None

1230 + )

1231 _publish_hisparse_pages(manager)

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

1233 manager.free(original)

@@ -1164,7 +1237,8 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies(

1237 resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)

1238 computed, num_computed, _ = manager.get_computed_blocks(resumed)

1239 assert num_computed == 3 * HISPARSE_BLOCK_SIZE

1167 - assert manager.allocate_slots(

1240 + assert _allocate_scheduled(

1241 + manager,

1242 resumed,

1243 num_new_tokens=len(tokens) - num_computed,

1244 num_new_computed_tokens=num_computed,

@@ -1190,14 +1264,17 @@ def test_hisparse_host_backed_request_accepts_local_prefix_hit():

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

1265 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

1266 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

1193 - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None

1267 + assert (

1268 + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None

1269 + )

1270 _publish_hisparse_pages(manager)

1271 manager.free(original)

1272

1273 resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)

1274 computed, num_computed, _ = manager.get_computed_blocks(resumed)

1275 assert num_computed == 3 * HISPARSE_BLOCK_SIZE

1200 - assert manager.allocate_slots(

1276 + assert _allocate_scheduled(

1277 + manager,

1278 resumed,

1279 num_new_tokens=len(tokens) - num_computed,

1280 num_new_computed_tokens=num_computed,

@@ -1211,7 +1288,9 @@ def test_hisparse_finished_request_leaves_free_copies():

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

1289 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

1290 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

1214 - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None

1291 + assert (

1292 + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None

1293 + )

1294 _publish_hisparse_pages(manager)

1295 manager.free(original)

1296

@@ -1241,7 +1320,9 @@ def test_hisparse_cleanup_publishes_only_finalized_terminal_pages(

1320 coordinator = get_hisparse_coordinator(manager)

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

1322 request.spec_token_ids = list(range(33, 49))

1244 - assert manager.allocate_slots(request, 49, num_lookahead_tokens=16) is not None

1323 + assert (

1324 + _allocate_scheduled(manager, request, 49, num_lookahead_tokens=16) is not None

1325 + )

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

1327 command = coordinator.build_offload_command()

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

@@ -1280,7 +1361,10 @@ def test_hisparse_prefill_reads_host_once_it_fills_admission_window(

1361 coordinator = get_hisparse_coordinator(manager)

1362 tokens = list(range(2 * window_pages * HISPARSE_BLOCK_SIZE))

1363 request = make_request("prefill", tokens, HISPARSE_BLOCK_SIZE, sha256)

1283 - assert manager.allocate_slots(request, num_pages * HISPARSE_BLOCK_SIZE) is not None

1364 + assert (

1365 + _allocate_scheduled(manager, request, num_pages * HISPARSE_BLOCK_SIZE)

1366 + is not None

1367 + )

1368 pool = coordinator.gpu_pool

1369 assert pool is not None

1370 assert pool.get_num_free_blocks() >= coordinator.transition_watermark

@@ -1303,10 +1387,10 @@ def test_hisparse_preempted_prefill_resumes_from_durable_prefix():

1387 chunk = 2 * HISPARSE_BLOCK_SIZE

1388 tokens = list(range(3 * chunk))

1389 request = make_request("long", tokens, HISPARSE_BLOCK_SIZE, sha256)

1306 - assert manager.allocate_slots(request, chunk) is not None

1390 + assert _allocate_scheduled(manager, request, chunk) is not None

1391 request.num_computed_tokens = chunk

1392 first_chunk_writes = coordinator.build_offload_command().page_transfers

1309 - assert manager.allocate_slots(request, chunk) is not None

1393 + assert _allocate_scheduled(manager, request, chunk) is not None

1394 request.num_computed_tokens = 2 * chunk

1395 acks = {transfer.transfer_id: 1 for transfer in first_chunk_writes}

1396 coordinator.update_spills(acks, acks)

@@ -1323,7 +1407,7 @@ def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity()

1407 host_pool = coordinator.get_host_block_pool()

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

1409 request = make_request("reused", tokens, HISPARSE_BLOCK_SIZE, sha256)

1326 - assert manager.allocate_slots(request, len(tokens)) is not None

1410 + assert _allocate_scheduled(manager, request, len(tokens)) is not None

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

1412 command = coordinator.build_offload_command()

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

@@ -1336,7 +1420,7 @@ def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity()

1420 replacement = make_request(

1421 "reused", list(range(100, 117)), HISPARSE_BLOCK_SIZE, sha256

1422 )

1339 - assert manager.allocate_slots(replacement, 17) is not None

1423 + assert _allocate_scheduled(manager, replacement, 17) is not None

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

1425 pressure = host_pool.get_new_blocks(host_pool.get_num_free_blocks())

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

@@ -1363,7 +1447,7 @@ def test_hisparse_reset_prefix_cache_drops_copies():

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

1448 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

1449 request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)

1366 - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None

1450 + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None

1451 _publish_hisparse_pages(manager)

1452 manager.free(request)

1453

@@ -1378,7 +1462,9 @@ def test_hisparse_reused_copy_is_not_adopted():

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

1463 tokens = list(range(4 * HISPARSE_BLOCK_SIZE))

1464 original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256)

1381 - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None

1465 + assert (

1466 + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None

1467 + )

1468 _publish_hisparse_pages(manager)

1469 indexer_blocks = list(manager.get_blocks("original").blocks[1])

1470 manager.free(original)

@@ -1391,7 +1477,8 @@ def test_hisparse_reused_copy_is_not_adopted():

1477 resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256)

1478 computed, num_computed, _ = manager.get_computed_blocks(resumed)

1479 assert num_computed == 3 * HISPARSE_BLOCK_SIZE

1394 - assert manager.allocate_slots(

1480 + assert _allocate_scheduled(

1481 + manager,

1482 resumed,

1483 num_new_tokens=len(tokens) - num_computed,

1484 num_new_computed_tokens=num_computed,

@@ -5291,7 +5378,8 @@ def test_can_fit_full_sequence_hisparse_caps_resident_pages():

5378 )

5379

5380 assert (

5294 - manager.allocate_slots(

5381 + _allocate_scheduled(

5382 + manager,

5383 request,

5384 num_new_tokens=HISPARSE_BLOCK_SIZE,

5385 full_sequence_must_fit=True,

@@ -6151,7 +6239,7 @@ def test_device_eviction_keeps_host_prefix():

6239 """Device eviction must not target a same-numbered host block."""

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

6241 request = make_request("producer", list(range(64)), HISPARSE_BLOCK_SIZE, sha256)

6154 - assert manager.allocate_slots(request, 64) is not None

6242 + assert _allocate_scheduled(manager, request, 64) is not None

6243 _publish_hisparse_pages(manager)

6244 host, indexer, _, _ = manager.get_blocks(request.request_id).blocks

6245 device_block = indexer[0]

tests/v1/kv_connector/unit/test_hisparse_connector.py

+53

@@ -114,3 +114,56 @@ def test_full_graph_step_prepares_host_mirror_outside_model():

114

115 worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None)

116 runtime.invalidate_written_slots.assert_called_once()

117 +

118 +

119 +def test_scheduled_prefix_hit_publishes_adopted_copies():

120 + """Copies adopted after scheduling must reach the worker's block table."""

121 + from tests.v1.core.test_prefix_caching import (

122 + HISPARSE_BLOCK_SIZE,

123 + _allocate_scheduled,

124 + _publish_hisparse_pages,

125 + make_hisparse_kv_cache_manager,

126 + make_request,

127 + sha256,

128 + )

129 + from vllm.v1.core.kv_cache_utils import init_none_hash

130 +

131 + init_none_hash(sha256)

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

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

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

135 + assert _allocate_scheduled(manager, original, num_new_tokens=len(tokens))

136 + _publish_hisparse_pages(manager)

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

138 + manager.free(original)

139 +

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

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

142 + num_new_tokens = len(tokens) - num_computed

143 + assert manager.allocate_slots(

144 + resumed,

145 + num_new_tokens=num_new_tokens,

146 + num_new_computed_tokens=num_computed,

147 + new_computed_blocks=computed,

148 + )

149 + scheduler = HiSparseConnectorScheduler(async_speculative=False)

150 + scheduler.bind_coordinator(get_hisparse_coordinator(manager))

151 + scheduler.requests[resumed.request_id] = resumed

152 + scheduler_output = SimpleNamespace(

153 + scheduled_new_reqs=[

154 + SimpleNamespace(

155 + req_id=resumed.request_id,

156 + num_computed_tokens=num_computed,

157 + block_ids=manager.get_block_ids(resumed.request_id),

158 + )

159 + ],

160 + scheduled_cached_reqs=SimpleNamespace(

161 + req_ids=[], num_computed_tokens=[], new_block_ids=[]

162 + ),

163 + num_scheduled_tokens={resumed.request_id: num_new_tokens},

164 + )

165 +

166 + scheduler.build_connector_meta(scheduler_output)

167 +

168 + resident_ids = scheduler_output.block_table_updates[resumed.request_id][2]

169 + assert resident_ids[:3] == copy_ids[:3]

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

+23 -15

@@ -106,21 +106,6 @@ class HiSparseConnectorScheduler:

106 # HiSparse rebinds worker state every step, so the load must start

107 # before the forward rather than being deferred to post-forward.

108 scheduler_output.has_sync_kv_loads = True

109 - scheduler_output.block_table_updates = (

110 - self.coordinator.take_block_table_updates() or None

111 - )

112 - command = self.coordinator.build_offload_command()

113 - host_block_copies = self.coordinator.take_host_block_copies()

114 - source_group_id = self.coordinator.host_group_id

115 - assert source_group_id is not None

116 - source_block_ids = [

117 - block_id

118 - for request in scheduler_output.scheduled_new_reqs

119 - for block_id in request.block_ids[source_group_id]

120 - ]

121 - for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids:

122 - if new_block_ids is not None:

123 - source_block_ids.extend(new_block_ids[source_group_id])

109 num_computed_tokens = {

110 request.req_id: request.num_computed_tokens

111 for request in scheduler_output.scheduled_new_reqs

@@ -141,6 +126,29 @@ class HiSparseConnectorScheduler:

126 scheduler_output.num_scheduled_tokens.items()

127 )

128 )

129 + # Runs first: it plans this step's transfers and table updates.

130 + self.coordinator.advance_scheduled(

131 + (

132 + request_id,

133 + min(start + count, self.requests[request_id].num_tokens),

134 + )

135 + for request_id, start, count in scheduled_requests

136 + )

137 + scheduler_output.block_table_updates = (

138 + self.coordinator.take_block_table_updates() or None

139 + )

140 + command = self.coordinator.build_offload_command()

141 + host_block_copies = self.coordinator.take_host_block_copies()

142 + source_group_id = self.coordinator.host_group_id

143 + assert source_group_id is not None

144 + source_block_ids = [

145 + block_id

146 + for request in scheduler_output.scheduled_new_reqs

147 + for block_id in request.block_ids[source_group_id]

148 + ]

149 + for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids:

150 + if new_block_ids is not None:

151 + source_block_ids.extend(new_block_ids[source_group_id])

152 row_mirrors = {}

153 for request_id, scheduled_start, scheduled_count in scheduled_requests:

154 mirror_start = scheduled_start

vllm/v1/core/single_type_kv_cache_manager.py

-18

@@ -2366,12 +2366,6 @@ class _HiSparseAuxiliaryManager(SingleTypeKVCacheManager):

2366

2367 coordinator: "HiSparseCoordinator | None" = None

2368

2369 - def __init__(self, kv_cache_spec: KVCacheSpec, **kwargs) -> None:

2370 - # Never prefix-cached, but the per-step ``cache_blocks`` hook is where

2371 - # residency work runs, so stay opted in regardless of prefix caching.

2372 - kwargs["enable_caching"] = True

2373 - super().__init__(kv_cache_spec, **kwargs)

2374 -

2369 def cache_blocks(

2370 self,

2371 request: Request,

@@ -2551,18 +2545,6 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):

2545 assert self.coordinator is not None

2546 self.coordinator.commit_computed_blocks(request_id, num_host_pages)

2547

2554 - def cache_blocks(

2555 - self,

2556 - request: Request,

2557 - num_tokens: int,

2558 - retention_interval: int | None = None,

2559 - *,

2560 - replay_boundaries: Sequence[int],

2561 - ) -> None:

2562 - assert self.coordinator is not None

2563 - self.coordinator.plan_prefix_materialization(request.request_id, num_tokens)

2564 - self.coordinator.update_residency(request.request_id)

2565 -

2548 def allocate_new_blocks(

2549 self, request_id: str, num_tokens: int, num_tokens_main_model: int

2550 ) -> list[KVCacheBlock]:

vllm/v1/hisparse/coordinator.py

+26 -4

@@ -99,9 +99,12 @@ class HiSparseCoordinator:

99 kv_cache_config: KVCacheConfig,

100 managers: tuple[SingleTypeKVCacheManager, ...],

101 max_model_len: int,

102 + *,

103 + num_reprefillable_tokens: int,

104 ) -> None:

105 self.managers = managers

106 self.max_model_len = max_model_len

107 + self.num_reprefillable_tokens = num_reprefillable_tokens

108 groups = kv_cache_config.kv_cache_groups

109

110 resident_managers: list[HiSparseResidentManager] = []

@@ -264,6 +267,7 @@ class HiSparseCoordinator:

267 if manager.adopt_resident_page(request_id, host_idx, block):

268 manager.block_pool.touch([block])

269 state.pinned_clean.add(host_idx)

270 + self.block_table_updates.add(request_id)

271

272 def _record_copies(self, request_id: str, num_computed_tokens: int) -> None:

273 """Index the GPU copies of just-published host blocks for later hits."""

@@ -431,6 +435,22 @@ class HiSparseCoordinator:

435 # Host publication and spills

436 # ------------------------------------------------------------------

437

438 + def advance_scheduled(self, requests: Iterable[tuple[str, int]]) -> None:

439 + """Run the per-step residency work for each scheduled request.

440 +

441 + ``requests`` pairs a request id with the number of tokens computed

442 + once the step completes. As in ``KVCacheCoordinator.cache_blocks``, the

443 + last ``num_reprefillable_tokens`` are not final yet and are not written

444 + back.

445 + """

446 + if not self.resident_managers:

447 + return

448 + for request_id, num_tokens in requests:

449 + self.plan_prefix_materialization(

450 + request_id, max(0, num_tokens - self.num_reprefillable_tokens)

451 + )

452 + self.update_residency(request_id)

453 +

454 def plan_prefix_materialization(

455 self, request_id: str, num_computed_tokens: int

456 ) -> None:

@@ -445,9 +465,8 @@ class HiSparseCoordinator:

465 )

466 state = self._get_request_state(request_id)

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

448 - for page_idx in range(num_pages):

449 - if page_idx < importing_pages:

450 - continue

468 + first_page = max(state.ready_prefix_pages, importing_pages)

469 + for page_idx in range(first_page, num_pages):

470 if budget == 0:

471 break

472 if page_idx in state.pending_pages:

@@ -852,7 +871,10 @@ def get_hisparse_coordinator(

871 assert isinstance(coordinator, HiSparseCoordinator)

872 return coordinator

873 coordinator = HiSparseCoordinator(

855 - kv_cache_manager.kv_cache_config, managers, kv_cache_manager.max_model_len

874 + kv_cache_manager.kv_cache_config,

875 + managers,

876 + kv_cache_manager.max_model_len,

877 + num_reprefillable_tokens=kv_cache_manager.coordinator.num_reprefillable_tokens,

878 )

879 if coordinator.host_manager is None:

880 raise ValueError("No HiSparse cache group is configured.")