| 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 |
| 7 | |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 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, |
| 22 | HiSparseConnectorScheduler, |
| 23 | ) |
| 24 | 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 |
| 37 | |
| 38 | |
| 39 | @pytest.mark.parametrize("nested", [False, True]) |
| 40 | def test_cache_manager_binding_preserves_hisparse_and_legacy_pool_hooks(nested): |
| 41 | """Composite binding must reach HiSparse and existing pool-only connectors.""" |
| 42 | hisparse = object.__new__(HiSparseConnector) |
| 43 | hisparse._role = KVConnectorRole.SCHEDULER |
| 44 | hisparse.connector_scheduler = HiSparseConnectorScheduler(async_speculative=False) |
| 45 | legacy = object.__new__(ExampleConnector) |
| 46 | legacy.bind_gpu_block_pool = MagicMock() |
| 47 | connector = object.__new__(MultiConnector) |
| 48 | connector._connectors = [hisparse, legacy] |
| 49 | if nested: |
| 50 | parent = object.__new__(MultiConnector) |
| 51 | parent._connectors = [connector] |
| 52 | connector = parent |
| 53 | from tests.v1.core.test_prefix_caching import make_hisparse_kv_cache_manager |
| 54 | |
| 55 | manager = make_hisparse_kv_cache_manager(16, 16) |
| 56 | |
| 57 | connector.bind_kv_cache_manager(manager) |
| 58 | |
| 59 | assert hisparse.connector_scheduler.coordinator is get_hisparse_coordinator(manager) |
| 60 | legacy.bind_gpu_block_pool.assert_called_once_with(manager.block_pool) |
| 61 | |
| 62 | |
| 63 | def test_hisparse_requires_block_outermost_device_layout(): |
| 64 | assert HiSparseConnector.get_required_kvcache_layout(MagicMock()) == "BLHNC" |
| 65 | |
| 66 | |
| 67 | def test_no_forward_enqueues_deferred_hisparse_transfers(): |
| 68 | """A zero-token step must still enqueue deferred post-forward transfers.""" |
| 69 | connector = object.__new__(ActiveKVConnector) |
| 70 | connector._disabled = False |
| 71 | connector.pre_forward = MagicMock() |
| 72 | connector.finish_forward = MagicMock() |
| 73 | connector.post_forward = MagicMock(return_value=None) |
| 74 | |
| 75 | scheduler_output = SimpleNamespace(finished_req_ids=set()) |
| 76 | connector.no_forward(scheduler_output) |
| 77 | |
| 78 | connector.pre_forward.assert_called_once_with(scheduler_output) |
| 79 | connector.finish_forward.assert_called_once_with() |
| 80 | |
| 81 | |
| 82 | def test_full_graph_step_prepares_host_mirror_outside_model(): |
| 83 | """Graph replay must restore host-mirror state cleared at step start.""" |
| 84 | runtime = SimpleNamespace( |
| 85 | is_group_leader=True, |
| 86 | eager_host_mirror=True, |
| 87 | begin_forward=MagicMock(), |
| 88 | invalidate_written_slots=MagicMock(), |
| 89 | ) |
| 90 | handle = HiSparseCacheHandle(runtime) |
| 91 | handle.mirror_slot_mapping = torch.tensor([4, 5]) |
| 92 | worker = object.__new__(HiSparseConnectorWorker) |
| 93 | worker.cache_layer_names = ["layer"] |
| 94 | worker.cache_handles = [handle] |
| 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( |
| 102 | side_effect=lambda *_args, **_kwargs: worker._clear_forward_mirror_state() |
| 103 | ) |
| 104 | |
| 105 | connector = object.__new__(HiSparseConnector) |
| 106 | connector.connector_worker = worker |
| 107 | connector._get_connector_metadata = MagicMock( |
| 108 | return_value=HiSparseConnectorMetadata(None, (), (), {}, True) |
| 109 | ) |
| 110 | req_id_per_token = torch.tensor([0, 1], dtype=torch.int32) |
| 111 | attn_metadata = SimpleNamespace( |
| 112 | num_actual_tokens=2, |
| 113 | num_decode_tokens=2, |
| 114 | num_reqs=2, |
| 115 | max_query_len=1, |
| 116 | req_id_per_token=req_id_per_token, |
| 117 | ) |
| 118 | |
| 119 | connector.start_load_kv( |
| 120 | SimpleNamespace(), |
| 121 | request_state_indices=None, |
| 122 | request_ids=[], |
| 123 | attn_metadata={"layer": attn_metadata}, |
| 124 | ) |
| 125 | |
| 126 | worker._enqueue_host_mirror() |
| 127 | |
| 128 | worker._enqueue_row_dma.assert_called_once_with((0,), ready_event=None) |
| 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 ( |
| 324 | HISPARSE_BLOCK_SIZE, |
| 325 | _allocate_scheduled, |
| 326 | _publish_hisparse_pages, |
| 327 | make_hisparse_kv_cache_manager, |
| 328 | make_request, |
| 329 | sha256, |
| 330 | ) |
| 331 | from vllm.v1.core.kv_cache_utils import init_none_hash |
| 332 | |
| 333 | init_none_hash(sha256) |
| 334 | manager = make_hisparse_kv_cache_manager(32, 16, enable_caching=True) |
| 335 | tokens = list(range(4 * HISPARSE_BLOCK_SIZE)) |
| 336 | original = make_request("original", tokens, HISPARSE_BLOCK_SIZE, sha256) |
| 337 | assert _allocate_scheduled(manager, original, num_new_tokens=len(tokens)) |
| 338 | _publish_hisparse_pages(manager) |
| 339 | copy_ids = [block.block_id for block in manager.get_blocks("original").blocks[2]] |
| 340 | manager.free(original) |
| 341 | |
| 342 | resumed = make_request("resumed", tokens, HISPARSE_BLOCK_SIZE, sha256) |
| 343 | computed, num_computed, _ = manager.get_computed_blocks(resumed) |
| 344 | num_new_tokens = len(tokens) - num_computed |
| 345 | assert manager.allocate_slots( |
| 346 | resumed, |
| 347 | num_new_tokens=num_new_tokens, |
| 348 | num_new_computed_tokens=num_computed, |
| 349 | new_computed_blocks=computed, |
| 350 | ) |
| 351 | scheduler = HiSparseConnectorScheduler(async_speculative=False) |
| 352 | scheduler.bind_coordinator(get_hisparse_coordinator(manager)) |
| 353 | scheduler.requests[resumed.request_id] = resumed |
| 354 | scheduler_output = SimpleNamespace( |
| 355 | scheduled_new_reqs=[ |
| 356 | SimpleNamespace( |
| 357 | req_id=resumed.request_id, |
| 358 | num_computed_tokens=num_computed, |
| 359 | block_ids=manager.get_block_ids(resumed.request_id), |
| 360 | ) |
| 361 | ], |
| 362 | scheduled_cached_reqs=SimpleNamespace( |
| 363 | req_ids=[], num_computed_tokens=[], new_block_ids=[] |
| 364 | ), |
| 365 | num_scheduled_tokens={resumed.request_id: num_new_tokens}, |
| 366 | ) |
| 367 | |
| 368 | scheduler.build_connector_meta(scheduler_output) |
| 369 | |
| 370 | resident_ids = scheduler_output.block_table_updates[resumed.request_id][2] |
| 371 | assert resident_ids[:3] == copy_ids[:3] |