[Bugfix][HiSparse] Fix MTP acceptance collapse under FULL graphs with a saturated GPU pool (#59309)

Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

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
 

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

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)

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 (

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]

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

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.

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(),

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,

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,
         )
 
 

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)