From d5f356dae093d2e47b97af7c9cd8d2fc1aedc56f Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Sat, 10 Oct 2026 00:25:44 -0400 Subject: [PATCH] [HiSparse] Own GPU residency instead of editing core block-table rows (#60647) Signed-off-by: Matthew Bonanni Co-authored-by: Claude Opus 5.5 --- csrc/libtorch_stable/hisparse_kernels.cu | 24 ++-- .../v1/attention/test_sparse_mla_backends.py | 63 ++++++++-- tests/v1/core/test_prefix_caching.py | 94 +++++++++++---- .../unit/test_hisparse_connector.py | 52 ++++++-- tests/v1/worker/test_utils.py | 3 + .../kv_connector/v1/hisparse/connector.py | 15 ++- .../kv_connector/v1/hisparse/worker.py | 19 ++- vllm/v1/attention/backends/mla/index_group.py | 12 +- vllm/v1/core/sched/output.py | 3 - vllm/v1/core/single_type_kv_cache_manager.py | 42 ++++++- vllm/v1/hisparse/binding.py | 28 ++++- vllm/v1/hisparse/coordinator.py | 111 +++++++++++------- vllm/v1/hisparse/runtime.py | 76 ++++++++++-- vllm/v1/hisparse/types.py | 9 ++ vllm/v1/worker/gpu/model_runner.py | 6 +- 15 files changed, 429 insertions(+), 128 deletions(-) diff --git a/csrc/libtorch_stable/hisparse_kernels.cu b/csrc/libtorch_stable/hisparse_kernels.cu index 069eddf3b0..5c7621c2f5 100644 --- a/csrc/libtorch_stable/hisparse_kernels.cu +++ b/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(request_row) * + resident_block_table[static_cast(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(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(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(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); diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index 160b766da0..6eb089ce3e 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/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(): diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index a89fd085f1..836286d127 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/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(): diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index f9598a24f8..9ae0d3bc8e 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/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]] diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 98395e90c4..84dd944e1d 100644 --- a/tests/v1/worker/test_utils.py +++ b/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, ) 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 863f52e025..4f17f1e737 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py +++ b/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: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py index b9dd01f054..bbccecd657 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/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() diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 39f205b91d..0dfda0ded9 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/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 diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index 9d16d2f710..a5af5a3d57 100644 --- a/vllm/v1/core/sched/output.py +++ b/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 diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index eed15427b2..77c6a4a5f7 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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, diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 60c60a3ff5..1420a45396 100644 --- a/vllm/v1/hisparse/binding.py +++ b/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] = [] diff --git a/vllm/v1/hisparse/coordinator.py b/vllm/v1/hisparse/coordinator.py index 390e0d27ce..374f100bc8 100644 --- a/vllm/v1/hisparse/coordinator.py +++ b/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( diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 1f7a5ad5ca..f1b0fa688e 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/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]: diff --git a/vllm/v1/hisparse/types.py b/vllm/v1/hisparse/types.py index db0ee06e93..37b25003df 100644 --- a/vllm/v1/hisparse/types.py +++ b/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.""" diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index f39fd0a6e6..674cf23527 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 )