diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 027bb67b18..3972fe5230 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -168,6 +168,7 @@ def make_hisparse_kv_cache_config( host_num_blocks: int, *, transfer_device_cache: bool = False, + num_resident_groups: int = 1, ) -> KVCacheConfig: source_spec = FullAttentionSpec( block_size=HISPARSE_BLOCK_SIZE, @@ -210,6 +211,13 @@ def make_hisparse_kv_cache_config( enable_kv_transfer=False, ) ) + for group_idx in range(1, num_resident_groups): + groups.extend( + replace( + group, layer_names=[f"{name}_{group_idx}" for name in group.layer_names] + ) + for group in groups[2:4] + ) return KVCacheConfig( num_blocks=num_blocks, hisparse_host_num_blocks=host_num_blocks, @@ -225,6 +233,7 @@ def make_hisparse_kv_cache_manager( max_model_len: int = 128, max_in_flight_tokens: int | None = None, enable_caching: bool = False, + num_prefill_lookahead: int = 0, **config_kwargs, ) -> KVCacheManager: config = make_hisparse_kv_cache_config(num_blocks, host_num_blocks, **config_kwargs) @@ -234,9 +243,53 @@ def make_hisparse_kv_cache_manager( max_in_flight_tokens=max_in_flight_tokens, enable_caching=enable_caching, hash_block_size=HISPARSE_BLOCK_SIZE, + num_prefill_lookahead=num_prefill_lookahead, ) +def _allocate_scheduled( + manager: KVCacheManager, + request: Request, + num_new_tokens: int, + *, + num_new_computed_tokens: int = 0, + num_external_computed_tokens: int = 0, + **kwargs, +) -> KVCacheBlocks | None: + """``allocate_slots`` plus the per-step residency work the connector runs.""" + blocks = manager.allocate_slots( + request, + num_new_tokens, + num_new_computed_tokens=num_new_computed_tokens, + num_external_computed_tokens=num_external_computed_tokens, + **kwargs, + ) + if blocks is None or kwargs.get("delay_cache_blocks"): + return blocks + num_computed_tokens = min( + request.num_computed_tokens + + num_new_computed_tokens + + num_external_computed_tokens, + manager.max_model_len, + ) + num_tokens = min(num_computed_tokens + num_new_tokens, request.num_tokens) + get_hisparse_coordinator(manager).advance_scheduled( + [(request.request_id, num_tokens)] + ) + return blocks + + +def test_hisparse_does_not_write_back_reprefillable_tokens(): + """Multi-module MTP re-prefills the last tokens, so they are not final.""" + manager = make_hisparse_kv_cache_manager(32, 16, num_prefill_lookahead=3) + request = make_request( + "request", list(range(2 * HISPARSE_BLOCK_SIZE)), HISPARSE_BLOCK_SIZE, sha256 + ) + assert _allocate_scheduled(manager, request, 2 * HISPARSE_BLOCK_SIZE) + transfers = get_hisparse_coordinator(manager).build_offload_command().page_transfers + assert len(transfers) == 1 + + def test_hisparse_builds_dma_row_mirrors_across_pages(): manager = make_hisparse_kv_cache_manager(32, 16) request = make_request( @@ -245,7 +298,7 @@ def test_hisparse_builds_dma_row_mirrors_across_pages(): HISPARSE_BLOCK_SIZE, sha256, ) - assert manager.allocate_slots(request, num_new_tokens=32) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=32) is not None coordinator = get_hisparse_coordinator(manager) resident_blocks = coordinator.resident_managers[0].req_to_blocks[request.request_id] host_blocks = coordinator.host_manager.req_to_blocks[request.request_id] @@ -308,14 +361,14 @@ def test_hisparse_reports_when_context_is_fully_resident(): manager = make_hisparse_kv_cache_manager(32, 16) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None coordinator = get_hisparse_coordinator(manager) scheduled = ((request.request_id, len(tokens) - 1, 1),) assert coordinator.all_context_pages_resident(scheduled) for hot_manager in coordinator.hot_managers: hot_manager.require_hot(request.request_id) - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None _publish_hisparse_pages(manager) pool = manager.block_pool pool.get_new_blocks(pool.get_num_free_blocks()) @@ -333,7 +386,9 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload(): ) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + assert ( + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None + ) spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers spill_counts = {spill.transfer_id: 1 for spill in spills} 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(): assert num_local == 2 * HISPARSE_BLOCK_SIZE assert [len(group_blocks) for group_blocks in blocks.blocks] == [3, 2, 0, 0] - allocated = manager.allocate_slots( + allocated = _allocate_scheduled( + manager, resumed, num_new_tokens=1, num_new_computed_tokens=num_local, @@ -379,7 +435,9 @@ def test_hisparse_indexer_offload_is_capped_by_missing_host_prefix(): ) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + assert ( + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None + ) spills = get_hisparse_coordinator(manager).build_offload_command().page_transfers spill_counts = {spill.transfer_id: 1 for spill in spills} 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( manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(64)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=64) is not None + assert _allocate_scheduled(manager, original, num_new_tokens=64) is not None _publish_hisparse_pages(manager) host, indexer, _, _ = manager.get_blocks(original.request_id).blocks missing_indexer = indexer[1].block_id @@ -457,7 +515,8 @@ def test_connector_completes_partial_prefix_without_importing_missing_host_kv( assert asynchronous == (external > 0) assert [len(group) for group in blocks.blocks] == [3, 1, 0, 0] assert ( - manager.allocate_slots( + _allocate_scheduled( + manager, resumed, int(external == 0), delay_cache_blocks=True, @@ -486,7 +545,8 @@ def test_connector_completes_partial_prefix_without_importing_missing_host_kv( def allocate_external_prefix( manager: KVCacheManager, request: Request, num_tokens: int ) -> KVCacheBlocks | None: - return manager.allocate_slots( + return _allocate_scheduled( + manager, request, num_new_tokens=0, num_external_computed_tokens=num_tokens, @@ -538,7 +598,7 @@ def test_nixl_hisparse_full_block_import_keeps_a_writable_tail(num_tokens): scheduler._mark_prefix_replay = MethodType(Scheduler._mark_prefix_replay, scheduler) Scheduler._update_waiting_for_remote_kv(scheduler, request) assert request.num_tokens - request.num_computed_tokens == 1 - assert manager.allocate_slots(request, num_new_tokens=1) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None assert coordinator.build_row_mirrors([(request.request_id, num_tokens - 1, 1)]) # A pending restore must not expose uninitialized GPU copies to prefix hits. 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 32, 32, max_model_len=256, enable_caching=enable_caching ) first = make_request("first", list(range(160)), HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(first, num_new_tokens=128) is not None + assert _allocate_scheduled(manager, first, num_new_tokens=128) is not None _publish_hisparse_pages(manager) first.num_computed_tokens = 128 pool = manager.block_pool held = pool.get_new_blocks(pool.get_num_free_blocks() - 3) - assert manager.allocate_slots(first, num_new_tokens=1) is not None + assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256) - manager.allocate_slots(second, num_new_tokens=16) + _allocate_scheduled(manager, second, num_new_tokens=16) coordinator = get_hisparse_coordinator(manager) scheduled = (("first", 128, 1),) assert coordinator.all_context_pages_resident(scheduled) @@ -600,7 +660,7 @@ def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_cach # Once another allocation finishes, A can acquire its hot buffer and # safely release its clean pages for other requests to reuse. pool.free_blocks(held) - assert manager.allocate_slots(first, num_new_tokens=1) is not None + assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None assert all(m.has_hot("first") for m in coordinator.hot_managers) pool.get_new_blocks(pool.get_num_free_blocks()) assert not coordinator.all_context_pages_resident(scheduled) @@ -611,21 +671,24 @@ def test_hisparse_full_pool_keeps_pages_pinned_until_preemption(): """Without room for a hot buffer, admission must defer instead of losing KV.""" manager = make_hisparse_kv_cache_manager(18, 18, max_model_len=160) first = make_request("first", list(range(128)), HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(first, num_new_tokens=128) is not None + assert _allocate_scheduled(manager, first, num_new_tokens=128) is not None _publish_hisparse_pages(manager) assert manager.block_pool.get_num_free_blocks() == 1 first.num_computed_tokens = 128 - assert manager.allocate_slots(first, num_new_tokens=1) is None + assert _allocate_scheduled(manager, first, num_new_tokens=1) is None assert get_hisparse_coordinator(manager).all_context_pages_resident( (("first", 127, 1),) ) manager.free(first) second = make_request("second", [1000] * 16, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(second, num_new_tokens=16) is not None + assert _allocate_scheduled(manager, second, num_new_tokens=16) is not None -def test_hisparse_materializes_prefix_without_allocating_hot_blocks(): +@pytest.mark.parametrize("num_resident_groups", [1, 3]) +def test_hisparse_materializes_prefix_without_allocating_hot_blocks( + num_resident_groups, +): """A host prefix becomes visible only when every page is durable. 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(): 32, 16, enable_caching=True, + num_resident_groups=num_resident_groups, ) tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) request = make_request("resident", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None blocks = manager.get_block_ids(request.request_id) assert len(blocks[2]) == 2 assert blocks[3] == [] @@ -675,19 +739,21 @@ def test_hisparse_materializes_prefix_without_allocating_hot_blocks(): assert blocks[3] == [] -def test_hisparse_materialization_respects_per_step_spill_budget(): +@pytest.mark.parametrize("num_resident_groups", [1, 3]) +def test_hisparse_materialization_respects_per_step_spill_budget(num_resident_groups): """Prefix publication must not bypass the configured spill batch limit.""" manager = make_hisparse_kv_cache_manager( 32, 16, enable_caching=True, + num_resident_groups=num_resident_groups, ) coordinator = get_hisparse_coordinator(manager) coordinator.max_spill_pages = 1 tokens = list(range(2 * HISPARSE_BLOCK_SIZE)) request = make_request("bounded", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None first = coordinator.build_offload_command().page_transfers assert len(first) == 1 @@ -705,7 +771,9 @@ def test_hisparse_deferred_free_retains_host_blocks_until_fence(): host_free = host_pool.get_num_free_blocks() device_free = manager.block_pool.get_num_free_blocks() request = make_request("free", list(range(16)), HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, 16, delay_cache_blocks=True) is not None + assert ( + _allocate_scheduled(manager, request, 16, delay_cache_blocks=True) is not None + ) blocks = manager.pop_blocks_for_free(request) assert host_pool.get_num_free_blocks() < host_free 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): """Pages are only allocated with host backing, so every one can be reclaimed.""" manager = make_hisparse_kv_cache_manager(32, 5, enable_caching=enable_caching) donor = make_request("donor", [99] * 16, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(donor, 16) is not None + assert _allocate_scheduled(manager, donor, 16) is not None request = make_request("resident", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) free_gpu = manager.block_pool.get_num_free_blocks() assert manager.allocate_slots(request, 64) is None @@ -916,7 +984,7 @@ def test_hisparse_host_exhaustion_defers_allocation(enable_caching): assert manager.allocate_slots(request, 64) is None _publish_hisparse_pages(manager) assert not coordinator.has_pending_reclamation() - assert manager.allocate_slots(request, 64) is not None + assert _allocate_scheduled(manager, request, 64) is not None host = coordinator.host_manager assert host is not None 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(): # allocation reaches the original import boundary before the worker runs. request.num_computed_tokens = 16 manager.cache_blocks(request, 16) - assert manager.allocate_slots(request, 16) is not None + assert _allocate_scheduled(manager, request, 16) is not None state = get_hisparse_coordinator(manager).request_states.get(request.request_id) assert state is None or 1 not in state.valid_pages @@ -995,12 +1063,12 @@ def test_hisparse_resident_request_can_grow_without_hot_capacity(): """Running resident history must not be mistaken for a new host-prefix hit.""" manager = make_hisparse_kv_cache_manager(8, 16) request = make_request("resident", list(range(48)), HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, num_new_tokens=32) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=32) is not None request.num_computed_tokens = 32 # Three free blocks fit the next indexer/resident pages, but not a hot region. assert manager.block_pool.get_num_free_blocks() == 3 - assert manager.allocate_slots(request, num_new_tokens=1) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=1) is not None assert manager.get_block_ids(request.request_id)[3] == [] @@ -1050,7 +1118,7 @@ def test_hisparse_events_report_host_and_device_placement(): enable_kv_cache_events=True, ) request = make_request("request", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, 64) is not None + assert _allocate_scheduled(manager, request, 64) is not None _publish_hisparse_pages(manager) stored = [ @@ -1113,7 +1181,9 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages(): manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + assert ( + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None + ) _publish_hisparse_pages(manager) original_resident_ids = [ 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(): resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) computed, num_computed, _ = manager.get_computed_blocks(resumed) assert num_computed == 3 * HISPARSE_BLOCK_SIZE - assert manager.allocate_slots( + assert _allocate_scheduled( + manager, resumed, num_new_tokens=len(tokens) - num_computed, num_new_computed_tokens=num_computed, @@ -1154,7 +1225,9 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + assert ( + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None + ) _publish_hisparse_pages(manager) copy_ids = [block.block_id for block in manager.get_blocks("original").blocks[2]] manager.free(original) @@ -1164,7 +1237,8 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies( resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) computed, num_computed, _ = manager.get_computed_blocks(resumed) assert num_computed == 3 * HISPARSE_BLOCK_SIZE - assert manager.allocate_slots( + assert _allocate_scheduled( + manager, resumed, num_new_tokens=len(tokens) - num_computed, num_new_computed_tokens=num_computed, @@ -1190,14 +1264,17 @@ def test_hisparse_host_backed_request_accepts_local_prefix_hit(): manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + assert ( + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None + ) _publish_hisparse_pages(manager) manager.free(original) resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) computed, num_computed, _ = manager.get_computed_blocks(resumed) assert num_computed == 3 * HISPARSE_BLOCK_SIZE - assert manager.allocate_slots( + assert _allocate_scheduled( + manager, resumed, num_new_tokens=len(tokens) - num_computed, num_new_computed_tokens=num_computed, @@ -1211,7 +1288,9 @@ def test_hisparse_finished_request_leaves_free_copies(): manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + assert ( + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None + ) _publish_hisparse_pages(manager) manager.free(original) @@ -1241,7 +1320,9 @@ def test_hisparse_cleanup_publishes_only_finalized_terminal_pages( 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 + assert ( + _allocate_scheduled(manager, 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} @@ -1280,7 +1361,10 @@ def test_hisparse_prefill_reads_host_once_it_fills_admission_window( coordinator = get_hisparse_coordinator(manager) tokens = list(range(2 * window_pages * HISPARSE_BLOCK_SIZE)) request = make_request("prefill", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, num_pages * HISPARSE_BLOCK_SIZE) is not None + assert ( + _allocate_scheduled(manager, request, num_pages * HISPARSE_BLOCK_SIZE) + is not None + ) pool = coordinator.gpu_pool assert pool is not None assert pool.get_num_free_blocks() >= coordinator.transition_watermark @@ -1303,10 +1387,10 @@ def test_hisparse_preempted_prefill_resumes_from_durable_prefix(): chunk = 2 * HISPARSE_BLOCK_SIZE tokens = list(range(3 * chunk)) request = make_request("long", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, chunk) is not None + assert _allocate_scheduled(manager, request, chunk) is not None request.num_computed_tokens = chunk first_chunk_writes = coordinator.build_offload_command().page_transfers - assert manager.allocate_slots(request, chunk) is not None + assert _allocate_scheduled(manager, request, chunk) is not None request.num_computed_tokens = 2 * chunk acks = {transfer.transfer_id: 1 for transfer in first_chunk_writes} coordinator.update_spills(acks, acks) @@ -1323,7 +1407,7 @@ def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity() 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 + assert _allocate_scheduled(manager, 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} @@ -1336,7 +1420,7 @@ def test_hisparse_terminal_prefix_waits_for_all_workers_and_preserves_identity() replacement = make_request( "reused", list(range(100, 117)), HISPARSE_BLOCK_SIZE, sha256 ) - assert manager.allocate_slots(replacement, 17) is not None + assert _allocate_scheduled(manager, 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]} & { @@ -1363,7 +1447,7 @@ def test_hisparse_reset_prefix_cache_drops_copies(): manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, num_new_tokens=len(tokens)) is not None + assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None _publish_hisparse_pages(manager) manager.free(request) @@ -1378,7 +1462,9 @@ def test_hisparse_reused_copy_is_not_adopted(): manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(original, num_new_tokens=len(tokens)) is not None + assert ( + _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) is not None + ) _publish_hisparse_pages(manager) indexer_blocks = list(manager.get_blocks("original").blocks[1]) manager.free(original) @@ -1391,7 +1477,8 @@ def test_hisparse_reused_copy_is_not_adopted(): resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) computed, num_computed, _ = manager.get_computed_blocks(resumed) assert num_computed == 3 * HISPARSE_BLOCK_SIZE - assert manager.allocate_slots( + assert _allocate_scheduled( + manager, resumed, num_new_tokens=len(tokens) - num_computed, num_new_computed_tokens=num_computed, @@ -5291,7 +5378,8 @@ def test_can_fit_full_sequence_hisparse_caps_resident_pages(): ) assert ( - manager.allocate_slots( + _allocate_scheduled( + manager, request, num_new_tokens=HISPARSE_BLOCK_SIZE, full_sequence_must_fit=True, @@ -6151,7 +6239,7 @@ def test_device_eviction_keeps_host_prefix(): """Device eviction must not target a same-numbered host block.""" manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) request = make_request("producer", list(range(64)), HISPARSE_BLOCK_SIZE, sha256) - assert manager.allocate_slots(request, 64) is not None + assert _allocate_scheduled(manager, request, 64) is not None _publish_hisparse_pages(manager) host, indexer, _, _ = manager.get_blocks(request.request_id).blocks device_block = indexer[0] diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index 6c4e79efbe..384460fa2c 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -114,3 +114,56 @@ def test_full_graph_step_prepares_host_mirror_outside_model(): worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None) runtime.invalidate_written_slots.assert_called_once() + + +def test_scheduled_prefix_hit_publishes_adopted_copies(): + """Copies adopted after scheduling must reach the worker's block table.""" + from tests.v1.core.test_prefix_caching import ( + HISPARSE_BLOCK_SIZE, + _allocate_scheduled, + _publish_hisparse_pages, + make_hisparse_kv_cache_manager, + make_request, + sha256, + ) + from vllm.v1.core.kv_cache_utils import init_none_hash + + init_none_hash(sha256) + manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) + tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) + original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) + assert _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) + _publish_hisparse_pages(manager) + copy_ids = [block.block_id for block in manager.get_blocks("original").blocks[2]] + manager.free(original) + + resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) + computed, num_computed, _ = manager.get_computed_blocks(resumed) + num_new_tokens = len(tokens) - num_computed + assert manager.allocate_slots( + resumed, + num_new_tokens=num_new_tokens, + num_new_computed_tokens=num_computed, + new_computed_blocks=computed, + ) + scheduler = HiSparseConnectorScheduler(async_speculative=False) + scheduler.bind_coordinator(get_hisparse_coordinator(manager)) + scheduler.requests[resumed.request_id] = resumed + scheduler_output = SimpleNamespace( + scheduled_new_reqs=[ + SimpleNamespace( + req_id=resumed.request_id, + num_computed_tokens=num_computed, + block_ids=manager.get_block_ids(resumed.request_id), + ) + ], + scheduled_cached_reqs=SimpleNamespace( + req_ids=[], num_computed_tokens=[], new_block_ids=[] + ), + num_scheduled_tokens={resumed.request_id: num_new_tokens}, + ) + + scheduler.build_connector_meta(scheduler_output) + + resident_ids = scheduler_output.block_table_updates[resumed.request_id][2] + assert resident_ids[:3] == copy_ids[:3] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py index 5acb7e3bbd..dabf9095db 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py @@ -106,21 +106,6 @@ class HiSparseConnectorScheduler: # HiSparse rebinds worker state every step, so the load must start # before the forward rather than being deferred to post-forward. scheduler_output.has_sync_kv_loads = True - scheduler_output.block_table_updates = ( - self.coordinator.take_block_table_updates() or None - ) - command = self.coordinator.build_offload_command() - host_block_copies = self.coordinator.take_host_block_copies() - source_group_id = self.coordinator.host_group_id - assert source_group_id is not None - source_block_ids = [ - block_id - for request in scheduler_output.scheduled_new_reqs - for block_id in request.block_ids[source_group_id] - ] - for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids: - if new_block_ids is not None: - source_block_ids.extend(new_block_ids[source_group_id]) num_computed_tokens = { request.req_id: request.num_computed_tokens for request in scheduler_output.scheduled_new_reqs @@ -141,6 +126,29 @@ class HiSparseConnectorScheduler: scheduler_output.num_scheduled_tokens.items() ) ) + # Runs first: it plans this step's transfers and table updates. + self.coordinator.advance_scheduled( + ( + request_id, + min(start + count, self.requests[request_id].num_tokens), + ) + for request_id, start, count in scheduled_requests + ) + scheduler_output.block_table_updates = ( + self.coordinator.take_block_table_updates() or None + ) + command = self.coordinator.build_offload_command() + host_block_copies = self.coordinator.take_host_block_copies() + source_group_id = self.coordinator.host_group_id + assert source_group_id is not None + source_block_ids = [ + block_id + for request in scheduler_output.scheduled_new_reqs + for block_id in request.block_ids[source_group_id] + ] + for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids: + if new_block_ids is not None: + source_block_ids.extend(new_block_ids[source_group_id]) row_mirrors = {} for request_id, scheduled_start, scheduled_count in scheduled_requests: mirror_start = scheduled_start diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index a097c9e924..7041600f8b 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -2366,12 +2366,6 @@ class _HiSparseAuxiliaryManager(SingleTypeKVCacheManager): coordinator: "HiSparseCoordinator | None" = None - def __init__(self, kv_cache_spec: KVCacheSpec, **kwargs) -> None: - # Never prefix-cached, but the per-step ``cache_blocks`` hook is where - # residency work runs, so stay opted in regardless of prefix caching. - kwargs["enable_caching"] = True - super().__init__(kv_cache_spec, **kwargs) - def cache_blocks( self, request: Request, @@ -2551,18 +2545,6 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager): assert self.coordinator is not None self.coordinator.commit_computed_blocks(request_id, num_host_pages) - def cache_blocks( - self, - request: Request, - num_tokens: int, - retention_interval: int | None = None, - *, - replay_boundaries: Sequence[int], - ) -> None: - assert self.coordinator is not None - self.coordinator.plan_prefix_materialization(request.request_id, num_tokens) - self.coordinator.update_residency(request.request_id) - def allocate_new_blocks( self, request_id: str, num_tokens: int, num_tokens_main_model: int ) -> list[KVCacheBlock]: diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 1a783b506b..bc9bff1a8b 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/vllm/v1/hisparse/coordinator.py @@ -99,9 +99,12 @@ class HiSparseCoordinator: kv_cache_config: KVCacheConfig, managers: tuple[SingleTypeKVCacheManager, ...], max_model_len: int, + *, + num_reprefillable_tokens: int, ) -> None: self.managers = managers self.max_model_len = max_model_len + self.num_reprefillable_tokens = num_reprefillable_tokens groups = kv_cache_config.kv_cache_groups resident_managers: list[HiSparseResidentManager] = [] @@ -264,6 +267,7 @@ class HiSparseCoordinator: if manager.adopt_resident_page(request_id, host_idx, block): manager.block_pool.touch([block]) state.pinned_clean.add(host_idx) + self.block_table_updates.add(request_id) def _record_copies(self, request_id: str, num_computed_tokens: int) -> None: """Index the GPU copies of just-published host blocks for later hits.""" @@ -431,6 +435,22 @@ class HiSparseCoordinator: # Host publication and spills # ------------------------------------------------------------------ + def advance_scheduled(self, requests: Iterable[tuple[str, int]]) -> None: + """Run the per-step residency work for each scheduled request. + + ``requests`` pairs a request id with the number of tokens computed + once the step completes. As in ``KVCacheCoordinator.cache_blocks``, the + last ``num_reprefillable_tokens`` are not final yet and are not written + back. + """ + if not self.resident_managers: + return + for request_id, num_tokens in requests: + self.plan_prefix_materialization( + request_id, max(0, num_tokens - self.num_reprefillable_tokens) + ) + self.update_residency(request_id) + def plan_prefix_materialization( self, request_id: str, num_computed_tokens: int ) -> None: @@ -445,9 +465,8 @@ class HiSparseCoordinator: ) state = self._get_request_state(request_id) budget = max(self.max_spill_pages - len(self.spills_to_send), 0) - for page_idx in range(num_pages): - if page_idx < importing_pages: - continue + first_page = max(state.ready_prefix_pages, importing_pages) + for page_idx in range(first_page, num_pages): if budget == 0: break if page_idx in state.pending_pages: @@ -852,7 +871,10 @@ def get_hisparse_coordinator( assert isinstance(coordinator, HiSparseCoordinator) return coordinator coordinator = HiSparseCoordinator( - kv_cache_manager.kv_cache_config, managers, kv_cache_manager.max_model_len + kv_cache_manager.kv_cache_config, + managers, + kv_cache_manager.max_model_len, + num_reprefillable_tokens=kv_cache_manager.coordinator.num_reprefillable_tokens, ) if coordinator.host_manager is None: raise ValueError("No HiSparse cache group is configured.")