diff --git a/docs/design/hisparse.md b/docs/design/hisparse.md index 2c72ddcbcc..8651b3a06d 100644 --- a/docs/design/hisparse.md +++ b/docs/design/hisparse.md @@ -133,7 +133,10 @@ duplicate mapping. Speculative decoding resolves all verification rows of a request in one pass: one block resolves the union of the rows' top-k against the request's hot-cache state, so rows that select the same host row share its hot row and no row -evicts a hot row another row of the step still reads. +evicts a hot row another row of the step still reads. Draft layers write their +rows after the target forward, so their host mirror runs at the start of the +next step, after the drafter, and the step's page transfers are submitted +behind it. ## P/D import target diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py index a6517430ad..b41db5a0f1 100644 --- a/tests/v1/attention/test_sparse_mla_backends.py +++ b/tests/v1/attention/test_sparse_mla_backends.py @@ -1943,6 +1943,10 @@ def _make_hisparse_cache_handle( return HiSparseCacheHandle(runtime) +def _valid_rows(num_rows: int) -> torch.Tensor: + return torch.tensor([num_rows], dtype=torch.int32, device=DEVICE_TYPE) + + @requires_hisparse_ops def test_hisparse_uses_graph_stable_request_state_mapping(): device = torch.device(DEVICE_TYPE) @@ -1969,6 +1973,7 @@ def test_hisparse_uses_graph_stable_request_state_mapping(): block_table=torch.tensor([[1]], dtype=torch.int32, device=device), logical_topk_indices=torch.tensor([[0]], dtype=torch.int32, device=device), block_size=block_size, + num_valid_rows=_valid_rows(1), ) torch.accelerator.synchronize() @@ -2011,6 +2016,7 @@ def test_hisparse_maps_speculative_rows_through_request_state(): block_table=torch.zeros((1, 1), dtype=torch.int32, device=device), logical_topk_indices=torch.zeros((4, 1), dtype=torch.int32, device=device), block_size=block_size, + num_valid_rows=_valid_rows(4), ) torch.accelerator.synchronize() @@ -2072,6 +2078,7 @@ def test_hisparse_speculative_rows_resolve_host_misses_consistently( block_table=block_table, logical_topk_indices=topk.clone(), block_size=block_size, + num_valid_rows=_valid_rows(num_rows), ) torch.accelerator.synchronize() @@ -2115,12 +2122,99 @@ def test_hisparse_speculative_rows_count_a_shared_load_once(): block_table=block_table, logical_topk_indices=topk, block_size=block_size, + num_valid_rows=_valid_rows(2), ) torch.accelerator.synchronize() assert runtime.index_group.swap_stats.cpu().tolist() == [top_k, top_k] +@requires_hisparse_ops +def test_hisparse_resolver_reads_current_request_ids(): + """The resolver must see this call's request ids, not the previous call's. + + It runs on the copy stream after waiting only on ``logical_topk_ready``, + which the indexer records before attention writes the ids. An MTP drafter + alternates four rows per request with one, so stale ids resolve rows + against other requests' KV. Stalling the compute stream exposes the race. + """ + device = torch.device(DEVICE_TYPE) + block_size, row_width, top_k, blocks_per_req, num_reqs = 64, 16, 64, 4, 2 + kv_pool = torch.randn( + (num_reqs * blocks_per_req + 1, block_size, row_width), dtype=torch.float32 + ).pin_memory() + runtime = _make_hisparse_runtime( + top_k=top_k, + device_buffer_size=4 * top_k, + max_num_reqs=num_reqs, + row_width=row_width, + block_size=block_size, + max_swap_rows=4 * num_reqs, + ) + runtime.bind_source_cache(kv_pool) + ready = torch.Event() + runtime.index_group.logical_topk_ready = ready + cache = HiSparseCacheHandle(runtime) + cache.all_context_pages_resident = False + cache.source_block_table = torch.arange( + 1, num_reqs * blocks_per_req + 1, dtype=torch.int32, device=device + ).view(num_reqs, blocks_per_req) + + index_group = object.__new__(HiSparseMLAIndexGroup) + index_group.caches = [cache] + index_group.physical_topk_indices = torch.empty( + (4 * num_reqs + 1, top_k), dtype=torch.int32, device=device + ) + + # Four rows per request, then one per request plus a CUDA-graph padding + # row, which maps to request 0 and must be skipped. + for rows_per_req, num_padding in ((4, 0), (1, 1)): + num_rows = num_reqs * rows_per_req + req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device) + req_ids = torch.cat( + [ + req_ids.repeat_interleave(rows_per_req), + torch.zeros(num_padding, dtype=torch.int32, device=device), + ] + ) + topk = torch.stack( + [ + torch.randperm(blocks_per_req * block_size, device=device)[:top_k] + for _ in range(req_ids.numel()) + ] + ).to(torch.int32) + query_start_loc = torch.arange( + 0, num_rows + 1, rows_per_req, dtype=torch.int32, device=device + ) + attn_metadata = SimpleNamespace( + req_id_per_token=req_ids, + query_start_loc=torch.cat([query_start_loc, query_start_loc[-1:]]), + block_size=block_size, + ) + runtime.begin_forward() + ready.record(current_stream()) + torch.cuda._sleep(200_000_000) + hot_indices = index_group.convert_logical_to_physical_topk( + 0, topk, attn_metadata, block_stride_rows=None, return_valid_counts=False + ) + torch.accelerator.synchronize() + + assert (hot_indices[num_rows:] == -1).all() + global_ref = _triton_convert_reference_impl( + req_ids[:num_rows], + cache.source_block_table, + topk[:num_rows], + block_size, + top_k, + ) + gathered = runtime.hot.attention_cache.reshape(-1, row_width)[ + hot_indices[:num_rows].to(torch.long) + ].cpu() + torch.testing.assert_close( + gathered, kv_pool.reshape(-1, row_width)[global_ref.cpu().to(torch.long)] + ) + + @requires_hisparse_ops def test_hisparse_union_table_scales_with_union_bound(): """Shared memory follows the per-request union bound, not the hot buffer. @@ -2157,6 +2251,7 @@ def test_hisparse_union_table_scales_with_union_bound(): block_table=block_table, logical_topk_indices=topk, block_size=block_size, + num_valid_rows=_valid_rows(req_ids.numel()), ) torch.accelerator.synchronize() @@ -2217,6 +2312,7 @@ def test_hisparse_resident_rows_bypass_hot_lru(): block_table=source_table, logical_topk_indices=topk, block_size=block_size, + num_valid_rows=_valid_rows(request_ids.numel()), return_valid_counts=True, ) torch.accelerator.synchronize() @@ -2280,6 +2376,7 @@ def test_hisparse_swap_in_preserves_rows_across_eviction(): block_table=block_table, logical_topk_indices=topk.clone(), block_size=block_size, + num_valid_rows=_valid_rows(req_ids.numel()), return_valid_counts=True, ) torch.accelerator.synchronize() @@ -2361,6 +2458,7 @@ def test_hisparse_multi_step_swaps_match_independent(): req_id_per_token=req_ids, block_table=block_table, block_size=block_size, + num_valid_rows=_valid_rows(req_ids.numel()), ) producer_indices = [ producer.swap_in(logical_topk_indices=topk.clone(), **kw) for topk in topks @@ -2435,6 +2533,7 @@ def test_hisparse_multi_step_writes_request_major_output(): block_table, logical[:, step], block_size=block_size, + num_valid_rows=_valid_rows(request_ids.numel()), return_valid_counts=True, attention_indices_out=physical[:, step], valid_counts_out=valid_counts[:, step], @@ -2619,6 +2718,7 @@ def test_hisparse_remaps_strided_hma_rows_for_attention(): [[0, 1, 2, 3]], dtype=torch.int32, device=device ), block_size=block_size, + num_valid_rows=_valid_rows(1), ) assert runtime.hot.attention_cache is not None @@ -2752,6 +2852,7 @@ def test_hisparse_newest_write_and_recycled_slot_invalidation(): block_table=block_table, logical_topk_indices=topk, block_size=block_size, + num_valid_rows=_valid_rows(req_ids.numel()), ) torch.accelerator.synchronize() stale_hot_slot = hot_indices[0, 0].item() @@ -2767,6 +2868,7 @@ def test_hisparse_newest_write_and_recycled_slot_invalidation(): block_table=block_table, logical_topk_indices=topk, block_size=block_size, + num_valid_rows=_valid_rows(req_ids.numel()), ) torch.accelerator.synchronize() idx = hot_indices.cpu().tolist()[0][0] @@ -3200,8 +3302,6 @@ def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention(): index_group.physical_topk_indices = torch.empty( (num_tokens + 1, 4), dtype=torch.int32, device=device ) - index_group.row_indices = torch.arange(num_tokens, dtype=torch.int32, device=device) - index_group.request_ids = torch.empty_like(index_group.row_indices) impl = SimpleNamespace( kv_lora_rank=1, index_group=index_group, @@ -3255,39 +3355,6 @@ def test_hisparse_keeps_speculative_rows_as_decodes(): assert builder.reorder_batch_threshold == 4 -def test_hisparse_decode_skips_padding_rows(): - """Padding rows are masked, and the ids stay in the group's persistent buffer. - - The resolver reads them on the HiSparse copy stream, where a compute-stream - temporary can already be recycled (CUDA graphs do), scrambling rows' requests. - """ - topk = torch.arange(20, dtype=torch.int32).view(10, 2) - cache = SimpleNamespace( - source_block_table=torch.zeros((2, 1), dtype=torch.int32), - swap_in=MagicMock(return_value=topk), - ) - index_group = object.__new__(HiSparseMLAIndexGroup) - index_group.caches = [cache] - index_group.physical_topk_indices = torch.empty((11, 2), dtype=torch.int32) - index_group.row_indices = torch.arange(10, dtype=torch.int32) - index_group.request_ids = torch.empty_like(index_group.row_indices) - metadata = SimpleNamespace( - query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32), - req_id_per_token=torch.tensor([0] * 8 + [1, 0], dtype=torch.int32), - block_size=64, - ) - - index_group.convert_logical_to_physical_topk( - 0, topk, metadata, block_stride_rows=None, return_valid_counts=False - ) - - request_ids = cache.swap_in.call_args.args[0] - assert request_ids.data_ptr() == index_group.request_ids.data_ptr() - torch.testing.assert_close( - request_ids, torch.tensor([0] * 8 + [1, -1], dtype=torch.int32) - ) - - def test_flashinfer_hisparse_decode_runs_batched_attention(): device = torch.device("cpu") num_tokens = 6 diff --git a/tests/v1/e2e/general/test_hisparse.py b/tests/v1/e2e/general/test_hisparse.py index 573b424ea7..2409b5cb69 100644 --- a/tests/v1/e2e/general/test_hisparse.py +++ b/tests/v1/e2e/general/test_hisparse.py @@ -95,6 +95,8 @@ def test_hisparse_spill_and_prefix_restore( A regression omitted attention metadata before FULL graph replay, so decode completed normally while its newly written KV rows were never copied to host. + MTP layers write after the target forward; their rows used to be mirrored + before the drafter wrote them, so the drafter later read stale host rows. """ capability = current_platform.get_device_capability() if capability is None or capability.major < 9: @@ -171,49 +173,70 @@ def test_hisparse_spill_and_prefix_restore( model_runner.cudagraph_manager, "run_fullgraph", record_fullgraph ) + # The MTP layer writes its rows after the target forward, so it is + # mirrored at the start of the next step instead of in finish_forward. + draft_layers = set(worker._draft_layers) + assert draft_layers original_finish_forward = worker.finish_forward - verified_full_replay_rows = 0 + original_finish_previous_step = worker._finish_previous_step + verified_rows = {"target": 0, "draft": 0} copied_nonzero_kv = False + draft_check_pending = False + + def check_host_rows(layer_indices: set[int], kind: str) -> None: + nonlocal copied_nonzero_kv + # Synchronize only the test's inspection, not the worker itself. + with gpu_sync_allowed(): + torch.accelerator.synchronize() + assert worker._row_mirrors + for layer_index in layer_indices: + cache = worker.cache_handles[layer_index] + source_index = cache.runtime.resident_source_index + for mirror in worker._row_mirrors: + source_row = mirror.source_starts[source_index] + destination_row = mirror.destination_start + num_rows = mirror.num_rows + source_block, source_offset = divmod( + source_row, worker.kernel_block_size + ) + with gpu_sync_allowed(): + gpu_rows = worker.resident_caches[layer_index][ + source_block, source_offset : source_offset + num_rows + ].cpu() + host_rows = worker.host_caches[layer_index][ + destination_row : destination_row + num_rows + ] + assert torch.equal(host_rows, gpu_rows) + copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item()) + verified_rows[kind] += num_rows def finish_forward(): - nonlocal full_replay_pending, verified_full_replay_rows - nonlocal copied_nonzero_kv + nonlocal full_replay_pending, draft_check_pending try: original_finish_forward() if not full_replay_pending: return - # Synchronize only the test's inspection, not finish_forward itself. - with gpu_sync_allowed(): - torch.accelerator.synchronize() - assert worker._row_mirrors - for layer_index, cache in enumerate(worker.cache_handles): - source_index = cache.runtime.resident_source_index - for mirror in worker._row_mirrors: - source_row = mirror.source_starts[source_index] - destination_row = mirror.destination_start - num_rows = mirror.num_rows - source_block, source_offset = divmod( - source_row, worker.kernel_block_size - ) - with gpu_sync_allowed(): - gpu_rows = worker.resident_caches[layer_index][ - source_block, source_offset : source_offset + num_rows - ].cpu() - host_rows = worker.host_caches[layer_index][ - destination_row : destination_row + num_rows - ] - assert torch.equal(host_rows, gpu_rows) - copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item()) - verified_full_replay_rows += num_rows + target_layers = set(range(len(worker.cache_handles))) - draft_layers + check_host_rows(target_layers, "target") + draft_check_pending = True finally: full_replay_pending = False + def finish_previous_step(): + nonlocal draft_check_pending + original_finish_previous_step() + if draft_check_pending: + draft_check_pending = False + check_host_rows(draft_layers, "draft") + monkeypatch.setattr(worker, "finish_forward", finish_forward) + monkeypatch.setattr(worker, "_finish_previous_step", finish_previous_step) expected = runner.generate_greedy([target], max_tokens=8) assert full_graph_calls > 0 - assert verified_full_replay_rows > 0 + assert verified_rows["target"] > 0 + assert verified_rows["draft"] > 0 assert copied_nonzero_kv runner.generate_greedy(pressure, max_tokens=8) diff --git a/tests/v1/kv_connector/unit/test_hisparse_connector.py b/tests/v1/kv_connector/unit/test_hisparse_connector.py index 384460fa2c..a2eeb8cd63 100644 --- a/tests/v1/kv_connector/unit/test_hisparse_connector.py +++ b/tests/v1/kv_connector/unit/test_hisparse_connector.py @@ -1,15 +1,21 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib +import ctypes from types import SimpleNamespace from unittest.mock import MagicMock import pytest import torch +from vllm.config.compilation import CUDAGraphMode from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import ( ExampleConnector, ) +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import ( + worker as worker_module, +) from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( HiSparseConnector, HiSparseConnectorMetadata, @@ -19,8 +25,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( HiSparseConnectorWorker, ) from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector +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.types import ( + SparseKVOffloadCommand, + SparseKVPageTransfer, + SparseKVRowMirror, +) from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector @@ -83,6 +95,7 @@ def test_full_graph_step_prepares_host_mirror_outside_model(): worker._group_leaders = (("layer", handle),) worker._per_layer_mirrored = set() worker._submitted_mirror_layers = set() + worker._draft_layers = () worker.is_host_writer = True worker._enqueue_row_dma = MagicMock() worker.start_step = MagicMock( @@ -116,6 +129,195 @@ def test_full_graph_step_prepares_host_mirror_outside_model(): runtime.invalidate_written_slots.assert_called_once() +class _FakeEvent: + log: list[tuple[str, object]] = [] + + def __init__(self, *args, **kwargs) -> None: + pass + + def record(self, stream=None) -> None: + self.log.append(("event", self)) + + def query(self) -> bool: + return True + + def synchronize(self) -> None: + pass + + +@pytest.mark.parametrize("cg_mode", [CUDAGraphMode.FULL, CUDAGraphMode.NONE]) +def test_draft_layer_rows_mirrored_after_drafter(monkeypatch, cg_mode): + """Draft-layer host rows must hold the drafter's writes before page release. + + The drafter writes its layers' rows after the target forward. Mirroring + them at the end of the target forward (the only mirror a FULL graph gets) + leaves pre-draft rows on the host, and the page is still reported clean. + """ + block_size, width = 4, 8 + resident_block, host_block, transfer_id = 2, 1, 7 + log: list[tuple[str, object]] = [] + _FakeEvent.log = log + + def swap_blocks_batch(src, dst, sizes): + for src_ptr, dst_ptr, size in zip(src.tolist(), dst.tolist(), sizes.tolist()): + ctypes.memmove(dst_ptr, src_ptr, size) + log.append(("dma", dst_ptr)) + + monkeypatch.setattr(torch, "Event", _FakeEvent) + monkeypatch.setattr(torch, "Stream", lambda *args, **kwargs: MagicMock()) + monkeypatch.setattr(torch.cuda, "Stream", lambda *args, **kwargs: MagicMock()) + monkeypatch.setattr(torch.cuda, "stream", lambda _: contextlib.nullcontext()) + compute_stream = MagicMock() + monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream) + monkeypatch.setattr(worker_module.ops, "swap_blocks_batch", swap_blocks_batch) + monkeypatch.setattr( + runtime_module, + "get_forward_context", + lambda: SimpleNamespace(cudagraph_runtime_mode=cg_mode), + ) + + slots = torch.arange(block_size) + resident_block * block_size + request_state_indices = torch.full((1,), -1, dtype=torch.int32) + handles = [] + for draft_layer in (False, True): + runtime = SimpleNamespace( + is_group_leader=True, + eager_host_mirror=True, + resident_source_index=0, + host_cache=torch.zeros(4 * block_size, width), + request_state_indices=request_state_indices, + begin_forward=MagicMock(), + invalidate_written_slots=MagicMock(), + ) + handle = HiSparseCacheHandle(runtime) + handle.view = SimpleNamespace( + block_size=block_size, cache=torch.zeros(4, block_size, width) + ) + handle.slot_mapping = slots + handle.mirror_slot_mapping = slots + handle.draft_layer = draft_layer + handles.append(handle) + target, draft = handles + # Rows the draft layer held before this step's drafter ran. + draft.view.cache[resident_block] = -1.0 + + worker = HiSparseConnectorWorker( + SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=16), + num_lookahead_tokens=3, + ), + MagicMock(), + ) + worker.initialize( + handles, + ["target", "draft"], + MagicMock(), + max_num_reqs=1, + host_num_blocks=4, + device=torch.device("cpu"), + pinned_host_pools=[], + ) + connector = object.__new__(HiSparseConnector) + connector.connector_worker = worker + connector._get_connector_metadata = MagicMock( + return_value=HiSparseConnectorMetadata( + SparseKVOffloadCommand( + [ + SparseKVPageTransfer( + transfer_id=transfer_id, + host_block_id=host_block, + resident_block_ids=(resident_block,), + after_forward=True, + ) + ] + ), + (), + (), + { + "req": ( + SparseKVRowMirror( + source_starts=(resident_block * block_size,), + destination_start=host_block * block_size, + num_rows=block_size, + ), + ) + }, + True, + ) + ) + # A verification step: four query tokens of one request fill the page. + attn_metadata = SimpleNamespace( + num_actual_tokens=block_size, + num_decode_tokens=0, + num_reqs=1, + max_query_len=block_size, + req_id_per_token=torch.zeros(block_size, dtype=torch.int32), + ) + + connector.start_load_kv( + SimpleNamespace(), + request_state_indices=None, + request_ids=["req"], + num_tokens=block_size, + attn_metadata={"target": attn_metadata, "draft": attn_metadata}, + ) + target.view.cache[resident_block] = 1.0 + if cg_mode != CUDAGraphMode.FULL: + target.finish_kv_update() + connector.finish_forward() + enqueued_before_draft = list(worker._enqueued_transfer_ids) + + # The drafter writes the draft layer's rows for the same positions. + if cg_mode != CUDAGraphMode.FULL: + draft.prepare_group_for_batch(attn_metadata) + draft.view.cache[resident_block] = 2.0 + if cg_mode != CUDAGraphMode.FULL: + draft.finish_kv_update() + # 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 + ) + connector.start_load_kv( + SimpleNamespace(), + request_state_indices=None, + request_ids=[], + num_tokens=0, + attn_metadata={}, + ) + ((completion_event, _),) = worker._pending_transfer_events + worker_meta = connector.build_connector_worker_meta() + + host_rows = slice(host_block * block_size, (host_block + 1) * block_size) + torch.testing.assert_close(target.runtime.host_cache[host_rows], torch.ones(4, 8)) + torch.testing.assert_close( + draft.runtime.host_cache[host_rows], torch.full((4, 8), 2.0) + ) + draft_host = draft.runtime.host_cache + draft_host_range = range( + draft_host.data_ptr(), + draft_host.data_ptr() + draft_host.numel() * draft_host.element_size(), + ) + draft_copies = [ + index + for index, (kind, ptr) in enumerate(log) + if kind == "dma" and ptr in draft_host_range + ] + # One draft-layer copy, and the page is not handed back before it. + assert len(draft_copies) == 1 + assert enqueued_before_draft == [] + assert log.index(("event", completion_event)) > draft_copies[0] + # The next forward waits for the draft copy: the scheduler may already + # have handed the blocks it reads and writes to another request. + draft_copy_event = next( + entry[1] for entry in log[draft_copies[0] :] if entry[0] == "event" + ) + assert compute_stream.wait_event.call_args_list[0].args == (draft_copy_event,) + assert worker_meta is not None + assert worker_meta.enqueued_transfer_counts == {transfer_id: 1} + assert worker_meta.completed_transfer_counts == {transfer_id: 1} + + def test_scheduled_prefix_hit_publishes_adopted_copies(): """Copies adopted after scheduling must reach the worker's block table.""" from tests.v1.core.test_prefix_caching import ( diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 43ac710f35..abcbdb24f2 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -44,6 +44,8 @@ def _make_hisparse_worker() -> HiSparseConnectorWorker: worker._row_mirror_num_rows = 0 worker._per_layer_mirrored = set() worker._submitted_mirror_layers = set() + worker._draft_layers = () + worker._draft_mirror_pending = False worker._pending_dma_descriptors = deque() worker._dma_free_descriptors = [] worker.host_write_events = (MagicMock(), MagicMock()) @@ -785,30 +787,6 @@ def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors(): worker._enqueue_row_dma.assert_not_called() -def test_hisparse_finish_forward_submits_lazy_post_forward_transfer(monkeypatch): - runtime = SimpleNamespace(eager_host_mirror=False) - worker = _make_hisparse_worker() - worker.is_host_writer = True - worker.cache_handles = [SimpleNamespace(runtime=runtime, num_actual_tokens=0)] - transfer = SparseKVPageTransfer(7, 2, (1,), after_forward=True) - worker._post_forward_transfers = [transfer] - worker._forward_ready_event = MagicMock() - worker._enqueue_host_mirror = MagicMock() - worker._submit_transfers = MagicMock() - worker._dma_submitted = False - worker._submitted_mirror_layers = set() - worker._finish_mirror_phase = MagicMock() - worker._release_completed_dma_descriptors = MagicMock() - stream = MagicMock() - monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) - - worker.finish_forward() - - worker._finish_mirror_phase.assert_called_once_with(worker._forward_ready_event) - worker._submit_transfers.assert_called_once_with([transfer]) - worker.host_write_event.record.assert_called_once_with(stream) - - def test_hisparse_prefill_mirrors_source_groups_and_flushes_partial_group(): slots = torch.tensor([7, 8], dtype=torch.int64) source_indices = [0, 0, 1, 1, 1, 1, 2] 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 b223616401..bf6329b3bf 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py @@ -301,6 +301,10 @@ class HiSparseConnectorWorker: self._submitted_mirror_layers: set[int] = set() self._layer_ready_events = tuple(torch.Event() for _ in cache_handles) self._forward_ready_event = torch.Event() + self._draft_layers = tuple( + index for index, cache in enumerate(cache_handles) if cache.draft_layer + ) + self._draft_mirror_pending = False self._set_row_mirrors(()) self.cache_layer_names = cache_layer_names self._group_leaders = tuple( @@ -344,9 +348,11 @@ class HiSparseConnectorWorker: self._init_dma() if self.is_host_writer: for layer_index, handle in enumerate(cache_handles): - handle.submit_layer_mirror = partial( - self._enqueue_layer_mirror, layer_index - ) + # Draft layers are mirrored once in the next start_step instead. + if not handle.draft_layer: + handle.submit_layer_mirror = partial( + self._enqueue_layer_mirror, layer_index + ) self._initialized = True def set_request_state_indices(self, indices: torch.Tensor) -> None: @@ -397,16 +403,17 @@ class HiSparseConnectorWorker: num_tokens: int = 0, ) -> None: self._stage_row_mirror_mapping(num_tokens) + self._finish_previous_step() previous_host_write_event = self.host_write_event self.host_write_event = self.host_write_events[self._next_host_write_event] self._next_host_write_event ^= 1 current_stream().wait_event(previous_host_write_event) self._release_completed_dma_descriptors() + self._dma_submitted = False mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids) if self._slot_mapping_staging is not None: self._slot_mapping_staging.candidates = mirrors self._set_row_mirrors(mirrors) - self._dma_submitted = False self._clear_forward_mirror_state() for handle in self.cache_handles: handle.all_context_pages_resident = metadata.all_context_pages_resident @@ -849,7 +856,12 @@ class HiSparseConnectorWorker: if num_rows == 0: return if self.is_host_writer: - expected_layers = {index for index, _ in active} + # The drafter has not written its layers' rows yet; the next + # start_step mirrors them after it runs. + self._draft_mirror_pending = bool(self._draft_layers) + expected_layers = {index for index, _ in active}.difference( + self._draft_layers + ) if self._per_layer_mirrored and self._per_layer_mirrored != expected_layers: raise RuntimeError( "HiSparse per-layer DMA did not mirror every active layer: " @@ -884,9 +896,6 @@ class HiSparseConnectorWorker: compute_stream = current_stream() self._forward_ready_event.record() self._finish_mirror_phase(self._forward_ready_event) - transfers = self._post_forward_transfers - self._post_forward_transfers = [] - self._submit_transfers(transfers) if self.is_host_writer: if self._dma_submitted: compute_stream.wait_event(self.host_write_event) @@ -895,6 +904,24 @@ class HiSparseConnectorWorker: self.host_write_event.record(compute_stream) self._release_completed_dma_descriptors() + def _finish_previous_step(self) -> None: + """Mirror the previous step's draft-layer rows, then hand its pages over. + + Runs before this step's work, so after the previous step's drafter. + Post-forward transfers are ordered behind the draft mirror, so a page is + never reported clean, and its resident block never released, before + its draft-layer rows reach the host. Both land in the previous step's + host write event, which this step waits on before its forward and its + host copies, since the scheduler may already have reused the + finished requests' blocks they read and write. + """ + if self._draft_mirror_pending: + self._draft_mirror_pending = False + self._enqueue_row_dma(self._draft_layers) + transfers = self._post_forward_transfers + self._post_forward_transfers = [] + self._submit_transfers(transfers) + def take_completed_host_copies(self) -> list[int]: """Drain host copies this worker has enqueued for this step.""" completed = self._completed_host_copy_dst_ids diff --git a/vllm/model_executor/layers/attention_layer_base.py b/vllm/model_executor/layers/attention_layer_base.py index cee57fdd04..d5332a47a8 100644 --- a/vllm/model_executor/layers/attention_layer_base.py +++ b/vllm/model_executor/layers/attention_layer_base.py @@ -21,6 +21,8 @@ class AttentionLayerBase(ABC): impl: "AttentionImpl" supports_dcp: bool = True + # Set by the speculator on its draft model's layers. + is_draft_layer: bool = False def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: """Bind the allocated KV cache tensor to this layer. diff --git a/vllm/v1/attention/backends/mla/index_group.py b/vllm/v1/attention/backends/mla/index_group.py index 21d578efd0..5a18332e41 100644 --- a/vllm/v1/attention/backends/mla/index_group.py +++ b/vllm/v1/attention/backends/mla/index_group.py @@ -39,8 +39,6 @@ class SparseMLAIndexGroup: logical_topk_indices: torch.Tensor physical_topk_indices: torch.Tensor valid_topk_counts: torch.Tensor - row_indices: torch.Tensor - request_ids: torch.Tensor side_stream: torch.Stream logical_topk_ready: torch.Event physical_topk_ready: torch.Event @@ -252,19 +250,13 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup): ) source_block_table = cache.source_block_table assert source_block_table is not None - # CUDA-graph padding rows past the batch's tokens map to request 0; - # mark them -1 so residency resolution skips them. - request_ids = self.request_ids[:num_tokens] - request_ids.copy_(req_id_per_token) - request_ids.masked_fill_( - self.row_indices[:num_tokens] >= attn_metadata.query_start_loc[-1], -1 - ) return cache.swap_in( - request_ids, + req_id_per_token, block_table=source_block_table, logical_topk_indices=logical_topk_indices, block_size=attn_metadata.block_size, return_valid_counts=return_valid_counts, + num_valid_rows=attn_metadata.query_start_loc[-1:], ) def stage_prefill_rows( @@ -384,16 +376,6 @@ class SparseMLAIndexGroupBuilder: dtype=torch.int32, device=self.logical_topk_indices.device, ), - row_indices=torch.arange( - workspace_rows, - dtype=torch.int32, - device=self.logical_topk_indices.device, - ), - request_ids=torch.empty( - workspace_rows, - dtype=torch.int32, - device=self.logical_topk_indices.device, - ), side_stream=_create_side_stream(self.logical_topk_indices.device), logical_topk_ready=_create_event(), physical_topk_ready=_create_event(), diff --git a/vllm/v1/hisparse/binding.py b/vllm/v1/hisparse/binding.py index 88c8976b56..d021d8b7be 100644 --- a/vllm/v1/hisparse/binding.py +++ b/vllm/v1/hisparse/binding.py @@ -284,6 +284,7 @@ def bind_hisparse_kv_caches( assert source_cache.untyped_storage().data_ptr() == ( host_pool.registered.untyped_storage().data_ptr() ) + cache_handle.draft_layer = forward_context[layer_name].is_draft_layer cache_handle.runtime.shared_host_region = host_pool.shared_region cache_handle.runtime.bind_source_cache( source_cache, diff --git a/vllm/v1/hisparse/runtime.py b/vllm/v1/hisparse/runtime.py index 2d9c6cc83f..fc415fa2ec 100644 --- a/vllm/v1/hisparse/runtime.py +++ b/vllm/v1/hisparse/runtime.py @@ -662,6 +662,10 @@ class HiSparseIndexGroup: self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device) self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous() self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k) + # Request ids the resolver reads, built on the copy stream: a + # compute-stream write would race with its logical_topk_ready wait. + self.row_indices = torch.arange(max_swap_rows, dtype=torch.int32, device=device) + self.request_ids = torch.empty_like(self.row_indices) self.copy_stream = ( copy_stream if copy_stream is not None else _create_copy_stream(device) ) @@ -997,6 +1001,7 @@ class HiSparseRuntime: shared_rows: slice, attention_indices_out: torch.Tensor | None, valid_counts_out: torch.Tensor | None, + num_valid_rows: torch.Tensor, ) -> None: group = self.index_group compute_stream = current_stream() @@ -1005,9 +1010,17 @@ class HiSparseRuntime: else: group.copy_stream.wait_stream(compute_stream) with group.copy_stream: + # CUDA-graph padding rows past the batch's tokens map to request 0; + # mark them -1 so residency resolution skips them. + num_tokens = logical_topk_indices.shape[0] + request_ids = group.request_ids[:num_tokens] + request_ids.copy_(req_id_per_token) + request_ids.masked_fill_( + group.row_indices[:num_tokens] >= num_valid_rows, -1 + ) self._resolve_residency( resident=resident, - req_id_per_token=req_id_per_token, + req_id_per_token=request_ids, block_table=block_table, topk_indices=logical_topk_indices, block_size=block_size, @@ -1051,6 +1064,7 @@ class HiSparseRuntime: block_table: torch.Tensor, logical_topk_indices: torch.Tensor, block_size: int, + num_valid_rows: torch.Tensor, return_valid_counts: bool = False, attention_indices_out: torch.Tensor | None = None, valid_counts_out: torch.Tensor | None = None, @@ -1074,6 +1088,7 @@ class HiSparseRuntime: shared_rows=shared_rows, attention_indices_out=attention_indices_out, valid_counts_out=valid_counts_out, + num_valid_rows=num_valid_rows, ) if not self._swap_staged: @@ -1110,6 +1125,8 @@ class HiSparseCacheHandle: self.mirror_staging_cache: torch.Tensor | None = None self.mirror_staging_slots: torch.Tensor | None = None self.submit_layer_mirror: Callable[[], None] | None = None + # Speculator layers write their rows after the target forward. + self.draft_layer = False self.index_group_caches: list[HiSparseCacheHandle] = [self] def prepare_group_for_batch(self, attn_metadata: Any | None) -> None: @@ -1203,6 +1220,7 @@ class HiSparseCacheHandle: logical_topk_indices: torch.Tensor, *, block_size: int, + num_valid_rows: torch.Tensor, return_valid_counts: bool = False, attention_indices_out: torch.Tensor | None = None, valid_counts_out: torch.Tensor | None = None, @@ -1216,6 +1234,7 @@ class HiSparseCacheHandle: return_valid_counts=return_valid_counts, attention_indices_out=attention_indices_out, valid_counts_out=valid_counts_out, + num_valid_rows=num_valid_rows, ) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 20a3b29eab..0728ddb87c 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -224,13 +224,13 @@ class DraftModelSpeculator(BaseSpeculator): self.model = self.load_draft_model(target_model, target_attn_layer_names) self._validate_local_argmax_reduction() - all_attn_layers = set[str]( - get_layers_from_vllm_config( - self.vllm_config, - AttentionLayerBase, # type: ignore[type-abstract] - ).keys() + all_attn_layers = get_layers_from_vllm_config( + self.vllm_config, + AttentionLayerBase, # type: ignore[type-abstract] ) - self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names + self.draft_attn_layer_names = set(all_attn_layers) - target_attn_layer_names + for layer_name in self.draft_attn_layer_names: + all_attn_layers[layer_name].is_draft_layer = True target_supports_mm = self.vllm_config.model_config.supports_multimodal_inputs draft_supports_mm = supports_multimodal_embeddings(self.model)