@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 4056c8ac1f

[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>

Lucas Wilkinson committed Oct 2, 2026 at 02:04 UTC 4056c8ac1f8a7e8fb50cf1c56ff96649fcadccc6

11 files changed +426 -122

docs/design/hisparse.md

+4 -1

@@ -133,7 +133,10 @@ duplicate mapping.

133 Speculative decoding resolves all verification rows of a request in one pass:

134 one block resolves the union of the rows' top-k against the request's hot-cache

135 state, so rows that select the same host row share its hot row and no row

136 -evicts a hot row another row of the step still reads.

136 +evicts a hot row another row of the step still reads. Draft layers write their

137 +rows after the target forward, so their host mirror runs at the start of the

138 +next step, after the drafter, and the step's page transfers are submitted

139 +behind it.

140

141 ## P/D import target

142

tests/v1/attention/test_sparse_mla_backends.py

+102 -35

@@ -1943,6 +1943,10 @@ def _make_hisparse_cache_handle(

1943 return HiSparseCacheHandle(runtime)

1944

1945

1946 +def _valid_rows(num_rows: int) -> torch.Tensor:

1947 + return torch.tensor([num_rows], dtype=torch.int32, device=DEVICE_TYPE)

1948 +

1949 +

1950 @requires_hisparse_ops

1951 def test_hisparse_uses_graph_stable_request_state_mapping():

1952 device = torch.device(DEVICE_TYPE)

@@ -1969,6 +1973,7 @@ def test_hisparse_uses_graph_stable_request_state_mapping():

1973 block_table=torch.tensor([[1]], dtype=torch.int32, device=device),

1974 logical_topk_indices=torch.tensor([[0]], dtype=torch.int32, device=device),

1975 block_size=block_size,

1976 + num_valid_rows=_valid_rows(1),

1977 )

1978 torch.accelerator.synchronize()

1979

@@ -2011,6 +2016,7 @@ def test_hisparse_maps_speculative_rows_through_request_state():

2016 block_table=torch.zeros((1, 1), dtype=torch.int32, device=device),

2017 logical_topk_indices=torch.zeros((4, 1), dtype=torch.int32, device=device),

2018 block_size=block_size,

2019 + num_valid_rows=_valid_rows(4),

2020 )

2021 torch.accelerator.synchronize()

2022

@@ -2072,6 +2078,7 @@ def test_hisparse_speculative_rows_resolve_host_misses_consistently(

2078 block_table=block_table,

2079 logical_topk_indices=topk.clone(),

2080 block_size=block_size,

2081 + num_valid_rows=_valid_rows(num_rows),

2082 )

2083 torch.accelerator.synchronize()

2084

@@ -2115,12 +2122,99 @@ def test_hisparse_speculative_rows_count_a_shared_load_once():

2122 block_table=block_table,

2123 logical_topk_indices=topk,

2124 block_size=block_size,

2125 + num_valid_rows=_valid_rows(2),

2126 )

2127 torch.accelerator.synchronize()

2128

2129 assert runtime.index_group.swap_stats.cpu().tolist() == [top_k, top_k]

2130

2131

2132 +@requires_hisparse_ops

2133 +def test_hisparse_resolver_reads_current_request_ids():

2134 + """The resolver must see this call's request ids, not the previous call's.

2135 +

2136 + It runs on the copy stream after waiting only on ``logical_topk_ready``,

2137 + which the indexer records before attention writes the ids. An MTP drafter

2138 + alternates four rows per request with one, so stale ids resolve rows

2139 + against other requests' KV. Stalling the compute stream exposes the race.

2140 + """

2141 + device = torch.device(DEVICE_TYPE)

2142 + block_size, row_width, top_k, blocks_per_req, num_reqs = 64, 16, 64, 4, 2

2143 + kv_pool = torch.randn(

2144 + (num_reqs * blocks_per_req + 1, block_size, row_width), dtype=torch.float32

2145 + ).pin_memory()

2146 + runtime = _make_hisparse_runtime(

2147 + top_k=top_k,

2148 + device_buffer_size=4 * top_k,

2149 + max_num_reqs=num_reqs,

2150 + row_width=row_width,

2151 + block_size=block_size,

2152 + max_swap_rows=4 * num_reqs,

2153 + )

2154 + runtime.bind_source_cache(kv_pool)

2155 + ready = torch.Event()

2156 + runtime.index_group.logical_topk_ready = ready

2157 + cache = HiSparseCacheHandle(runtime)

2158 + cache.all_context_pages_resident = False

2159 + cache.source_block_table = torch.arange(

2160 + 1, num_reqs * blocks_per_req + 1, dtype=torch.int32, device=device

2161 + ).view(num_reqs, blocks_per_req)

2162 +

2163 + index_group = object.__new__(HiSparseMLAIndexGroup)

2164 + index_group.caches = [cache]

2165 + index_group.physical_topk_indices = torch.empty(

2166 + (4 * num_reqs + 1, top_k), dtype=torch.int32, device=device

2167 + )

2168 +

2169 + # Four rows per request, then one per request plus a CUDA-graph padding

2170 + # row, which maps to request 0 and must be skipped.

2171 + for rows_per_req, num_padding in ((4, 0), (1, 1)):

2172 + num_rows = num_reqs * rows_per_req

2173 + req_ids = torch.arange(num_reqs, dtype=torch.int32, device=device)

2174 + req_ids = torch.cat(

2175 + [

2176 + req_ids.repeat_interleave(rows_per_req),

2177 + torch.zeros(num_padding, dtype=torch.int32, device=device),

2178 + ]

2179 + )

2180 + topk = torch.stack(

2181 + [

2182 + torch.randperm(blocks_per_req * block_size, device=device)[:top_k]

2183 + for _ in range(req_ids.numel())

2184 + ]

2185 + ).to(torch.int32)

2186 + query_start_loc = torch.arange(

2187 + 0, num_rows + 1, rows_per_req, dtype=torch.int32, device=device

2188 + )

2189 + attn_metadata = SimpleNamespace(

2190 + req_id_per_token=req_ids,

2191 + query_start_loc=torch.cat([query_start_loc, query_start_loc[-1:]]),

2192 + block_size=block_size,

2193 + )

2194 + runtime.begin_forward()

2195 + ready.record(current_stream())

2196 + torch.cuda._sleep(200_000_000)

2197 + hot_indices = index_group.convert_logical_to_physical_topk(

2198 + 0, topk, attn_metadata, block_stride_rows=None, return_valid_counts=False

2199 + )

2200 + torch.accelerator.synchronize()

2201 +

2202 + assert (hot_indices[num_rows:] == -1).all()

2203 + global_ref = _triton_convert_reference_impl(

2204 + req_ids[:num_rows],

2205 + cache.source_block_table,

2206 + topk[:num_rows],

2207 + block_size,

2208 + top_k,

2209 + )

2210 + gathered = runtime.hot.attention_cache.reshape(-1, row_width)[

2211 + hot_indices[:num_rows].to(torch.long)

2212 + ].cpu()

2213 + torch.testing.assert_close(

2214 + gathered, kv_pool.reshape(-1, row_width)[global_ref.cpu().to(torch.long)]

2215 + )

2216 +

2217 +

2218 @requires_hisparse_ops

2219 def test_hisparse_union_table_scales_with_union_bound():

2220 """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():

2251 block_table=block_table,

2252 logical_topk_indices=topk,

2253 block_size=block_size,

2254 + num_valid_rows=_valid_rows(req_ids.numel()),

2255 )

2256 torch.accelerator.synchronize()

2257

@@ -2217,6 +2312,7 @@ def test_hisparse_resident_rows_bypass_hot_lru():

2312 block_table=source_table,

2313 logical_topk_indices=topk,

2314 block_size=block_size,

2315 + num_valid_rows=_valid_rows(request_ids.numel()),

2316 return_valid_counts=True,

2317 )

2318 torch.accelerator.synchronize()

@@ -2280,6 +2376,7 @@ def test_hisparse_swap_in_preserves_rows_across_eviction():

2376 block_table=block_table,

2377 logical_topk_indices=topk.clone(),

2378 block_size=block_size,

2379 + num_valid_rows=_valid_rows(req_ids.numel()),

2380 return_valid_counts=True,

2381 )

2382 torch.accelerator.synchronize()

@@ -2361,6 +2458,7 @@ def test_hisparse_multi_step_swaps_match_independent():

2458 req_id_per_token=req_ids,

2459 block_table=block_table,

2460 block_size=block_size,

2461 + num_valid_rows=_valid_rows(req_ids.numel()),

2462 )

2463 producer_indices = [

2464 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():

2533 block_table,

2534 logical[:, step],

2535 block_size=block_size,

2536 + num_valid_rows=_valid_rows(request_ids.numel()),

2537 return_valid_counts=True,

2538 attention_indices_out=physical[:, step],

2539 valid_counts_out=valid_counts[:, step],

@@ -2619,6 +2718,7 @@ def test_hisparse_remaps_strided_hma_rows_for_attention():

2718 [[0, 1, 2, 3]], dtype=torch.int32, device=device

2719 ),

2720 block_size=block_size,

2721 + num_valid_rows=_valid_rows(1),

2722 )

2723

2724 assert runtime.hot.attention_cache is not None

@@ -2752,6 +2852,7 @@ def test_hisparse_newest_write_and_recycled_slot_invalidation():

2852 block_table=block_table,

2853 logical_topk_indices=topk,

2854 block_size=block_size,

2855 + num_valid_rows=_valid_rows(req_ids.numel()),

2856 )

2857 torch.accelerator.synchronize()

2858 stale_hot_slot = hot_indices[0, 0].item()

@@ -2767,6 +2868,7 @@ def test_hisparse_newest_write_and_recycled_slot_invalidation():

2868 block_table=block_table,

2869 logical_topk_indices=topk,

2870 block_size=block_size,

2871 + num_valid_rows=_valid_rows(req_ids.numel()),

2872 )

2873 torch.accelerator.synchronize()

2874 idx = hot_indices.cpu().tolist()[0][0]

@@ -3200,8 +3302,6 @@ def test_hisparse_fp8_decode_resolves_rows_once_then_runs_batched_attention():

3302 index_group.physical_topk_indices = torch.empty(

3303 (num_tokens + 1, 4), dtype=torch.int32, device=device

3304 )

3203 - index_group.row_indices = torch.arange(num_tokens, dtype=torch.int32, device=device)

3204 - index_group.request_ids = torch.empty_like(index_group.row_indices)

3305 impl = SimpleNamespace(

3306 kv_lora_rank=1,

3307 index_group=index_group,

@@ -3255,39 +3355,6 @@ def test_hisparse_keeps_speculative_rows_as_decodes():

3355 assert builder.reorder_batch_threshold == 4

3356

3357

3258 -def test_hisparse_decode_skips_padding_rows():

3259 - """Padding rows are masked, and the ids stay in the group's persistent buffer.

3260 -

3261 - The resolver reads them on the HiSparse copy stream, where a compute-stream

3262 - temporary can already be recycled (CUDA graphs do), scrambling rows' requests.

3263 - """

3264 - topk = torch.arange(20, dtype=torch.int32).view(10, 2)

3265 - cache = SimpleNamespace(

3266 - source_block_table=torch.zeros((2, 1), dtype=torch.int32),

3267 - swap_in=MagicMock(return_value=topk),

3268 - )

3269 - index_group = object.__new__(HiSparseMLAIndexGroup)

3270 - index_group.caches = [cache]

3271 - index_group.physical_topk_indices = torch.empty((11, 2), dtype=torch.int32)

3272 - index_group.row_indices = torch.arange(10, dtype=torch.int32)

3273 - index_group.request_ids = torch.empty_like(index_group.row_indices)

3274 - metadata = SimpleNamespace(

3275 - query_start_loc=torch.tensor([0, 8, 9], dtype=torch.int32),

3276 - req_id_per_token=torch.tensor([0] * 8 + [1, 0], dtype=torch.int32),

3277 - block_size=64,

3278 - )

3279 -

3280 - index_group.convert_logical_to_physical_topk(

3281 - 0, topk, metadata, block_stride_rows=None, return_valid_counts=False

3282 - )

3283 -

3284 - request_ids = cache.swap_in.call_args.args[0]

3285 - assert request_ids.data_ptr() == index_group.request_ids.data_ptr()

3286 - torch.testing.assert_close(

3287 - request_ids, torch.tensor([0] * 8 + [1, -1], dtype=torch.int32)

3288 - )

3289 -

3290 -

3358 def test_flashinfer_hisparse_decode_runs_batched_attention():

3359 device = torch.device("cpu")

3360 num_tokens = 6

tests/v1/e2e/general/test_hisparse.py

+50 -27

@@ -95,6 +95,8 @@ def test_hisparse_spill_and_prefix_restore(

96 A regression omitted attention metadata before FULL graph replay, so decode

97 completed normally while its newly written KV rows were never copied to host.

98 + MTP layers write after the target forward; their rows used to be mirrored

99 + before the drafter wrote them, so the drafter later read stale host rows.

100 """

101 capability = current_platform.get_device_capability()

102 if capability is None or capability.major < 9:

@@ -171,49 +173,70 @@ def test_hisparse_spill_and_prefix_restore(

173 model_runner.cudagraph_manager, "run_fullgraph", record_fullgraph

174 )

175

176 + # The MTP layer writes its rows after the target forward, so it is

177 + # mirrored at the start of the next step instead of in finish_forward.

178 + draft_layers = set(worker._draft_layers)

179 + assert draft_layers

180 original_finish_forward = worker.finish_forward

175 - verified_full_replay_rows = 0

181 + original_finish_previous_step = worker._finish_previous_step

182 + verified_rows = {"target": 0, "draft": 0}

183 copied_nonzero_kv = False

184 + draft_check_pending = False

185

178 - def finish_forward():

179 - nonlocal full_replay_pending, verified_full_replay_rows

186 + def check_host_rows(layer_indices: set[int], kind: str) -> None:

187 nonlocal copied_nonzero_kv

188 + # Synchronize only the test's inspection, not the worker itself.

189 + with gpu_sync_allowed():

190 + torch.accelerator.synchronize()

191 + assert worker._row_mirrors

192 + for layer_index in layer_indices:

193 + cache = worker.cache_handles[layer_index]

194 + source_index = cache.runtime.resident_source_index

195 + for mirror in worker._row_mirrors:

196 + source_row = mirror.source_starts[source_index]

197 + destination_row = mirror.destination_start

198 + num_rows = mirror.num_rows

199 + source_block, source_offset = divmod(

200 + source_row, worker.kernel_block_size

201 + )

202 + with gpu_sync_allowed():

203 + gpu_rows = worker.resident_caches[layer_index][

204 + source_block, source_offset : source_offset + num_rows

205 + ].cpu()

206 + host_rows = worker.host_caches[layer_index][

207 + destination_row : destination_row + num_rows

208 + ]

209 + assert torch.equal(host_rows, gpu_rows)

210 + copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item())

211 + verified_rows[kind] += num_rows

212 +

213 + def finish_forward():

214 + nonlocal full_replay_pending, draft_check_pending

215 try:

216 original_finish_forward()

217 if not full_replay_pending:

218 return

185 - # Synchronize only the test's inspection, not finish_forward itself.

186 - with gpu_sync_allowed():

187 - torch.accelerator.synchronize()

188 - assert worker._row_mirrors

189 - for layer_index, cache in enumerate(worker.cache_handles):

190 - source_index = cache.runtime.resident_source_index

191 - for mirror in worker._row_mirrors:

192 - source_row = mirror.source_starts[source_index]

193 - destination_row = mirror.destination_start

194 - num_rows = mirror.num_rows

195 - source_block, source_offset = divmod(

196 - source_row, worker.kernel_block_size

197 - )

198 - with gpu_sync_allowed():

199 - gpu_rows = worker.resident_caches[layer_index][

200 - source_block, source_offset : source_offset + num_rows

201 - ].cpu()

202 - host_rows = worker.host_caches[layer_index][

203 - destination_row : destination_row + num_rows

204 - ]

205 - assert torch.equal(host_rows, gpu_rows)

206 - copied_nonzero_kv |= bool(torch.count_nonzero(gpu_rows).item())

207 - verified_full_replay_rows += num_rows

219 + target_layers = set(range(len(worker.cache_handles))) - draft_layers

220 + check_host_rows(target_layers, "target")

221 + draft_check_pending = True

222 finally:

223 full_replay_pending = False

224

225 + def finish_previous_step():

226 + nonlocal draft_check_pending

227 + original_finish_previous_step()

228 + if draft_check_pending:

229 + draft_check_pending = False

230 + check_host_rows(draft_layers, "draft")

231 +

232 monkeypatch.setattr(worker, "finish_forward", finish_forward)

233 + monkeypatch.setattr(worker, "_finish_previous_step", finish_previous_step)

234

235 expected = runner.generate_greedy([target], max_tokens=8)

236

237 assert full_graph_calls > 0

216 - assert verified_full_replay_rows > 0

238 + assert verified_rows["target"] > 0

239 + assert verified_rows["draft"] > 0

240 assert copied_nonzero_kv

241

242 runner.generate_greedy(pressure, max_tokens=8)

tests/v1/kv_connector/unit/test_hisparse_connector.py

+202

@@ -1,15 +1,21 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import contextlib

4 +import ctypes

5 from types import SimpleNamespace

6 from unittest.mock import MagicMock

8 import pytest

9 import torch

11 +from vllm.config.compilation import CUDAGraphMode

12 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole

13 from vllm.distributed.kv_transfer.kv_connector.v1.example_connector import (

14 ExampleConnector,

15 )

16 +from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import (

17 + worker as worker_module,

18 +)

19 from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import (

20 HiSparseConnector,

21 HiSparseConnectorMetadata,

@@ -19,8 +25,14 @@ from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (

25 HiSparseConnectorWorker,

26 )

27 from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import MultiConnector

28 +from vllm.v1.hisparse import runtime as runtime_module

29 from vllm.v1.hisparse.coordinator import get_hisparse_coordinator

30 from vllm.v1.hisparse.runtime import HiSparseCacheHandle

31 +from vllm.v1.hisparse.types import (

32 + SparseKVOffloadCommand,

33 + SparseKVPageTransfer,

34 + SparseKVRowMirror,

35 +)

36 from vllm.v1.worker.gpu.kv_connector import ActiveKVConnector

@@ -83,6 +95,7 @@ def test_full_graph_step_prepares_host_mirror_outside_model():

95 worker._group_leaders = (("layer", handle),)

96 worker._per_layer_mirrored = set()

97 worker._submitted_mirror_layers = set()

98 + worker._draft_layers = ()

99 worker.is_host_writer = True

100 worker._enqueue_row_dma = MagicMock()

101 worker.start_step = MagicMock(

@@ -116,6 +129,195 @@ def test_full_graph_step_prepares_host_mirror_outside_model():

129 runtime.invalidate_written_slots.assert_called_once()

130

131

132 +class _FakeEvent:

133 + log: list[tuple[str, object]] = []

134 +

135 + def __init__(self, *args, **kwargs) -> None:

136 + pass

137 +

138 + def record(self, stream=None) -> None:

139 + self.log.append(("event", self))

140 +

141 + def query(self) -> bool:

142 + return True

143 +

144 + def synchronize(self) -> None:

145 + pass

146 +

147 +

148 +@pytest.mark.parametrize("cg_mode", [CUDAGraphMode.FULL, CUDAGraphMode.NONE])

149 +def test_draft_layer_rows_mirrored_after_drafter(monkeypatch, cg_mode):

150 + """Draft-layer host rows must hold the drafter's writes before page release.

151 +

152 + The drafter writes its layers' rows after the target forward. Mirroring

153 + them at the end of the target forward (the only mirror a FULL graph gets)

154 + leaves pre-draft rows on the host, and the page is still reported clean.

155 + """

156 + block_size, width = 4, 8

157 + resident_block, host_block, transfer_id = 2, 1, 7

158 + log: list[tuple[str, object]] = []

159 + _FakeEvent.log = log

160 +

161 + def swap_blocks_batch(src, dst, sizes):

162 + for src_ptr, dst_ptr, size in zip(src.tolist(), dst.tolist(), sizes.tolist()):

163 + ctypes.memmove(dst_ptr, src_ptr, size)

164 + log.append(("dma", dst_ptr))

165 +

166 + monkeypatch.setattr(torch, "Event", _FakeEvent)

167 + monkeypatch.setattr(torch, "Stream", lambda *args, **kwargs: MagicMock())

168 + monkeypatch.setattr(torch.cuda, "Stream", lambda *args, **kwargs: MagicMock())

169 + monkeypatch.setattr(torch.cuda, "stream", lambda _: contextlib.nullcontext())

170 + compute_stream = MagicMock()

171 + monkeypatch.setattr(worker_module, "current_stream", lambda: compute_stream)

172 + monkeypatch.setattr(worker_module.ops, "swap_blocks_batch", swap_blocks_batch)

173 + monkeypatch.setattr(

174 + runtime_module,

175 + "get_forward_context",

176 + lambda: SimpleNamespace(cudagraph_runtime_mode=cg_mode),

177 + )

178 +

179 + slots = torch.arange(block_size) + resident_block * block_size

180 + request_state_indices = torch.full((1,), -1, dtype=torch.int32)

181 + handles = []

182 + for draft_layer in (False, True):

183 + runtime = SimpleNamespace(

184 + is_group_leader=True,

185 + eager_host_mirror=True,

186 + resident_source_index=0,

187 + host_cache=torch.zeros(4 * block_size, width),

188 + request_state_indices=request_state_indices,

189 + begin_forward=MagicMock(),

190 + invalidate_written_slots=MagicMock(),

191 + )

192 + handle = HiSparseCacheHandle(runtime)

193 + handle.view = SimpleNamespace(

194 + block_size=block_size, cache=torch.zeros(4, block_size, width)

195 + )

196 + handle.slot_mapping = slots

197 + handle.mirror_slot_mapping = slots

198 + handle.draft_layer = draft_layer

199 + handles.append(handle)

200 + target, draft = handles

201 + # Rows the draft layer held before this step's drafter ran.

202 + draft.view.cache[resident_block] = -1.0

203 +

204 + worker = HiSparseConnectorWorker(

205 + SimpleNamespace(

206 + scheduler_config=SimpleNamespace(max_num_batched_tokens=16),

207 + num_lookahead_tokens=3,

208 + ),

209 + MagicMock(),

210 + )

211 + worker.initialize(

212 + handles,

213 + ["target", "draft"],

214 + MagicMock(),

215 + max_num_reqs=1,

216 + host_num_blocks=4,

217 + device=torch.device("cpu"),

218 + pinned_host_pools=[],

219 + )

220 + connector = object.__new__(HiSparseConnector)

221 + connector.connector_worker = worker

222 + connector._get_connector_metadata = MagicMock(

223 + return_value=HiSparseConnectorMetadata(

224 + SparseKVOffloadCommand(

225 + [

226 + SparseKVPageTransfer(

227 + transfer_id=transfer_id,

228 + host_block_id=host_block,

229 + resident_block_ids=(resident_block,),

230 + after_forward=True,

231 + )

232 + ]

233 + ),

234 + (),

235 + (),

236 + {

237 + "req": (

238 + SparseKVRowMirror(

239 + source_starts=(resident_block * block_size,),

240 + destination_start=host_block * block_size,

241 + num_rows=block_size,

242 + ),

243 + )

244 + },

245 + True,

246 + )

247 + )

248 + # A verification step: four query tokens of one request fill the page.

249 + attn_metadata = SimpleNamespace(

250 + num_actual_tokens=block_size,

251 + num_decode_tokens=0,

252 + num_reqs=1,

253 + max_query_len=block_size,

254 + req_id_per_token=torch.zeros(block_size, dtype=torch.int32),

255 + )

256 +

257 + connector.start_load_kv(

258 + SimpleNamespace(),

259 + request_state_indices=None,

260 + request_ids=["req"],

261 + num_tokens=block_size,

262 + attn_metadata={"target": attn_metadata, "draft": attn_metadata},

263 + )

264 + target.view.cache[resident_block] = 1.0

265 + if cg_mode != CUDAGraphMode.FULL:

266 + target.finish_kv_update()

267 + connector.finish_forward()

268 + enqueued_before_draft = list(worker._enqueued_transfer_ids)

269 +

270 + # The drafter writes the draft layer's rows for the same positions.

271 + if cg_mode != CUDAGraphMode.FULL:

272 + draft.prepare_group_for_batch(attn_metadata)

273 + draft.view.cache[resident_block] = 2.0

274 + if cg_mode != CUDAGraphMode.FULL:

275 + draft.finish_kv_update()

276 + # The next step's start mirrors the drafter's rows, then hands the page over.

277 + compute_stream.wait_event.reset_mock()

278 + connector._get_connector_metadata.return_value = HiSparseConnectorMetadata(

279 + None, (), (), {}, True

280 + )

281 + connector.start_load_kv(

282 + SimpleNamespace(),

283 + request_state_indices=None,

284 + request_ids=[],

285 + num_tokens=0,

286 + attn_metadata={},

287 + )

288 + ((completion_event, _),) = worker._pending_transfer_events

289 + worker_meta = connector.build_connector_worker_meta()

290 +

291 + host_rows = slice(host_block * block_size, (host_block + 1) * block_size)

292 + torch.testing.assert_close(target.runtime.host_cache[host_rows], torch.ones(4, 8))

293 + torch.testing.assert_close(

294 + draft.runtime.host_cache[host_rows], torch.full((4, 8), 2.0)

295 + )

296 + draft_host = draft.runtime.host_cache

297 + draft_host_range = range(

298 + draft_host.data_ptr(),

299 + draft_host.data_ptr() + draft_host.numel() * draft_host.element_size(),

300 + )

301 + draft_copies = [

302 + index

303 + for index, (kind, ptr) in enumerate(log)

304 + if kind == "dma" and ptr in draft_host_range

305 + ]

306 + # One draft-layer copy, and the page is not handed back before it.

307 + assert len(draft_copies) == 1

308 + assert enqueued_before_draft == []

309 + assert log.index(("event", completion_event)) > draft_copies[0]

310 + # The next forward waits for the draft copy: the scheduler may already

311 + # have handed the blocks it reads and writes to another request.

312 + draft_copy_event = next(

313 + entry[1] for entry in log[draft_copies[0] :] if entry[0] == "event"

314 + )

315 + assert compute_stream.wait_event.call_args_list[0].args == (draft_copy_event,)

316 + assert worker_meta is not None

317 + assert worker_meta.enqueued_transfer_counts == {transfer_id: 1}

318 + assert worker_meta.completed_transfer_counts == {transfer_id: 1}

319 +

320 +

321 def test_scheduled_prefix_hit_publishes_adopted_copies():

322 """Copies adopted after scheduling must reach the worker's block table."""

323 from tests.v1.core.test_prefix_caching import (

tests/v1/worker/test_utils.py

+2 -24

@@ -44,6 +44,8 @@ def _make_hisparse_worker() -> HiSparseConnectorWorker:

44 worker._row_mirror_num_rows = 0

45 worker._per_layer_mirrored = set()

46 worker._submitted_mirror_layers = set()

47 + worker._draft_layers = ()

48 + worker._draft_mirror_pending = False

49 worker._pending_dma_descriptors = deque()

50 worker._dma_free_descriptors = []

51 worker.host_write_events = (MagicMock(), MagicMock())

@@ -785,30 +787,6 @@ def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors():

787 worker._enqueue_row_dma.assert_not_called()

788

789

788 -def test_hisparse_finish_forward_submits_lazy_post_forward_transfer(monkeypatch):

789 - runtime = SimpleNamespace(eager_host_mirror=False)

790 - worker = _make_hisparse_worker()

791 - worker.is_host_writer = True

792 - worker.cache_handles = [SimpleNamespace(runtime=runtime, num_actual_tokens=0)]

793 - transfer = SparseKVPageTransfer(7, 2, (1,), after_forward=True)

794 - worker._post_forward_transfers = [transfer]

795 - worker._forward_ready_event = MagicMock()

796 - worker._enqueue_host_mirror = MagicMock()

797 - worker._submit_transfers = MagicMock()

798 - worker._dma_submitted = False

799 - worker._submitted_mirror_layers = set()

800 - worker._finish_mirror_phase = MagicMock()

801 - worker._release_completed_dma_descriptors = MagicMock()

802 - stream = MagicMock()

803 - monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream)

804 -

805 - worker.finish_forward()

806 -

807 - worker._finish_mirror_phase.assert_called_once_with(worker._forward_ready_event)

808 - worker._submit_transfers.assert_called_once_with([transfer])

809 - worker.host_write_event.record.assert_called_once_with(stream)

810 -

811 -

790 def test_hisparse_prefill_mirrors_source_groups_and_flushes_partial_group():

791 slots = torch.tensor([7, 8], dtype=torch.int64)

792 source_indices = [0, 0, 1, 1, 1, 1, 2]

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

+35 -8

@@ -301,6 +301,10 @@ class HiSparseConnectorWorker:

301 self._submitted_mirror_layers: set[int] = set()

302 self._layer_ready_events = tuple(torch.Event() for _ in cache_handles)

303 self._forward_ready_event = torch.Event()

304 + self._draft_layers = tuple(

305 + index for index, cache in enumerate(cache_handles) if cache.draft_layer

306 + )

307 + self._draft_mirror_pending = False

308 self._set_row_mirrors(())

309 self.cache_layer_names = cache_layer_names

310 self._group_leaders = tuple(

@@ -344,9 +348,11 @@ class HiSparseConnectorWorker:

348 self._init_dma()

349 if self.is_host_writer:

350 for layer_index, handle in enumerate(cache_handles):

347 - handle.submit_layer_mirror = partial(

348 - self._enqueue_layer_mirror, layer_index

349 - )

351 + # Draft layers are mirrored once in the next start_step instead.

352 + if not handle.draft_layer:

353 + handle.submit_layer_mirror = partial(

354 + self._enqueue_layer_mirror, layer_index

355 + )

356 self._initialized = True

357

358 def set_request_state_indices(self, indices: torch.Tensor) -> None:

@@ -397,16 +403,17 @@ class HiSparseConnectorWorker:

403 num_tokens: int = 0,

404 ) -> None:

405 self._stage_row_mirror_mapping(num_tokens)

406 + self._finish_previous_step()

407 previous_host_write_event = self.host_write_event

408 self.host_write_event = self.host_write_events[self._next_host_write_event]

409 self._next_host_write_event ^= 1

410 current_stream().wait_event(previous_host_write_event)

411 self._release_completed_dma_descriptors()

412 + self._dma_submitted = False

413 mirrors = _flatten_row_mirrors(metadata.row_mirrors, request_ids)

414 if self._slot_mapping_staging is not None:

415 self._slot_mapping_staging.candidates = mirrors

416 self._set_row_mirrors(mirrors)

409 - self._dma_submitted = False

417 self._clear_forward_mirror_state()

418 for handle in self.cache_handles:

419 handle.all_context_pages_resident = metadata.all_context_pages_resident

@@ -849,7 +856,12 @@ class HiSparseConnectorWorker:

856 if num_rows == 0:

857 return

858 if self.is_host_writer:

852 - expected_layers = {index for index, _ in active}

859 + # The drafter has not written its layers' rows yet; the next

860 + # start_step mirrors them after it runs.

861 + self._draft_mirror_pending = bool(self._draft_layers)

862 + expected_layers = {index for index, _ in active}.difference(

863 + self._draft_layers

864 + )

865 if self._per_layer_mirrored and self._per_layer_mirrored != expected_layers:

866 raise RuntimeError(

867 "HiSparse per-layer DMA did not mirror every active layer: "

@@ -884,9 +896,6 @@ class HiSparseConnectorWorker:

896 compute_stream = current_stream()

897 self._forward_ready_event.record()

898 self._finish_mirror_phase(self._forward_ready_event)

887 - transfers = self._post_forward_transfers

888 - self._post_forward_transfers = []

889 - self._submit_transfers(transfers)

899 if self.is_host_writer:

900 if self._dma_submitted:

901 compute_stream.wait_event(self.host_write_event)

@@ -895,6 +904,24 @@ class HiSparseConnectorWorker:

904 self.host_write_event.record(compute_stream)

905 self._release_completed_dma_descriptors()

906

907 + def _finish_previous_step(self) -> None:

908 + """Mirror the previous step's draft-layer rows, then hand its pages over.

909 +

910 + Runs before this step's work, so after the previous step's drafter.

911 + Post-forward transfers are ordered behind the draft mirror, so a page is

912 + never reported clean, and its resident block never released, before

913 + its draft-layer rows reach the host. Both land in the previous step's

914 + host write event, which this step waits on before its forward and its

915 + host copies, since the scheduler may already have reused the

916 + finished requests' blocks they read and write.

917 + """

918 + if self._draft_mirror_pending:

919 + self._draft_mirror_pending = False

920 + self._enqueue_row_dma(self._draft_layers)

921 + transfers = self._post_forward_transfers

922 + self._post_forward_transfers = []

923 + self._submit_transfers(transfers)

924 +

925 def take_completed_host_copies(self) -> list[int]:

926 """Drain host copies this worker has enqueued for this step."""

927 completed = self._completed_host_copy_dst_ids

vllm/model_executor/layers/attention_layer_base.py

@@ -21,6 +21,8 @@ class AttentionLayerBase(ABC):

22 impl: "AttentionImpl"

23 supports_dcp: bool = True

24 + # Set by the speculator on its draft model's layers.

25 + is_draft_layer: bool = False

27 def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:

28 """Bind the allocated KV cache tensor to this layer.

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

+2 -20

@@ -39,8 +39,6 @@ class SparseMLAIndexGroup:

39 logical_topk_indices: torch.Tensor

40 physical_topk_indices: torch.Tensor

41 valid_topk_counts: torch.Tensor

42 - row_indices: torch.Tensor

43 - request_ids: torch.Tensor

42 side_stream: torch.Stream

43 logical_topk_ready: torch.Event

44 physical_topk_ready: torch.Event

@@ -252,19 +250,13 @@ class HiSparseMLAIndexGroup(SparseMLAIndexGroup):

250 )

251 source_block_table = cache.source_block_table

252 assert source_block_table is not None

255 - # CUDA-graph padding rows past the batch's tokens map to request 0;

256 - # mark them -1 so residency resolution skips them.

257 - request_ids = self.request_ids[:num_tokens]

258 - request_ids.copy_(req_id_per_token)

259 - request_ids.masked_fill_(

260 - self.row_indices[:num_tokens] >= attn_metadata.query_start_loc[-1], -1

261 - )

253 return cache.swap_in(

263 - request_ids,

254 + req_id_per_token,

255 block_table=source_block_table,

256 logical_topk_indices=logical_topk_indices,

257 block_size=attn_metadata.block_size,

258 return_valid_counts=return_valid_counts,

259 + num_valid_rows=attn_metadata.query_start_loc[-1:],

260 )

261

262 def stage_prefill_rows(

@@ -384,16 +376,6 @@ class SparseMLAIndexGroupBuilder:

376 dtype=torch.int32,

377 device=self.logical_topk_indices.device,

378 ),

387 - row_indices=torch.arange(

388 - workspace_rows,

389 - dtype=torch.int32,

390 - device=self.logical_topk_indices.device,

391 - ),

392 - request_ids=torch.empty(

393 - workspace_rows,

394 - dtype=torch.int32,

395 - device=self.logical_topk_indices.device,

396 - ),

379 side_stream=_create_side_stream(self.logical_topk_indices.device),

380 logical_topk_ready=_create_event(),

381 physical_topk_ready=_create_event(),

vllm/v1/hisparse/binding.py

@@ -284,6 +284,7 @@ def bind_hisparse_kv_caches(

284 assert source_cache.untyped_storage().data_ptr() == (

285 host_pool.registered.untyped_storage().data_ptr()

286 )

287 + cache_handle.draft_layer = forward_context[layer_name].is_draft_layer

288 cache_handle.runtime.shared_host_region = host_pool.shared_region

289 cache_handle.runtime.bind_source_cache(

290 source_cache,

vllm/v1/hisparse/runtime.py

+20 -1

@@ -662,6 +662,10 @@ class HiSparseIndexGroup:

662 self.lru_init = torch.arange(region_stride, dtype=torch.int16, device=device)

663 self.lru_slots = self.lru_init.repeat(max_num_reqs, 1).contiguous()

664 self.shared_topk = _create_shared_topk_state(device, max_swap_rows, top_k)

665 + # Request ids the resolver reads, built on the copy stream: a

666 + # compute-stream write would race with its logical_topk_ready wait.

667 + self.row_indices = torch.arange(max_swap_rows, dtype=torch.int32, device=device)

668 + self.request_ids = torch.empty_like(self.row_indices)

669 self.copy_stream = (

670 copy_stream if copy_stream is not None else _create_copy_stream(device)

671 )

@@ -997,6 +1001,7 @@ class HiSparseRuntime:

1001 shared_rows: slice,

1002 attention_indices_out: torch.Tensor | None,

1003 valid_counts_out: torch.Tensor | None,

1004 + num_valid_rows: torch.Tensor,

1005 ) -> None:

1006 group = self.index_group

1007 compute_stream = current_stream()

@@ -1005,9 +1010,17 @@ class HiSparseRuntime:

1010 else:

1011 group.copy_stream.wait_stream(compute_stream)

1012 with group.copy_stream:

1013 + # CUDA-graph padding rows past the batch's tokens map to request 0;

1014 + # mark them -1 so residency resolution skips them.

1015 + num_tokens = logical_topk_indices.shape[0]

1016 + request_ids = group.request_ids[:num_tokens]

1017 + request_ids.copy_(req_id_per_token)

1018 + request_ids.masked_fill_(

1019 + group.row_indices[:num_tokens] >= num_valid_rows, -1

1020 + )

1021 self._resolve_residency(

1022 resident=resident,

1010 - req_id_per_token=req_id_per_token,

1023 + req_id_per_token=request_ids,

1024 block_table=block_table,

1025 topk_indices=logical_topk_indices,

1026 block_size=block_size,

@@ -1051,6 +1064,7 @@ class HiSparseRuntime:

1064 block_table: torch.Tensor,

1065 logical_topk_indices: torch.Tensor,

1066 block_size: int,

1067 + num_valid_rows: torch.Tensor,

1068 return_valid_counts: bool = False,

1069 attention_indices_out: torch.Tensor | None = None,

1070 valid_counts_out: torch.Tensor | None = None,

@@ -1074,6 +1088,7 @@ class HiSparseRuntime:

1088 shared_rows=shared_rows,

1089 attention_indices_out=attention_indices_out,

1090 valid_counts_out=valid_counts_out,

1091 + num_valid_rows=num_valid_rows,

1092 )

1093

1094 if not self._swap_staged:

@@ -1110,6 +1125,8 @@ class HiSparseCacheHandle:

1125 self.mirror_staging_cache: torch.Tensor | None = None

1126 self.mirror_staging_slots: torch.Tensor | None = None

1127 self.submit_layer_mirror: Callable[[], None] | None = None

1128 + # Speculator layers write their rows after the target forward.

1129 + self.draft_layer = False

1130 self.index_group_caches: list[HiSparseCacheHandle] = [self]

1131

1132 def prepare_group_for_batch(self, attn_metadata: Any | None) -> None:

@@ -1203,6 +1220,7 @@ class HiSparseCacheHandle:

1220 logical_topk_indices: torch.Tensor,

1221 *,

1222 block_size: int,

1223 + num_valid_rows: torch.Tensor,

1224 return_valid_counts: bool = False,

1225 attention_indices_out: torch.Tensor | None = None,

1226 valid_counts_out: torch.Tensor | None = None,

@@ -1216,6 +1234,7 @@ class HiSparseCacheHandle:

1234 return_valid_counts=return_valid_counts,

1235 attention_indices_out=attention_indices_out,

1236 valid_counts_out=valid_counts_out,

1237 + num_valid_rows=num_valid_rows,

1238 )

1239

1240

vllm/v1/worker/gpu/spec_decode/speculator.py

+6 -6

@@ -224,13 +224,13 @@ class DraftModelSpeculator(BaseSpeculator):

224 self.model = self.load_draft_model(target_model, target_attn_layer_names)

225 self._validate_local_argmax_reduction()

226

227 - all_attn_layers = set[str](

228 - get_layers_from_vllm_config(

229 - self.vllm_config,

230 - AttentionLayerBase, # type: ignore[type-abstract]

231 - ).keys()

227 + all_attn_layers = get_layers_from_vllm_config(

228 + self.vllm_config,

229 + AttentionLayerBase, # type: ignore[type-abstract]

230 )

233 - self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names

231 + self.draft_attn_layer_names = set(all_attn_layers) - target_attn_layer_names

232 + for layer_name in self.draft_attn_layer_names:

233 + all_attn_layers[layer_name].is_draft_layer = True

234

235 target_supports_mm = self.vllm_config.model_config.supports_multimodal_inputs

236 draft_supports_mm = supports_multimodal_embeddings(self.model)