[HiSparse] Own GPU residency instead of editing core block-table rows (#60647)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

csrc/libtorch_stable/hisparse_kernels.cu

@@ -194,7 +194,7 @@ constexpr int kHashLoadInverse = 4;
 // the page is not resident) and returns its host row (-1 when not host-backed).
 __device__ __forceinline__ int32_t translate_topk_entry(
     const int32_t token_index, const int32_t request_row,
-    const int32_t* __restrict__ source_block_table,
+    const int32_t state_row, const int32_t* __restrict__ source_block_table,
     const int32_t* __restrict__ resident_block_table, const int64_t host_rows,
     const int64_t source_bt_stride, const int32_t source_num_reqs,
     const int32_t source_num_blocks, const int32_t source_block_size,
@@ -221,10 +221,10 @@ __device__ __forceinline__ int32_t translate_topk_entry(
     if (resident_block_table != nullptr) {
       const int32_t resident_block =
           token_index >= 0 ? token_index / resident_block_size : -1;
-      if (request_row >= 0 && request_row < resident_num_reqs &&
+      if (state_row >= 0 && state_row < resident_num_reqs &&
           resident_block >= 0 && resident_block < resident_num_blocks) {
         const int32_t physical_block =
-            resident_block_table[static_cast<int64_t>(request_row) *
+            resident_block_table[static_cast<int64_t>(state_row) *
                                      resident_bt_stride +
                                  resident_block];
         if (physical_block != resident_null_block && physical_block >= 0) {
@@ -424,11 +424,11 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
     for (int i = tid; i < top_k; i += blockDim.x) {
       int32_t resident_row;
       const int32_t g = translate_topk_entry(
-          row_topk[i], request_row, source_block_table, resident_block_table,
-          host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
-          source_block_size, resident_bt_stride, resident_num_reqs,
-          resident_num_blocks, resident_block_size, resident_null_block,
-          resident_row);
+          row_topk[i], request_row, state_row, source_block_table,
+          resident_block_table, host_rows, source_bt_stride, source_num_reqs,
+          source_num_blocks, source_block_size, resident_bt_stride,
+          resident_num_reqs, resident_num_blocks, resident_block_size,
+          resident_null_block, resident_row);
       if (resolved_global_indices != nullptr) {
         resolved_global_indices[static_cast<int64_t>(row) * top_k + i] = g;
       }
@@ -602,8 +602,8 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
       g = translate_topk_entry(
           global_indices[static_cast<int64_t>(row) * input_row_stride +
                          position % top_k],
-          request_row, source_block_table, resident_block_table, host_rows,
-          source_bt_stride, source_num_reqs, source_num_blocks,
+          request_row, state_row, source_block_table, resident_block_table,
+          host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
           source_block_size, resident_bt_stride, resident_num_reqs,
           resident_num_blocks, resident_block_size, resident_null_block,
           resident_row);
@@ -708,8 +708,8 @@ __global__ __launch_bounds__(1024) void hisparse_resolve_residency_kernel(
       const int32_t g = translate_topk_entry(
           global_indices[static_cast<int64_t>(row) * input_row_stride +
                          position % top_k],
-          request_row, source_block_table, resident_block_table, host_rows,
-          source_bt_stride, source_num_reqs, source_num_blocks,
+          request_row, state_row, source_block_table, resident_block_table,
+          host_rows, source_bt_stride, source_num_reqs, source_num_blocks,
           source_block_size, resident_bt_stride, resident_num_reqs,
           resident_num_blocks, resident_block_size, resident_null_block,
           resident_row);

tests/v1/attention/test_sparse_mla_backends.py

@@ -3241,7 +3241,8 @@ def test_hisparse_prefill_staging_plan_resolves_resident_sources():
         [[11, 12, 0, 13, 0, 0], [21, 0, 22, 23, 0, 0]], dtype=torch.int32
     )
 
-    plan.ensure_gpu_sources(resident_table, resident_block_size)
+    state_indices = torch.arange(2, dtype=torch.int32)
+    plan.ensure_gpu_sources(resident_table, state_indices, resident_block_size)
 
     assert plan.gpu_row_ids is not None
     unique_hosts = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist()
@@ -3271,12 +3272,51 @@ def test_hisparse_prefill_staging_plan_resolves_resident_sources():
         check_dtype=False,
     )
 
-    plan.ensure_gpu_sources(torch.zeros_like(resident_table), resident_block_size)
+    plan.ensure_gpu_sources(
+        torch.zeros_like(resident_table), state_indices, resident_block_size
+    )
 
     assert plan.gpu_row_ids is not None
     assert (plan.gpu_row_ids == -1).all()
 
 
+def test_hisparse_prefill_resident_sources_follow_each_resident_group():
+    """Layers of different resident groups share one batch staging plan.
+
+    Each group's layers must stage from that group's resident blocks. Keying
+    the plan's cached resident sources on a per-call batch gather let a later
+    group's gather reuse the freed address of an earlier one, so its layers
+    read the earlier group's block ids.
+    """
+    block_size = 4
+    plan = build_hisparse_prefill_staging_plan(
+        torch.tensor([[5, 2]], dtype=torch.int32),
+        torch.tensor([8], dtype=torch.int32),
+        block_size,
+        staging_block_capacity=2,
+    )
+    # The prefill runs in request state row 1; one table holds both groups.
+    # Rows are wide so that a cache keyed on a freed per-call gather of a row
+    # would see the next gather reuse its address.
+    state_rows = torch.zeros((4, 2, 1024), dtype=torch.int32)
+    state_rows[1, 0, :2] = torch.tensor([11, 12])
+    state_rows[1, 1, :2] = torch.tensor([21, 22])
+    state_indices = torch.tensor([1], dtype=torch.int32)
+    host_ids = (plan.row_ids[0].view(-1, block_size)[:, 0] // block_size).tolist()
+
+    for group in (0, 1, 0, 1):
+        plan.ensure_gpu_sources(state_rows[:, group], state_indices, block_size)
+        assert plan.gpu_row_ids is not None
+        gpu_blocks = plan.gpu_row_ids[0].view(-1, block_size)[:, 0] // block_size
+        staged = {
+            host_id: gpu_block
+            for host_id, gpu_block in zip(host_ids, gpu_blocks.tolist())
+            if host_id > 0
+        }
+        first_page, second_page = state_rows[1, group, :2].tolist()
+        assert staged == {5: first_page, 2: second_page}
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
 def test_hisparse_gather_prefill_cache_prefers_resident_rows():
     """Staged rows come from the resident cache when a shadow page exists."""
@@ -3294,7 +3334,11 @@ def test_hisparse_gather_prefill_cache_prefers_resident_rows():
         dtype=torch.int32,
         device=device,
     )
-    plan.ensure_gpu_sources(resident_table, resident_block_size)
+    plan.ensure_gpu_sources(
+        resident_table,
+        torch.arange(2, dtype=torch.int32, device=device),
+        resident_block_size,
+    )
 
     num_host_blocks, num_res_blocks = 10, 24
     host_cache = (
@@ -3677,6 +3721,7 @@ def test_hisparse_resident_prefill_uses_attention_block_stride():
         all_context_pages_resident=True,
         view=SimpleNamespace(block_size=64, attention_block_stride=832),
         block_table=torch.tensor([[3]], dtype=torch.int32),
+        batch_block_table=lambda: torch.tensor([[3]], dtype=torch.int32),
         runtime=SimpleNamespace(max_swap_rows=1),
     )
     index_group = object.__new__(HiSparseMLAIndexGroup)
@@ -4091,6 +4136,7 @@ def test_hisparse_prefill_reuses_builder_staging_plan():
     cache = SimpleNamespace(
         runtime=SimpleNamespace(
             gather_prefill_cache=gather,
+            request_state_indices=torch.tensor([0], dtype=torch.int32),
         ),
         view=SimpleNamespace(cache=resident_cache, block_size=1),
         block_table=resident_block_table,
@@ -4113,10 +4159,6 @@ def test_hisparse_prefill_reuses_builder_staging_plan():
     assert result is staged
     assert block_table is plan.block_table
     torch.testing.assert_close(request_ids, metadata.req_id_per_token)
-    plan.ensure_gpu_sources.assert_called_once()
-    args = plan.ensure_gpu_sources.call_args.args
-    torch.testing.assert_close(args[0], resident_block_table)
-    assert args[1] == 1
     assert calls == [(source, plan, resident_cache)]
 
 
@@ -4134,6 +4176,9 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch):
     cache = SimpleNamespace(
         view=SimpleNamespace(cache=resident_cache, block_size=4),
         block_table=torch.tensor([[2, 3]], dtype=torch.int32),
+        runtime=SimpleNamespace(
+            request_state_indices=torch.tensor([0], dtype=torch.int32)
+        ),
     )
     index_group = object.__new__(HiSparseMLAIndexGroup)
     index_group.caches = [cache]
@@ -4178,10 +4223,6 @@ def test_hisparse_fp8_prefill_gather_uses_dedicated_stream(monkeypatch):
     assert args.kwargs["host_cache"].data_ptr() == source.data_ptr()
     assert args.kwargs["host_row_ids"] is plan.row_ids
     assert args.kwargs["device_row_ids"] is plan.gpu_row_ids
-    plan.ensure_gpu_sources.assert_called_once()
-    ensure_args = plan.ensure_gpu_sources.call_args.args
-    torch.testing.assert_close(ensure_args[0], cache.block_table)
-    assert ensure_args[1] == 4
 
 
 def test_sparse_impl_observes_repointed_indexer_buffer():

tests/v1/core/test_prefix_caching.py

@@ -279,6 +279,12 @@ def _allocate_scheduled(
     return blocks
 
 
+def _resident_block_ids(manager: KVCacheManager, request_id: str) -> list[int]:
+    """The GPU block HiSparse reads each resident page from; 0 if host-only."""
+    resident = get_hisparse_coordinator(manager).resident_managers[0]
+    return [block.block_id for block in resident.get_residency_row(request_id)]
+
+
 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)
@@ -319,7 +325,7 @@ def test_hisparse_builds_dma_row_mirrors_across_pages():
 def test_hisparse_async_speculation_mirrors_uncertain_position_range():
     """Unresolved drafts must not leave gaps in the eager host mirror."""
     coordinator = MagicMock(host_group_id=0)
-    coordinator.take_block_table_updates.return_value = {}
+    coordinator.take_residency_updates.return_value = {}
     coordinator.build_offload_command.return_value = None
     coordinator.build_row_mirrors.return_value = ()
     scheduler = HiSparseConnectorScheduler(
@@ -332,7 +338,6 @@ def test_hisparse_async_speculation_mirrors_uncertain_position_range():
     connector.connector_scheduler = scheduler
     connector.update_state_after_alloc(request, None, 0)
     scheduler_output = SimpleNamespace(
-        block_table_updates=None,
         kv_cache_block_copies=None,
         scheduled_new_reqs=[],
         scheduled_cached_reqs=SimpleNamespace(
@@ -370,11 +375,66 @@ def test_hisparse_reports_when_context_is_fully_resident():
         hot_manager.require_hot(request.request_id)
     assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens)) is not None
     _publish_hisparse_pages(manager)
+    coordinator.take_residency_updates([request.request_id])
+    core_row = manager.get_block_ids(request.request_id)[2]
     pool = manager.block_pool
     pool.get_new_blocks(pool.get_num_free_blocks())
 
     assert not coordinator.all_context_pages_resident(scheduled)
-    assert coordinator.take_block_table_updates().keys() == {request.request_id}
+    # The block-table row stays append-only; the worker learns of the lost
+    # pages from the residency update instead.
+    assert manager.get_block_ids(request.request_id)[2] == core_row
+    resident_ids = _resident_block_ids(manager, request.request_id)
+    lost = [page for page, block_id in enumerate(resident_ids) if block_id == 0]
+    assert lost
+    update = coordinator.take_residency_updates([request.request_id])
+    assert update[request.request_id].pages == lost
+    assert update[request.request_id].block_ids[0] == [0] * len(lost)
+
+
+def test_hisparse_residency_update_sends_only_changed_pages():
+    """Losing one early page must not resend the pages after it: per-step
+    residency updates would otherwise grow with the context length."""
+    manager = make_hisparse_kv_cache_manager(32, 16)
+    tokens = list(range(6 * HISPARSE_BLOCK_SIZE))
+    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens))
+    coordinator = get_hisparse_coordinator(manager)
+    for hot_manager in coordinator.hot_managers:
+        hot_manager.require_hot(request.request_id)
+    assert _allocate_scheduled(manager, request, num_new_tokens=len(tokens))
+    coordinator.take_residency_updates([request.request_id])
+    pool = manager.block_pool
+    num_never_used = pool.get_num_free_blocks()
+    # Pages unpin as their host copies complete, so page 0 is reused first.
+    _publish_hisparse_pages(manager)
+    pool.get_new_blocks(num_never_used + 1)
+
+    update = coordinator.take_residency_updates([request.request_id])
+
+    assert update[request.request_id].pages == [0]
+    assert update[request.request_id].block_ids == ([0],)
+
+
+def test_hisparse_readmitted_request_resends_its_full_residency():
+    """A preempted request may resume in a worker state row holding another
+    request's residency, so its first update must cover every page again."""
+    manager = make_hisparse_kv_cache_manager(32, 16)
+    coordinator = get_hisparse_coordinator(manager)
+    tokens = list(range(2 * HISPARSE_BLOCK_SIZE))
+    request = make_request("request", tokens, HISPARSE_BLOCK_SIZE, sha256)
+    assert _allocate_scheduled(manager, request, len(tokens)) is not None
+    assert request.request_id in coordinator.take_residency_updates(
+        [request.request_id]
+    )
+    manager.free(request)
+
+    assert _allocate_scheduled(manager, request, len(tokens)) is not None
+    update = coordinator.take_residency_updates([request.request_id])
+    assert update[request.request_id].pages == [0, 1]
+    assert update[request.request_id].block_ids[0] == _resident_block_ids(
+        manager, request.request_id
+    )
 
 
 def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():
@@ -416,9 +476,7 @@ def test_hisparse_host_prefix_can_be_completed_by_indexer_offload():
     source, indexer, resident, hot = manager.get_blocks(resumed.request_id).blocks
     assert len(source) == len(indexer) == len(resident) == 4
     assert len(hot) == 2
-    assert not any(block.is_null for block in resident[:2])
-    assert not resident[2].is_null
-    assert not resident[3].is_null
+    assert 0 not in _resident_block_ids(manager, resumed.request_id)
     coordinator = get_hisparse_coordinator(manager)
     assert not coordinator.build_offload_command().page_transfers
     coordinator.finish_host_import(resumed.request_id, failed=False)
@@ -666,9 +724,10 @@ def test_hisparse_keeps_resident_pages_until_hot_buffer_is_allocated(enable_cach
     pool.free_blocks(held)
     assert _allocate_scheduled(manager, first, num_new_tokens=1) is not None
     assert all(m.has_hot("first") for m in coordinator.hot_managers)
+    coordinator.take_residency_updates(["first"])
     pool.get_new_blocks(pool.get_num_free_blocks())
     assert not coordinator.all_context_pages_resident(scheduled)
-    assert "first" in coordinator.take_block_table_updates()
+    assert "first" in coordinator.take_residency_updates(["first"])
 
 
 def test_hisparse_full_pool_keeps_pages_pinned_until_preemption():
@@ -1204,11 +1263,7 @@ def test_hisparse_prefix_hit_adopts_gpu_shadow_pages():
         num_new_computed_tokens=num_computed,
         new_computed_blocks=computed,
     )
-    resident_blocks = manager.get_blocks("resumed").blocks[2]
-    assert [block.block_id for block in resident_blocks[:3]] == (
-        original_resident_ids[:3]
-    )
-    assert not any(block.is_null for block in resident_blocks[:3])
+    assert _resident_block_ids(manager, "resumed")[:3] == original_resident_ids[:3]
     assert get_hisparse_coordinator(manager).all_context_pages_resident(
         ((resumed.request_id, num_computed, len(tokens) - num_computed),)
     )
@@ -1249,17 +1304,13 @@ def test_hisparse_prefix_hit_under_pressure_adopts_surviving_copies(
         new_computed_blocks=computed,
     )
 
-    resumed_blocks = manager.get_blocks("resumed").blocks
-    assert [block.block_id for block in resumed_blocks[2][:3]] == [
+    resident_ids = _resident_block_ids(manager, "resumed")
+    assert resident_ids[:3] == [
         copy_ids[page] if page in adopted_pages else pool.null_block.block_id
         for page in range(3)
     ]
-    gpu_ids = [
-        block.block_id
-        for group in resumed_blocks[1:]
-        for block in group
-        if not block.is_null
-    ]
+    _, indexer, _, hot = manager.get_block_ids("resumed")
+    gpu_ids = [block_id for block_id in indexer + resident_ids + hot if block_id]
     assert len(gpu_ids) == len(set(gpu_ids))
 
 
@@ -1488,8 +1539,7 @@ def test_hisparse_reused_copy_is_not_adopted():
         num_new_computed_tokens=num_computed,
         new_computed_blocks=computed,
     )
-    resident_blocks = manager.get_blocks("resumed").blocks[2]
-    assert all(block.is_null for block in resident_blocks[:3])
+    assert _resident_block_ids(manager, "resumed")[:3] == [0, 0, 0]
 
 
 def test_hisparse_external_import_uses_hard_gpu_footprint():

tests/v1/kv_connector/unit/test_hisparse_connector.py

@@ -28,10 +28,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiCo
 from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
 from vllm.v1.hisparse import runtime as runtime_module
 from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
-from vllm.v1.hisparse.runtime import HiSparseCacheHandle
+from vllm.v1.hisparse.runtime import (
+    HiSparseCacheHandle,
+    update_hisparse_residency,
+)
 from vllm.v1.hisparse.types import (
     SparseKVOffloadCommand,
     SparseKVPageTransfer,
+    SparseKVResidencyUpdate,
     SparseKVRowMirror,
 )
 from vllm.v1.metrics.stats import KVCacheEvictionEvent
@@ -183,7 +187,7 @@ def test_full_graph_step_prepares_host_mirror_outside_model():
     connector = object.__new__(HiSparseConnector)
     connector.connector_worker = worker
     connector._get_connector_metadata = MagicMock(
-        return_value=HiSparseConnectorMetadata(None, (), (), {}, True)
+        return_value=HiSparseConnectorMetadata(None, (), (), {}, True, {})
     )
     req_id_per_token = torch.tensor([0, 1], dtype=torch.int32)
     attn_metadata = SimpleNamespace(
@@ -321,6 +325,7 @@ def test_draft_layer_rows_mirrored_after_drafter(monkeypatch, cg_mode):
                 )
             },
             True,
+            {},
         )
     )
     # A verification step: four query tokens of one request fill the page.
@@ -354,7 +359,7 @@ def test_draft_layer_rows_mirrored_after_drafter(monkeypatch, cg_mode):
     # The next step's start mirrors the drafter's rows, then hands the page over.
     compute_stream.wait_event.reset_mock()
     connector._get_connector_metadata.return_value = HiSparseConnectorMetadata(
-        None, (), (), {}, True
+        None, (), (), {}, True, {}
     )
     connector.start_load_kv(
         SimpleNamespace(),
@@ -397,7 +402,8 @@ def test_draft_layer_rows_mirrored_after_drafter(monkeypatch, cg_mode):
 
 
 def test_scheduled_prefix_hit_publishes_adopted_copies():
-    """Copies adopted after scheduling must reach the worker's block table."""
+    """Copies adopted after scheduling reach the worker as a residency update,
+    leaving the block-table row the scheduler output carries untouched."""
     from tests.v1.core.test_prefix_caching import (
         HISPARSE_BLOCK_SIZE,
         _allocate_scheduled,
@@ -443,7 +449,39 @@ def test_scheduled_prefix_hit_publishes_adopted_copies():
         num_scheduled_tokens={resumed.request_id: num_new_tokens},
     )
 
-    scheduler.build_connector_meta(scheduler_output)
+    core_row = list(scheduler_output.scheduled_new_reqs[0].block_ids[2])
 
-    resident_ids = scheduler_output.block_table_updates[resumed.request_id][2]
-    assert resident_ids[:3] == copy_ids[:3]
+    metadata = scheduler.build_connector_meta(scheduler_output)
+
+    update = metadata.residency_updates[resumed.request_id]
+    assert update.pages == [0, 1, 2, 3]
+    assert update.block_ids[0][:3] == copy_ids[:3]
+    assert manager.get_block_ids(resumed.request_id)[2] == core_row
+    assert core_row[:3] == [0, 0, 0]
+
+
+def test_residency_updates_persist_by_state_row():
+    """Updates name requests by batch row but are stored by state row.
+
+    An update for a lost page and an appended page must leave the other pages
+    and the other request's row intact across steps that reorder the batch.
+    """
+    table = torch.zeros((4, 2, 4), dtype=torch.int32)
+    update_hisparse_residency(
+        table,
+        {
+            "a": SparseKVResidencyUpdate([0, 1, 2], ([1, 2, 3], [11, 12, 13])),
+            "b": SparseKVResidencyUpdate([0, 1], ([4, 5], [14, 15])),
+        },
+        ["a", "b"],
+        torch.tensor([2, 0], dtype=torch.int32),
+    )
+    update_hisparse_residency(
+        table,
+        {"a": SparseKVResidencyUpdate([1, 3], ([0, 6], [0, 16]))},
+        ["b", "a"],
+        torch.tensor([0, 2], dtype=torch.int32),
+    )
+
+    assert table[[0, 2], 0].tolist() == [[4, 5, 0, 0], [1, 0, 3, 6]]
+    assert table[[0, 2], 1].tolist() == [[14, 15, 0, 0], [11, 0, 13, 16]]

tests/v1/worker/test_utils.py

@@ -1068,6 +1068,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer
             row_mirrors={},
             all_context_pages_resident=True,
             row_mirrors_from_resident=False,
+            residency_updates={},
         ),
         None,
     )
@@ -1079,6 +1080,7 @@ def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer
             row_mirrors={},
             all_context_pages_resident=True,
             row_mirrors_from_resident=False,
+            residency_updates={},
         ),
         None,
     )
@@ -1152,6 +1154,7 @@ def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch):
             row_mirrors={},
             all_context_pages_resident=True,
             row_mirrors_from_resident=False,
+            residency_updates={},
         ),
         None,
     )

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

@@ -36,7 +36,11 @@ from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.hisparse.coordinator import get_hisparse_coordinator
 from vllm.v1.hisparse.layout import get_hisparse_steady_state_concurrency
-from vllm.v1.hisparse.types import SparseKVOffloadCommand, SparseKVRowMirror
+from vllm.v1.hisparse.types import (
+    SparseKVOffloadCommand,
+    SparseKVResidencyUpdate,
+    SparseKVRowMirror,
+)
 from vllm.v1.outputs import KVConnectorOutput
 
 if TYPE_CHECKING:
@@ -56,6 +60,7 @@ class HiSparseConnectorMetadata(KVConnectorMetadata):
     source_block_ids: tuple[int, ...]
     row_mirrors: dict[str, tuple[SparseKVRowMirror, ...]]
     all_context_pages_resident: bool
+    residency_updates: dict[str, SparseKVResidencyUpdate]
 
 
 @dataclass
@@ -143,7 +148,7 @@ class HiSparseConnectorScheduler:
                 scheduler_output.num_scheduled_tokens.items()
             )
         )
-        # Runs first: it plans this step's transfers and table updates.
+        # Runs first: it plans this step's transfers and residency changes.
         self.coordinator.advance_scheduled(
             (
                 request_id,
@@ -151,9 +156,6 @@ class HiSparseConnectorScheduler:
             )
             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
@@ -192,6 +194,9 @@ class HiSparseConnectorScheduler:
             tuple(source_block_ids),
             row_mirrors,
             self.coordinator.all_context_pages_resident(scheduled_requests),
+            self.coordinator.take_residency_updates(
+                scheduler_output.num_scheduled_tokens
+            ),
         )
 
     def update_connector_output(self, connector_output: KVConnectorOutput) -> None:

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

@@ -26,7 +26,11 @@ from vllm.utils.gpu_sync_debug import gpu_sync_allowed
 from vllm.utils.torch_utils import current_stream
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.hisparse.layout import HISPARSE_HOT_SUFFIX
-from vllm.v1.hisparse.runtime import HiSparseCacheHandle, release_pinned_state
+from vllm.v1.hisparse.runtime import (
+    HiSparseCacheHandle,
+    release_pinned_state,
+    update_hisparse_residency,
+)
 from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror
 from vllm.v1.kv_cache_interface import (
     HiSparseHotSpec,
@@ -334,6 +338,10 @@ class HiSparseConnectorWorker:
                 "HiSparse request-state mapping does not match max_num_seqs."
             )
         self.hot_backing = hot_backing
+        self.residency: torch.Tensor | None = next(
+            (cache.residency for cache in cache_handles if cache.residency is not None),
+            None,
+        )
         self._pending_invalid_block_ids: list[int] = []
         # Destination block ids of host copies this worker has run.
         self._completed_host_copy_dst_ids: list[int] = []
@@ -437,6 +445,15 @@ class HiSparseConnectorWorker:
         self._pending_invalid_block_ids.extend(metadata.source_block_ids)
         if request_state_indices is not None:
             self.set_request_state_indices(request_state_indices)
+        if metadata.residency_updates:
+            assert self.residency is not None
+            assert request_state_indices is not None and request_ids is not None
+            update_hisparse_residency(
+                self.residency,
+                metadata.residency_updates,
+                request_ids,
+                request_state_indices,
+            )
 
     def _clear_forward_mirror_state(self) -> None:
         self._per_layer_mirrored.clear()

vllm/v1/attention/backends/mla/index_group.py

@@ -243,7 +243,7 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
                 layer_index,
                 logical_topk_indices,
                 req_id_per_token,
-                leader.block_table,
+                leader.batch_block_table(),
                 leader.view.block_size,
                 block_stride_rows=leader.view.attention_block_stride,
                 return_valid_counts=return_valid_counts,
@@ -268,8 +268,11 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
         assert staging_plan is not None
         resident_cache = None
         if cache.view is not None and cache.block_table is not None:
+            state_indices = cache.runtime.request_state_indices
+            assert state_indices is not None
             staging_plan.ensure_gpu_sources(
-                cache.block_table[attn_metadata.num_decodes :],
+                cache.block_table,
+                state_indices[attn_metadata.num_decodes :],
                 cache.view.block_size,
             )
             resident_cache = cache.view.cache
@@ -291,8 +294,11 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):
         plan = prefill.host_staging_plan if prefill is not None else None
         assert plan is not None
         assert cache.view is not None and cache.block_table is not None
+        state_indices = cache.runtime.request_state_indices
+        assert state_indices is not None
         plan.ensure_gpu_sources(
-            cache.block_table[attn_metadata.num_decodes :],
+            cache.block_table,
+            state_indices[attn_metadata.num_decodes :],
             cache.view.block_size,
         )
         assert plan.gpu_row_ids is not None

vllm/v1/core/sched/output.py

@@ -304,9 +304,6 @@ class SchedulerOutput:
     # CoW copies to apply after zeroing new blocks and before forward.
     kv_cache_block_copies: list[KVCacheBlockCopy] | None = None
 
-    # Complete block-table rows that replace incrementally appended block IDs.
-    block_table_updates: dict[str, tuple[list[int], ...]] | None = None
-
     # Scheduler-local; always None by the time this reaches a worker.
     kv_connector_block_state: KVConnectorBlockState | None = None
 

vllm/v1/core/single_type_kv_cache_manager.py

@@ -2503,6 +2503,10 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):
 
     max_admission_blocks_per_request: int
 
+    def __init__(self, kv_cache_spec: HiSparseResidentSpec, **kwargs) -> None:
+        super().__init__(kv_cache_spec, **kwargs)
+        self.residency_changes: dict[str, dict[int, KVCacheBlock]] = {}
+
     def get_num_blocks_to_allocate(
         self,
         request_id: str,
@@ -2546,7 +2550,7 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):
         tail_page = (num_tokens - 1) // self.block_size
         if len(blocks) <= tail_page:
             blocks.extend([self._null_block] * (tail_page + 1 - len(blocks)))
-        if blocks[tail_page].is_null:
+        if self.get_resident_page(request_id, tail_page) is None:
             blocks[tail_page] = self.block_pool.get_new_blocks(1)[0]
 
     def add_local_computed_blocks(
@@ -2580,25 +2584,55 @@ class HiSparseResidentManager(_HiSparseAuxiliaryManager):
     def pop_blocks_for_free(self, request_id: str) -> list[KVCacheBlock]:
         assert self.coordinator is not None
         self.coordinator.free(request_id)
-        return super().pop_blocks_for_free(request_id)
+        blocks = self.get_residency_row(request_id)
+        self.residency_changes.pop(request_id, None)
+        super().pop_blocks_for_free(request_id)
+        return blocks
 
     def adopt_resident_page(
         self, request_id: str, block_idx: int, block: KVCacheBlock
     ) -> bool:
         """Point a null prefix page at a pinned GPU copy of its contents."""
         blocks = self.req_to_blocks.get(request_id)
-        if blocks is None or block_idx >= len(blocks) or not blocks[block_idx].is_null:
+        if blocks is None or block_idx >= len(blocks):
             return False
-        blocks[block_idx] = block
+        if self.get_resident_page(request_id, block_idx) is not None:
+            return False
+        self.residency_changes.setdefault(request_id, {})[block_idx] = block
         return True
 
+    def drop_resident_page(
+        self, request_id: str, block_idx: int
+    ) -> KVCacheBlock | None:
+        """Stop reading a page from the GPU; return the block it was read from."""
+        block = self.get_resident_page(request_id, block_idx)
+        if block is not None:
+            changes = self.residency_changes.setdefault(request_id, {})
+            changes[block_idx] = self._null_block
+        return block
+
     def get_resident_page(self, request_id: str, block_idx: int) -> KVCacheBlock | None:
         blocks = self.req_to_blocks.get(request_id)
         if blocks is None or block_idx >= len(blocks):
             return None
         block = blocks[block_idx]
+        changes = self.residency_changes.get(request_id)
+        if changes:
+            block = changes.get(block_idx, block)
         return None if block.is_null else block
 
+    def get_residency_row(self, request_id: str, start: int = 0) -> list[KVCacheBlock]:
+        """The request's row from ``start`` with residency changes applied. Unlike
+        ``get_resident_page``, host-only pages stay as the null block."""
+        blocks = self.req_to_blocks.get(request_id, [])[start:]
+        changes = self.residency_changes.get(request_id)
+        if not changes:
+            return blocks
+        return [
+            changes.get(block_idx, block)
+            for block_idx, block in enumerate(blocks, start)
+        ]
+
 
 def get_manager_for_kv_cache_spec(
     kv_cache_spec: KVCacheSpec,

vllm/v1/hisparse/binding.py

@@ -214,10 +214,26 @@ def bind_hisparse_kv_caches(
         for tensor_config in kv_cache_config.kv_cache_tensors
         for layer_index, name in enumerate(tensor_config.layers)
     }
-    resident_source_index = 0
-    for group_id, group in enumerate(kv_cache_config.kv_cache_groups):
-        if not isinstance(group.kv_cache_spec, HiSparseResidentSpec):
-            continue
+    resident_groups = [
+        (group_id, group)
+        for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
+        if isinstance(group.kv_cache_spec, HiSparseResidentSpec)
+    ]
+    residency: torch.Tensor | None = None
+    if resident_groups:
+        table_shapes = {
+            block_tables.input_block_tables[group_id].shape
+            for group_id, _ in resident_groups
+        }
+        assert len(table_shapes) == 1
+        max_num_reqs, max_num_pages = table_shapes.pop()
+        residency = torch.zeros(
+            (max_num_reqs, len(resident_groups), max_num_pages),
+            dtype=torch.int32,
+            device=block_tables.input_block_tables[resident_groups[0][0]].device,
+        )
+    for resident_source_index, (group_id, group) in enumerate(resident_groups):
+        assert residency is not None
         for cache_name in group.layer_names:
             assert cache_name.endswith(HISPARSE_RESIDENT_SUFFIX)
             layer_name = cache_name[: -len(HISPARSE_RESIDENT_SUFFIX)]
@@ -230,13 +246,13 @@ def bind_hisparse_kv_caches(
                 block_stride=tensor_config.block_stride,
                 num_blocks=kv_cache_config.num_blocks,
                 block_size=group.kv_cache_spec.block_size,
-                block_table=block_tables.input_block_tables[group_id],
+                block_table=residency[:, resident_source_index],
                 slot_mapping=block_tables.slot_mappings[group_id],
             )
+            cache_handle.residency = residency
             assert cache_handle.view is not None
             kv_caches[cache_name] = cache_handle.view.cache
             cache_handle.runtime.resident_source_index = resident_source_index
-        resident_source_index += 1
 
     hot_backing: torch.Tensor | None = None
     cache_handles: list[HiSparseCacheHandle] = []

vllm/v1/hisparse/coordinator.py

@@ -22,6 +22,7 @@ from vllm.v1.hisparse.types import (
     ACTIVE_TAIL_PAGES,
     SparseKVOffloadCommand,
     SparseKVPageTransfer,
+    SparseKVResidencyUpdate,
     SparseKVRowMirror,
 )
 from vllm.v1.kv_cache_interface import (
@@ -65,6 +66,10 @@ class _HiSparseRequestState:
     unpinned_pages: set[int] = field(default_factory=set)
     # Prefix pages whose GPU copies are adopted after the admitting allocation.
     pages_to_adopt: int = 0
+    # The worker holds this request's residency for every page below
+    # synced_pages, except the stale pages changed since it was sent.
+    synced_pages: int = 0
+    stale_pages: set[int] = field(default_factory=set)
 
 
 @dataclass
@@ -86,9 +91,10 @@ class HiSparseCoordinator:
     GPU-resident pages are a write-back cache of the host tier. Once a page's
     host copy is durable and its owner can read from host (it has, or has
     asked for, a hot region), the page's allocation reference is released with
-    ``BlockPool.unpin_blocks``: the block stays in the block table and keeps
-    being read, but the pool counts it as free and may hand it out, at which
-    point the owner's page is nulled and its block table republished.
+    ``BlockPool.unpin_blocks``: the block keeps being read, but the pool counts
+    it as free and may hand it out, at which point the owner stops reading it.
+    Residency changes reach the worker through ``take_residency_updates``, not
+    the block table, which stays append-only.
     """
 
     def __init__(
@@ -167,7 +173,6 @@ class HiSparseCoordinator:
         self._retained_copies: dict[int, tuple[KVCacheBlock, KVCacheBlock]] = {}
         # request -> prefix length an external load is still filling in.
         self._pending_imports: dict[str, int] = {}
-        self.block_table_updates: set[str] = set()
         # Requests whose resident pages were all non-null when last scanned.
         # Pages only become null at admission or import (before the first
         # scan) and when the pool reuses an unpinned page (_lose_page).
@@ -269,7 +274,8 @@ 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)
+                    if host_idx < state.synced_pages:
+                        state.stale_pages.add(host_idx)
 
     def _record_copies(self, request_id: str, num_computed_tokens: int) -> None:
         """Index the GPU copies of just-published host blocks for later hits."""
@@ -342,14 +348,11 @@ class HiSparseCoordinator:
     ) -> list[KVCacheBlock] | None:
         blocks: list[KVCacheBlock] = []
         for manager in self.resident_managers:
-            req_blocks = manager.req_to_blocks.get(request_id)
-            if (
-                req_blocks is None
-                or page_idx >= len(req_blocks) - ACTIVE_TAIL_PAGES
-                or req_blocks[page_idx].is_null
-            ):
+            num_pages = len(manager.req_to_blocks.get(request_id, ()))
+            block = manager.get_resident_page(request_id, page_idx)
+            if page_idx >= num_pages - ACTIVE_TAIL_PAGES or block is None:
                 return None
-            blocks.append(req_blocks[page_idx])
+            blocks.append(block)
         return blocks
 
     def _unpin_page(
@@ -392,16 +395,15 @@ class HiSparseCoordinator:
         if state is None:
             return
         for manager in self.resident_managers:
-            blocks = manager.req_to_blocks.get(request_id)
-            if blocks is None or page_idx >= len(blocks) or blocks[page_idx].is_null:
+            block = manager.drop_resident_page(request_id, page_idx)
+            if block is None:
                 continue
-            block = blocks[page_idx]
-            blocks[page_idx] = manager._null_block
             owners = self._owners.get(block.block_id)
             if owners is not None:
                 owners.discard((request_id, page_idx))
         state.unpinned_pages.discard(page_idx)
-        self.block_table_updates.add(request_id)
+        if page_idx < state.synced_pages:
+            state.stale_pages.add(page_idx)
         for hot_manager in self.hot_managers:
             hot_manager.require_hot(request_id)
 
@@ -627,11 +629,10 @@ class HiSparseCoordinator:
         mirrors: list[SparseKVRowMirror] = []
         for request_id, num_computed_tokens, num_scheduled_tokens in requests:
             host_blocks = self.host_manager.req_to_blocks.get(request_id)
-            resident_blocks = [
-                manager.req_to_blocks.get(request_id)
+            if host_blocks is None or any(
+                request_id not in manager.req_to_blocks
                 for manager in self.resident_managers
-            ]
-            if host_blocks is None or any(blocks is None for blocks in resident_blocks):
+            ):
                 continue
             token_position = num_computed_tokens
             end_position = token_position + num_scheduled_tokens
@@ -644,14 +645,12 @@ class HiSparseCoordinator:
                 # window only ever moves forward, so dropping them here would
                 # leave their host rows permanently stale.
                 source_starts = []
-                for blocks in resident_blocks:
-                    assert blocks is not None
-                    if page_idx >= len(blocks) or blocks[page_idx].is_null:
+                for manager in self.resident_managers:
+                    block = manager.get_resident_page(request_id, page_idx)
+                    if block is None:
                         break
-                    source_starts.append(
-                        blocks[page_idx].block_id * block_size + row_offset
-                    )
-                if len(source_starts) != len(resident_blocks):
+                    source_starts.append(block.block_id * block_size + row_offset)
+                if len(source_starts) != len(self.resident_managers):
                     token_position += num_rows
                     continue
                 host_block_idx = page_idx
@@ -693,21 +692,53 @@ class HiSparseCoordinator:
             if any(
                 block.is_null
                 for manager in self.resident_managers
-                for block in manager.req_to_blocks.get(request_id, ())
+                for block in manager.get_residency_row(request_id)
             ):
                 return False
             self._resident_requests.add(request_id)
         return True
 
-    def take_block_table_updates(self) -> dict[str, tuple[list[int], ...]]:
-        updates = {
-            request_id: tuple(
-                [block.block_id for block in manager.req_to_blocks.get(request_id, [])]
-                for manager in self.managers
+    def take_residency_updates(
+        self, request_ids: Iterable[str]
+    ) -> dict[str, SparseKVResidencyUpdate]:
+        """Resident block ids the worker lacks for the scheduled requests:
+        stale pages plus pages appended since the last update.
+
+        Requests that are not scheduled keep their changes until they are:
+        only a scheduled request's residency is read.
+        """
+        updates: dict[str, SparseKVResidencyUpdate] = {}
+        if not self.resident_managers:
+            return updates
+        first = self.resident_managers[0]
+        for request_id in request_ids:
+            state = self._get_request_state(request_id)
+            num_pages = len(first.req_to_blocks.get(request_id, ()))
+            start = state.synced_pages
+            if not state.stale_pages and start >= num_pages:
+                continue
+            stale_pages = sorted(state.stale_pages)
+            block_ids: list[list[int]] = []
+            for manager in self.resident_managers:
+                null_block_id = manager.block_pool.null_block.block_id
+                group_block_ids = [
+                    null_block_id if block is None else block.block_id
+                    for block in (
+                        manager.get_resident_page(request_id, page_idx)
+                        for page_idx in stale_pages
+                    )
+                ]
+                group_block_ids.extend(
+                    block.block_id
+                    for block in manager.get_residency_row(request_id, start)
+                )
+                block_ids.append(group_block_ids)
+            updates[request_id] = SparseKVResidencyUpdate(
+                pages=[*stale_pages, *range(start, num_pages)],
+                block_ids=tuple(block_ids),
             )
-            for request_id in self.block_table_updates
-        }
-        self.block_table_updates.clear()
+            state.stale_pages.clear()
+            state.synced_pages = num_pages
         return updates
 
     def build_offload_command(self) -> SparseKVOffloadCommand | None:
@@ -846,9 +877,7 @@ class HiSparseCoordinator:
             # planned, so its copies do not prove that the KV was computed.
             state.publication = None
         for manager in self.resident_managers:
-            blocks = manager.req_to_blocks.get(request_id)
-            if not blocks:
-                continue
+            blocks = manager.get_residency_row(request_id)
             # Tail first, so a prefix loses its tail pages before its head.
             for page_idx in reversed(range(len(blocks))):
                 block = blocks[page_idx]
@@ -865,7 +894,7 @@ class HiSparseCoordinator:
                     manager.block_pool.unpin_blocks([block], self._on_block_reused)
                 else:
                     continue
-                blocks[page_idx] = manager._null_block
+                manager.drop_resident_page(request_id, page_idx)
 
 
 def get_hisparse_coordinator(

vllm/v1/hisparse/runtime.py

@@ -7,10 +7,11 @@ from __future__ import annotations
 import math
 import mmap
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, TypeAlias
 
+import numpy as np
 import psutil
 import torch
 
@@ -21,8 +22,9 @@ from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv, round_up
 from vllm.utils.mem_utils import get_max_shared_memory_bytes
-from vllm.utils.torch_utils import current_stream
+from vllm.utils.torch_utils import async_tensor_h2d, current_stream
 from vllm.v1.attention.backend import max_decode_query_len
+from vllm.v1.hisparse.types import SparseKVResidencyUpdate
 from vllm.v1.kv_offload.cpu.shared_offload_region import SharedOffloadRegion
 from vllm.v1.simple_kv_offload.cuda_mem_ops import pin_tensor
 
@@ -499,16 +501,19 @@ class HiSparsePrefillStagingPlan:
 
     def ensure_gpu_sources(
         self,
-        resident_block_table: torch.Tensor,
+        resident_state_rows: torch.Tensor,
+        state_indices: torch.Tensor,
         resident_block_size: int,
     ) -> None:
         """Resolve which staged rows can be served from the resident cache.
 
-        Computed once per plan (the resident block table is shared by every
-        layer in the group); non-null resident pages become miss_mask=0 rows
-        gathered device-to-device by ``gather_prefill_cache``.
+        ``resident_state_rows`` is a resident group's persistent table by
+        request state row and ``state_indices`` holds each staged request's
+        state row. Computed once per plan and resident group (every layer in
+        the group shares the table); non-null resident pages become
+        miss_mask=0 rows gathered device-to-device by ``gather_prefill_cache``.
         """
-        source_key = (resident_block_table.data_ptr(), resident_block_size)
+        source_key = (resident_state_rows.data_ptr(), resident_block_size)
         if self.gpu_source_key == source_key:
             return
         block_size = self.block_size
@@ -519,8 +524,11 @@ class HiSparsePrefillStagingPlan:
         host_ids = self.row_ids[0].view(num_unique, block_size)[:, 0] // block_size
         new_bt = self.block_table.to(torch.int64)
         num_rows, num_cols = new_bt.shape
-        if num_rows == 0 or resident_block_table.shape[0] < num_rows:
+        if num_rows == 0 or state_indices.shape[0] < num_rows:
             return
+        resident_block_table = resident_state_rows.index_select(
+            0, state_indices[:num_rows].clamp(min=0)
+        )
         # One representative (row, col) per unique host block: any request
         # referencing the block holds an equivalent (refcounted) resident view.
         flat_pos = torch.arange(num_rows * num_cols, device=device)
@@ -1107,11 +1115,50 @@ class HiSparseRuntime:
         return physical_topk_indices
 
 
+def update_hisparse_residency(
+    residency: torch.Tensor,
+    updates: Mapping[str, SparseKVResidencyUpdate],
+    request_ids: Sequence[str],
+    request_state_indices: torch.Tensor,
+) -> None:
+    """Apply scheduled requests' residency changes.
+
+    ``residency`` is ``[state rows, resident groups, pages]``: the GPU block of
+    each page, or block 0 when the page is read from the host.
+    """
+    if not updates:
+        return
+    batch_rows = {request_id: row for row, request_id in enumerate(request_ids)}
+    num_groups = residency.shape[1]
+    rows: list[int] = []
+    pages: list[int] = []
+    block_ids: list[list[int]] = [[] for _ in range(num_groups)]
+    for request_id, update in updates.items():
+        rows.append(batch_rows[request_id])
+        pages.extend(update.pages)
+        for group_block_ids, update_block_ids in zip(
+            block_ids, update.block_ids, strict=True
+        ):
+            group_block_ids.extend(update_block_ids)
+    host_values = np.empty((len(pages), 2 + num_groups), dtype=np.int32)
+    host_values[:, 0] = np.repeat(
+        rows, [len(update.pages) for update in updates.values()]
+    )
+    host_values[:, 1] = pages
+    host_values[:, 2:] = np.asarray(block_ids, dtype=np.int32).T
+    values = async_tensor_h2d(
+        host_values, device=request_state_indices.device, dtype=torch.int32
+    )
+    state_rows = request_state_indices[values[:, 0]]
+    residency[state_rows, :, values[:, 1]] = values[:, 2:]
+
+
 class HiSparseCacheHandle:
     """Attention-facing handle for resident KV and sparse offload state."""
 
     def __init__(self, runtime: HiSparseRuntime) -> None:
         self.view: PagedCacheView | None = None
+        self.residency: torch.Tensor | None = None
         self.block_table: torch.Tensor | None = None
         self.source_block_table: torch.Tensor | None = None
         self.slot_mapping: torch.Tensor | None = None
@@ -1128,6 +1175,7 @@ class HiSparseCacheHandle:
         # Speculator layers write their rows after the target forward.
         self.draft_layer = False
         self.index_group_caches: list[HiSparseCacheHandle] = [self]
+        self._batch_block_table: torch.Tensor | None = None
 
     def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:
         assert self.runtime.is_group_leader
@@ -1136,6 +1184,7 @@ class HiSparseCacheHandle:
 
     def _prepare_for_batch(self, attn_metadata: Any | None) -> None:
         self.dummy_batch = attn_metadata is None
+        self._batch_block_table = None
         self.runtime.begin_forward()
         self.num_actual_tokens = (
             attn_metadata.num_actual_tokens if attn_metadata is not None else 0
@@ -1151,6 +1200,17 @@ class HiSparseCacheHandle:
             not self.decode_batch or self.runtime.eager_host_mirror
         )
 
+    def batch_block_table(self) -> torch.Tensor:
+        """Resident rows by this step's batch row, for the prefill paths."""
+        if self._batch_block_table is None:
+            assert self.block_table is not None
+            indices = self.runtime.request_state_indices
+            assert indices is not None
+            self._batch_block_table = self.block_table.index_select(
+                0, indices.clamp(min=0)
+            )
+        return self._batch_block_table
+
     def write_target(
         self, num_input_rows: int, num_slot_rows: int
     ) -> tuple[torch.Tensor, torch.Tensor, int]:

vllm/v1/hisparse/types.py

@@ -28,6 +28,15 @@ class SparseKVRowMirror:
     num_rows: int
 
 
+@dataclass(frozen=True)
+class SparseKVResidencyUpdate:
+    """GPU block ids of some of a request's resident pages, per resident group.
+    A null block id means the page is read from the host."""
+
+    pages: list[int]
+    block_ids: tuple[list[int], ...]
+
+
 @dataclass
 class SparseKVOffloadCommand:
     """Opaque scheduler-to-worker command for one model step."""

vllm/v1/worker/gpu/model_runner.py

@@ -1224,17 +1224,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
     def update_requests(self, scheduler_output: SchedulerOutput) -> None:
         # Add new blocks and update num_computed_tokens for the existing requests.
         reqs = scheduler_output.scheduled_cached_reqs
-        table_updates = scheduler_output.block_table_updates or {}
-        for req_id, block_ids in table_updates.items():
-            if (idx := self.req_states.req_id_to_index.get(req_id)) is not None:
-                self.block_tables.append_block_ids(idx, block_ids, overwrite=True)
         num_computed_tokens_np = self.req_states.num_computed_tokens_np
         for req_id, num_computed_tokens, req_new_block_ids in zip(
             reqs.req_ids, reqs.num_computed_tokens, reqs.new_block_ids
         ):
             req_index = self.req_states.req_id_to_index[req_id]
             num_computed_tokens_np[req_index] = num_computed_tokens
-            if req_new_block_ids is not None and req_id not in table_updates:
+            if req_new_block_ids is not None:
                 self.block_tables.append_block_ids(
                     req_index, req_new_block_ids, overwrite=False
                 )