[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>
2026-10-01 21:48:18 +00:00
parent 7566d83bd3
commit ca65eb67d9
5 changed files with 236 additions and 83 deletions
@@ -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]
@@ -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]
@@ -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
@@ -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]:
@@ -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.")