| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import mmap |
| 5 | from collections import deque |
| 6 | from contextlib import nullcontext |
| 7 | from types import SimpleNamespace |
| 8 | from unittest.mock import MagicMock, call |
| 9 | |
| 10 | import numpy as np |
| 11 | import pytest |
| 12 | import torch |
| 13 | |
| 14 | import vllm.v1.attention.backends.mla.index_group as index_group_module |
| 15 | import vllm.v1.hisparse.runtime as hisparse_runtime_module |
| 16 | from vllm.config import CacheConfig, CUDAGraphMode |
| 17 | from vllm.config.mamba import MambaBackendEnum, MambaConfig |
| 18 | from vllm.distributed.kv_transfer.kv_connector.v1.hisparse import ( |
| 19 | worker as hisparse_worker_module, |
| 20 | ) |
| 21 | from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( |
| 22 | HiSparseConnectorWorker, |
| 23 | _flatten_row_mirrors, |
| 24 | _select_written_row_mirrors, |
| 25 | _SlotMappingStaging, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 |
| 28 | from vllm.utils.mem_constants import GiB_bytes |
| 29 | from vllm.utils.mem_utils import MemorySnapshot |
| 30 | from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy |
| 31 | from vllm.v1.hisparse.types import SparseKVPageTransfer, SparseKVRowMirror |
| 32 | from vllm.v1.worker.utils import ( |
| 33 | bind_kv_cache, |
| 34 | bind_kv_cache_to_layers, |
| 35 | copy_kv_cache_blocks_inplace, |
| 36 | request_memory, |
| 37 | ) |
| 38 | |
| 39 | |
| 40 | def _make_hisparse_worker() -> HiSparseConnectorWorker: |
| 41 | worker = object.__new__(HiSparseConnectorWorker) |
| 42 | worker.refines_row_mirrors = False |
| 43 | worker._slot_mapping_staging = None |
| 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()) |
| 52 | worker.host_write_event = worker.host_write_events[1] |
| 53 | worker._next_host_write_event = 0 |
| 54 | worker.dma_stream = None |
| 55 | worker.shared_host_region = None |
| 56 | worker._completed_host_copy_dst_ids = [] |
| 57 | worker._metrics_calls = 0 |
| 58 | worker._metrics_event = MagicMock() |
| 59 | worker._metrics_pending = False |
| 60 | worker.leader_runtimes = [] |
| 61 | return worker |
| 62 | |
| 63 | |
| 64 | def test_hisparse_worker_get_kv_connector_stats_reads_completed_snapshot(monkeypatch): |
| 65 | worker = _make_hisparse_worker() |
| 66 | worker._metrics_calls = hisparse_worker_module._METRICS_INTERVAL - 1 |
| 67 | worker._metrics_pending = False |
| 68 | worker._metrics_event = MagicMock() |
| 69 | worker._metrics_event.query.return_value = True |
| 70 | compute_stream = MagicMock() |
| 71 | group = SimpleNamespace( |
| 72 | swap_stats=torch.tensor([12, 4], dtype=torch.uint64), |
| 73 | swap_stats_host=torch.empty(2, dtype=torch.uint64), |
| 74 | stats_row_bytes=16, |
| 75 | copy_stream=MagicMock(), |
| 76 | ) |
| 77 | worker.leader_runtimes = [SimpleNamespace(index_group=group)] |
| 78 | monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) |
| 79 | monkeypatch.setattr( |
| 80 | hisparse_worker_module, "current_stream", lambda: compute_stream |
| 81 | ) |
| 82 | |
| 83 | assert worker.get_kv_connector_stats() is None |
| 84 | compute_stream.wait_stream.assert_called_once_with(group.copy_stream) |
| 85 | group.copy_stream.wait_stream.assert_called_once_with(compute_stream) |
| 86 | worker._metrics_event.record.assert_called_once_with() |
| 87 | assert group.swap_stats.tolist() == [0, 0] |
| 88 | |
| 89 | stats = worker.get_kv_connector_stats() |
| 90 | assert stats is not None |
| 91 | assert stats.data == { |
| 92 | "cache_hits": [12], |
| 93 | "cache_misses": [4], |
| 94 | "host_to_device_bytes": [64], |
| 95 | } |
| 96 | |
| 97 | |
| 98 | def test_hisparse_row_mirrors_follow_runner_request_order(): |
| 99 | first = SparseKVRowMirror((1,), 10, 1) |
| 100 | second = SparseKVRowMirror((2,), 20, 1) |
| 101 | row_mirrors = {"first": (first,), "second": (second,)} |
| 102 | |
| 103 | assert _flatten_row_mirrors(row_mirrors, ["second", "first"]) == ( |
| 104 | second, |
| 105 | first, |
| 106 | ) |
| 107 | |
| 108 | |
| 109 | def test_hisparse_gpu_slots_select_written_rows_from_async_envelope(): |
| 110 | """Rejected speculative rows must never become durable host KV.""" |
| 111 | candidates = ( |
| 112 | SparseKVRowMirror((100, 500), 1000, 4), |
| 113 | SparseKVRowMirror((20, 300), 2000, 4), |
| 114 | ) |
| 115 | source_slots = np.array([102, 103, -1, 21, 22, 23], dtype=np.int64) |
| 116 | |
| 117 | assert _select_written_row_mirrors(candidates, source_slots, 0) == ( |
| 118 | SparseKVRowMirror((102, 502), 1002, 2), |
| 119 | SparseKVRowMirror((21, 301), 2001, 3), |
| 120 | ) |
| 121 | |
| 122 | |
| 123 | def test_hisparse_written_rows_preserve_scheduler_page_boundaries(): |
| 124 | candidates = ( |
| 125 | SparseKVRowMirror((0,), 100, 4), |
| 126 | SparseKVRowMirror((4,), 104, 4), |
| 127 | ) |
| 128 | |
| 129 | assert _select_written_row_mirrors( |
| 130 | candidates, np.array([2, 3, 4, 5], dtype=np.int64), 0 |
| 131 | ) == ( |
| 132 | SparseKVRowMirror((2,), 102, 2), |
| 133 | SparseKVRowMirror((4,), 104, 2), |
| 134 | ) |
| 135 | |
| 136 | |
| 137 | def test_hisparse_written_rows_without_host_destinations_are_skipped(): |
| 138 | assert ( |
| 139 | _select_written_row_mirrors((), np.array([100, 101], dtype=np.int64), 0) == () |
| 140 | ) |
| 141 | |
| 142 | |
| 143 | def test_hisparse_appends_reference_slots_within_a_mirror_phase(monkeypatch): |
| 144 | """Separate context and query writes must share one mirror phase.""" |
| 145 | worker = _make_hisparse_worker() |
| 146 | worker.is_host_writer = True |
| 147 | worker.refines_row_mirrors = True |
| 148 | state = _SlotMappingStaging( |
| 149 | MagicMock(), MagicMock(), torch.empty(4, dtype=torch.int64) |
| 150 | ) |
| 151 | worker._slot_mapping_staging = state |
| 152 | handle = SimpleNamespace(runtime=SimpleNamespace(resident_source_index=1)) |
| 153 | worker.cache_layer_names = ["layer"] |
| 154 | worker.cache_handles = [handle] |
| 155 | worker._layer_mirror_callbacks = (MagicMock(),) |
| 156 | worker._row_mirror_num_rows = 1 |
| 157 | worker._submitted_mirror_layers = set() |
| 158 | main_stream = MagicMock() |
| 159 | monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: main_stream) |
| 160 | monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) |
| 161 | |
| 162 | handle.slot_mapping = torch.tensor([7, 8], dtype=torch.int64) |
| 163 | worker._stage_row_mirror_mapping(2) |
| 164 | handle.slot_mapping = torch.tensor([9, 10], dtype=torch.int64) |
| 165 | worker._stage_row_mirror_mapping(2) |
| 166 | |
| 167 | torch.testing.assert_close(state.slots, torch.tensor([7, 8, 9, 10])) |
| 168 | assert state.stream.wait_stream.call_args_list == [call(main_stream)] * 2 |
| 169 | assert state.event.record.call_args_list == [call(state.stream)] * 2 |
| 170 | assert state.num_tokens == 4 |
| 171 | assert state.source_index == 1 |
| 172 | |
| 173 | |
| 174 | @pytest.mark.parametrize( |
| 175 | ("mode", "expected_calls"), |
| 176 | [(CUDAGraphMode.PIECEWISE, 1), (CUDAGraphMode.FULL, 0)], |
| 177 | ) |
| 178 | def test_hisparse_submits_layer_mirror_at_replayed_attention_boundary( |
| 179 | monkeypatch, mode, expected_calls |
| 180 | ): |
| 181 | """Piecewise replay must pipeline DMA; FULL replay uses the batch fallback.""" |
| 182 | handle = object.__new__(hisparse_runtime_module.HiSparseCacheHandle) |
| 183 | handle.dummy_batch = False |
| 184 | handle.decode_batch = False |
| 185 | handle.submit_layer_mirror = MagicMock() |
| 186 | monkeypatch.setattr( |
| 187 | hisparse_runtime_module, |
| 188 | "get_forward_context", |
| 189 | lambda: SimpleNamespace(cudagraph_runtime_mode=mode), |
| 190 | ) |
| 191 | |
| 192 | handle.finish_kv_update() |
| 193 | |
| 194 | assert handle.submit_layer_mirror.call_count == expected_calls |
| 195 | |
| 196 | |
| 197 | def _hisparse_parallel_config(**overrides): |
| 198 | values = { |
| 199 | "tensor_parallel_size": 2, |
| 200 | "pipeline_parallel_size": 1, |
| 201 | "prefill_context_parallel_size": 1, |
| 202 | "decode_context_parallel_size": 1, |
| 203 | "world_size": 2, |
| 204 | "distributed_executor_backend": "mp", |
| 205 | "nnodes_within_dp": 1, |
| 206 | "data_parallel_index": 3, |
| 207 | } |
| 208 | values.update(overrides) |
| 209 | return SimpleNamespace(**values) |
| 210 | |
| 211 | |
| 212 | def test_hisparse_shares_host_pool_only_for_local_tp(monkeypatch): |
| 213 | monkeypatch.setattr( |
| 214 | hisparse_runtime_module.current_platform, "is_cuda_alike", lambda: True |
| 215 | ) |
| 216 | config = SimpleNamespace(parallel_config=_hisparse_parallel_config()) |
| 217 | |
| 218 | assert hisparse_runtime_module.use_shared_hisparse_host_pool(config) |
| 219 | |
| 220 | unsupported = ( |
| 221 | {"tensor_parallel_size": 1, "world_size": 1}, |
| 222 | {"pipeline_parallel_size": 2, "world_size": 4}, |
| 223 | {"prefill_context_parallel_size": 2, "world_size": 4}, |
| 224 | {"decode_context_parallel_size": 2}, |
| 225 | {"world_size": 4}, |
| 226 | {"distributed_executor_backend": "ray"}, |
| 227 | {"nnodes_within_dp": 2}, |
| 228 | ) |
| 229 | for overrides in unsupported: |
| 230 | config.parallel_config = _hisparse_parallel_config(**overrides) |
| 231 | assert not hisparse_runtime_module.use_shared_hisparse_host_pool(config) |
| 232 | |
| 233 | |
| 234 | @pytest.mark.skip_global_cleanup |
| 235 | def test_hisparse_shared_host_pool_uses_one_replicated_mmap(monkeypatch): |
| 236 | page = mmap.PAGESIZE |
| 237 | tp_group = MagicMock() |
| 238 | monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", lambda: tp_group) |
| 239 | |
| 240 | class FakeSharedOffloadRegion: |
| 241 | BLOCK_SIZE_ALIGNMENT = page |
| 242 | |
| 243 | def __init__(self, **kwargs): |
| 244 | self.kwargs = kwargs |
| 245 | self.total_size_bytes = kwargs["num_chunks"] * kwargs["kv_bytes_per_chunk"] |
| 246 | self.base_tensor = torch.empty(self.total_size_bytes, dtype=torch.int8) |
| 247 | self.is_pinned = False |
| 248 | self.pinned_addresses = [] |
| 249 | self.view_sizes = [] |
| 250 | self.offset = 0 |
| 251 | |
| 252 | def create_next_canonical_view(self, size): |
| 253 | self.view_sizes.append(size) |
| 254 | view = torch.as_strided( |
| 255 | self.base_tensor, |
| 256 | size=(self.kwargs["num_chunks"], size), |
| 257 | stride=(self.kwargs["kv_bytes_per_chunk"], 1), |
| 258 | storage_offset=self.offset, |
| 259 | ) |
| 260 | self.offset += size |
| 261 | return view |
| 262 | |
| 263 | def cleanup(self): |
| 264 | pass |
| 265 | |
| 266 | monkeypatch.setattr( |
| 267 | hisparse_runtime_module, "SharedOffloadRegion", FakeSharedOffloadRegion |
| 268 | ) |
| 269 | registration_ranges = MagicMock(return_value=((0, page), (page, 4 * page))) |
| 270 | monkeypatch.setattr( |
| 271 | hisparse_runtime_module, |
| 272 | "_hisparse_registration_ranges", |
| 273 | registration_ranges, |
| 274 | ) |
| 275 | pinned: list[torch.Tensor] = [] |
| 276 | |
| 277 | def pin_tensor(tensor): |
| 278 | pinned.append(tensor) |
| 279 | |
| 280 | monkeypatch.setattr(hisparse_runtime_module, "pin_tensor", pin_tensor) |
| 281 | config = SimpleNamespace( |
| 282 | instance_id="instance", |
| 283 | parallel_config=_hisparse_parallel_config( |
| 284 | tensor_parallel_size=1, |
| 285 | world_size=1, |
| 286 | distributed_executor_backend="uni", |
| 287 | ), |
| 288 | ) |
| 289 | |
| 290 | pools, private_pools, region = hisparse_runtime_module.allocate_hisparse_host_pools( |
| 291 | config, |
| 292 | [24, 40], |
| 293 | num_blocks=4, |
| 294 | host_block_stride=page, |
| 295 | use_shared_host_pool=True, |
| 296 | ) |
| 297 | |
| 298 | assert region is not None |
| 299 | assert private_pools == [] |
| 300 | assert region.kwargs == { |
| 301 | "engine_id": "hisparse_instance_dp3", |
| 302 | "num_chunks": 1, |
| 303 | "rank": 0, |
| 304 | "kv_bytes_per_chunk": 4 * page, |
| 305 | "cpu_page_size": 64, |
| 306 | "barrier": tp_group.barrier, |
| 307 | "creator_memory_check": hisparse_runtime_module.check_hisparse_host_memory, |
| 308 | "populate_only_on_creator": True, |
| 309 | } |
| 310 | assert region.view_sizes == [24, 40] |
| 311 | assert [pool.shape for pool in pools] == [(24,), (40,)] |
| 312 | registration_ranges.assert_called_once_with([24, 40], 4, page) |
| 313 | assert [ |
| 314 | (tensor.data_ptr() - region.base_tensor.data_ptr(), tensor.nbytes) |
| 315 | for tensor in pinned |
| 316 | ] == [(0, page), (page, 3 * page)] |
| 317 | assert region.is_pinned |
| 318 | |
| 319 | |
| 320 | @pytest.mark.parametrize("fail_registration", [None, 1, 2]) |
| 321 | def test_shared_host_pool_tracks_successful_registrations( |
| 322 | monkeypatch, fail_registration |
| 323 | ): |
| 324 | """Only successful registrations may be unregistered on allocation failure.""" |
| 325 | page = mmap.PAGESIZE |
| 326 | monkeypatch.setattr(hisparse_runtime_module, "get_tp_group", MagicMock()) |
| 327 | backing = torch.empty(3 * page, dtype=torch.uint8) |
| 328 | region = MagicMock(base_tensor=backing, pinned_addresses=[], is_pinned=False) |
| 329 | monkeypatch.setattr( |
| 330 | hisparse_runtime_module, "SharedOffloadRegion", lambda **kw: region |
| 331 | ) |
| 332 | monkeypatch.setattr( |
| 333 | hisparse_runtime_module, |
| 334 | "_hisparse_registration_ranges", |
| 335 | lambda *args: ((0, page), (page, 3 * page)), |
| 336 | ) |
| 337 | cudart = MagicMock() |
| 338 | cudart.cudaHostRegister.side_effect = [ |
| 339 | SimpleNamespace(value=int(fail_registration == i)) for i in (1, 2) |
| 340 | ] |
| 341 | monkeypatch.setattr(torch.cuda, "cudart", lambda: cudart) |
| 342 | config = SimpleNamespace( |
| 343 | instance_id="test", parallel_config=SimpleNamespace(data_parallel_index=0) |
| 344 | ) |
| 345 | context = ( |
| 346 | pytest.raises(RuntimeError, match="cudaHostRegister failed") |
| 347 | if fail_registration |
| 348 | else nullcontext() |
| 349 | ) |
| 350 | with context: |
| 351 | hisparse_runtime_module.allocate_hisparse_host_pools( |
| 352 | config, [3 * page], 3, page, use_shared_host_pool=True |
| 353 | ) |
| 354 | successful = 2 if fail_registration is None else fail_registration - 1 |
| 355 | assert region.pinned_addresses == [ |
| 356 | backing.data_ptr() + i * page for i in range(successful) |
| 357 | ] |
| 358 | assert region.is_pinned == bool(successful) |
| 359 | assert region.cleanup.call_count == int(fail_registration is not None) |
| 360 | attempts = 2 if fail_registration is None else fail_registration |
| 361 | assert ( |
| 362 | cudart.cudaHostRegister.call_args_list |
| 363 | == [ |
| 364 | call(backing.data_ptr(), page, 0), |
| 365 | call(backing.data_ptr() + page, 2 * page, 0), |
| 366 | ][:attempts] |
| 367 | ) |
| 368 | |
| 369 | |
| 370 | def test_shared_host_pool_registers_layer_spans_in_one_backing(monkeypatch): |
| 371 | """Aliased tensor configs must not count the shared backing once per view.""" |
| 372 | region = SimpleNamespace(base_tensor=torch.empty(64, dtype=torch.int8)) |
| 373 | allocate = MagicMock(return_value=([], [], region)) |
| 374 | monkeypatch.setattr( |
| 375 | hisparse_runtime_module, "allocate_hisparse_host_pools", allocate |
| 376 | ) |
| 377 | config = SimpleNamespace( |
| 378 | hisparse_shared_host_pool=True, |
| 379 | hisparse_host_num_blocks=4, |
| 380 | hisparse_host_block_stride=16, |
| 381 | kv_cache_tensors=[ |
| 382 | SimpleNamespace( |
| 383 | host_resident=True, |
| 384 | layers=["a", "b"], |
| 385 | offset=0, |
| 386 | layer_stride=12, |
| 387 | size=44, |
| 388 | ), |
| 389 | SimpleNamespace( |
| 390 | host_resident=True, layers=["c"], offset=24, layer_stride=20, size=44 |
| 391 | ), |
| 392 | ], |
| 393 | ) |
| 394 | vllm_config = SimpleNamespace() |
| 395 | pool = hisparse_runtime_module.HiSparseHostPool(vllm_config, config) |
| 396 | backing = pool.allocate(44) |
| 397 | allocate.assert_called_once_with( |
| 398 | vllm_config, [12, 12, 20], 4, 16, use_shared_host_pool=True |
| 399 | ) |
| 400 | assert backing.numel() == 44 |
| 401 | assert backing.data_ptr() == region.base_tensor.data_ptr() |
| 402 | assert pool.registered is region.base_tensor |
| 403 | |
| 404 | |
| 405 | def test_hisparse_registration_chunks_end_between_host_blocks(): |
| 406 | """Registration seams must not bisect a DMA-addressable host block.""" |
| 407 | page = mmap.PAGESIZE |
| 408 | ranges = hisparse_runtime_module._hisparse_registration_ranges( |
| 409 | tensor_sizes=[4 * 3 * page, 4 * 5 * page], |
| 410 | num_blocks=4, |
| 411 | host_block_stride=8 * page, |
| 412 | max_chunk_bytes=10 * page, |
| 413 | ) |
| 414 | |
| 415 | assert ranges == ( |
| 416 | (0, 9 * page), |
| 417 | (9 * page, 17 * page), |
| 418 | (17 * page, 27 * page), |
| 419 | (27 * page, 32 * page), |
| 420 | ) |
| 421 | |
| 422 | |
| 423 | def test_copy_cpu_kv_cache_logical_blocks_ignores_storage_padding(): |
| 424 | backing = torch.full((10, 2, 3), -1, dtype=torch.float32) |
| 425 | cache = backing[1:9] |
| 426 | cache[2:4] = 7 |
| 427 | cache[6:8] = 11 |
| 428 | |
| 429 | copy_kv_cache_blocks_inplace( |
| 430 | [cache], |
| 431 | num_blocks=4, |
| 432 | kv_cache_block_copies=[ |
| 433 | KVCacheBlockCopy(1, 0), |
| 434 | KVCacheBlockCopy(3, 2), |
| 435 | ], |
| 436 | ) |
| 437 | |
| 438 | torch.testing.assert_close(cache[0:2], torch.full_like(cache[0:2], 7)) |
| 439 | torch.testing.assert_close(cache[4:6], torch.full_like(cache[4:6], 11)) |
| 440 | assert (backing[0] == -1).all() |
| 441 | assert (backing[9] == -1).all() |
| 442 | |
| 443 | |
| 444 | @pytest.mark.parametrize("shared,rank", [(False, 0), (True, 0), (True, 1)]) |
| 445 | @pytest.mark.parametrize("has_copies", [False, True]) |
| 446 | def test_hisparse_host_copy_waits_for_writes_on_writer_only( |
| 447 | monkeypatch, shared, rank, has_copies |
| 448 | ): |
| 449 | """Host CoW must read completed writes; shared non-writers only barrier.""" |
| 450 | worker = _make_hisparse_worker() |
| 451 | cache = torch.zeros((2, 1, 1)) |
| 452 | worker.host_caches = [cache] |
| 453 | worker.host_num_blocks = 2 |
| 454 | worker.shared_host_region = object() if shared else None |
| 455 | event = MagicMock() |
| 456 | event.synchronize.side_effect = lambda: cache[0].fill_(7) |
| 457 | barrier = MagicMock() |
| 458 | monkeypatch.setattr( |
| 459 | hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: rank |
| 460 | ) |
| 461 | monkeypatch.setattr( |
| 462 | hisparse_worker_module, "get_tp_group", lambda: SimpleNamespace(barrier=barrier) |
| 463 | ) |
| 464 | copies = [KVCacheBlockCopy(0, 1)] if has_copies else [] |
| 465 | |
| 466 | worker._copy_host_blocks(copies, event) |
| 467 | |
| 468 | writes = has_copies and (not shared or rank == 0) |
| 469 | assert cache[1].item() == (7 if writes else 0) |
| 470 | assert event.synchronize.call_count == int(writes) |
| 471 | assert barrier.call_count == int(shared and has_copies) |
| 472 | assert worker._completed_host_copy_dst_ids == ([1] if has_copies else []) |
| 473 | |
| 474 | |
| 475 | def test_hisparse_worker_updates_request_state_mapping_in_place(monkeypatch): |
| 476 | worker = _make_hisparse_worker() |
| 477 | worker.request_state_indices = torch.arange(4, dtype=torch.int32) |
| 478 | worker._pending_invalid_block_ids = [5] |
| 479 | invalidations = [] |
| 480 | worker.invalidate_blocks = lambda blocks, states: invalidations.append( |
| 481 | (blocks.copy(), states.clone()) |
| 482 | ) |
| 483 | original_ptr = worker.request_state_indices.data_ptr() |
| 484 | monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) |
| 485 | |
| 486 | worker.set_request_state_indices(torch.tensor([3, 1], dtype=torch.int32)) |
| 487 | |
| 488 | assert worker.request_state_indices.data_ptr() == original_ptr |
| 489 | assert worker.request_state_indices.tolist() == [3, 1, -1, -1] |
| 490 | assert len(invalidations) == 1 |
| 491 | assert invalidations[0][0] == [5] |
| 492 | torch.testing.assert_close( |
| 493 | invalidations[0][1], torch.tensor([3, 1], dtype=torch.int32) |
| 494 | ) |
| 495 | assert worker._pending_invalid_block_ids == [] |
| 496 | |
| 497 | |
| 498 | def test_hisparse_pre_forward_transfer_builds_page_descriptors(): |
| 499 | worker = _make_hisparse_worker() |
| 500 | worker.is_host_writer = True |
| 501 | worker.kernel_block_size = 2 |
| 502 | source = torch.empty((3, 2, 4), dtype=torch.uint8) |
| 503 | destination = torch.empty((6, 4), dtype=torch.uint8) |
| 504 | worker.resident_caches = (source,) |
| 505 | worker.host_caches = (destination,) |
| 506 | worker.cache_handles = [ |
| 507 | SimpleNamespace(runtime=SimpleNamespace(resident_source_index=0)) |
| 508 | ] |
| 509 | worker._dma_free_descriptors = [] |
| 510 | worker._submit_dma_descriptors = MagicMock() |
| 511 | |
| 512 | worker._enqueue_transfers([SparseKVPageTransfer(7, 2, (1,), after_forward=False)]) |
| 513 | |
| 514 | descriptors, count = worker._submit_dma_descriptors.call_args.args |
| 515 | assert count == 1 |
| 516 | assert descriptors.src[:count].tolist() == [ |
| 517 | source.data_ptr() + source.stride(0) * source.element_size() |
| 518 | ] |
| 519 | assert descriptors.dst[:count].tolist() == [destination.data_ptr() + 4 * 4] |
| 520 | assert descriptors.sizes[:count].tolist() == [2 * 4] |
| 521 | assert worker._submit_dma_descriptors.call_args.kwargs["transfer_ids"] == (7,) |
| 522 | |
| 523 | |
| 524 | def test_hisparse_eager_mirror_records_transfer_without_page_copy(): |
| 525 | worker = _make_hisparse_worker() |
| 526 | worker.cache_handles = [ |
| 527 | SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=True)) |
| 528 | ] |
| 529 | worker._record_transfer_completion = MagicMock() |
| 530 | worker._enqueue_transfers = MagicMock() |
| 531 | transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)] |
| 532 | |
| 533 | worker._submit_transfers(transfers) |
| 534 | |
| 535 | worker._record_transfer_completion.assert_called_once_with(transfers) |
| 536 | worker._enqueue_transfers.assert_not_called() |
| 537 | |
| 538 | |
| 539 | @pytest.mark.parametrize("is_host_writer", [False, True]) |
| 540 | @pytest.mark.parametrize("device", ["cpu", "cuda"]) |
| 541 | def test_hisparse_tail_restore_preserves_imported_rows( |
| 542 | monkeypatch, request, is_host_writer, device |
| 543 | ): |
| 544 | """Restore externally pinned host rows on every rank under sync checking.""" |
| 545 | if device == "cuda" and not torch.cuda.is_available(): |
| 546 | pytest.skip("CUDA is required for the device copy check") |
| 547 | worker = _make_hisparse_worker() |
| 548 | worker.is_host_writer = is_host_writer |
| 549 | worker.kernel_block_size = 4 |
| 550 | worker.host_caches = tuple( |
| 551 | torch.arange(24, dtype=torch.float32).reshape(12, 2) + layer * 100 |
| 552 | for layer in range(2) |
| 553 | ) |
| 554 | if device == "cuda": |
| 555 | # Match production's cudaHostRegister pool, which torch.is_pinned() |
| 556 | # does not recognize, rather than torch's caching pinned allocator. |
| 557 | host_pool, registered = hisparse_runtime_module.allocate_pinned_host_pool( |
| 558 | sum(cache.nbytes for cache in worker.host_caches) |
| 559 | ) |
| 560 | request.addfinalizer( |
| 561 | lambda: hisparse_runtime_module.release_pinned_state([], [registered]) |
| 562 | ) |
| 563 | host_caches = host_pool.view(torch.float32).view(2, 12, 2).unbind() |
| 564 | for destination, source in zip(host_caches, worker.host_caches): |
| 565 | destination.copy_(source) |
| 566 | worker.host_caches = host_caches |
| 567 | worker.resident_caches = tuple( |
| 568 | torch.full((4, 4, 2), -1.0, device=device) for _ in range(2) |
| 569 | ) |
| 570 | worker.cache_handles = [ |
| 571 | SimpleNamespace( |
| 572 | runtime=SimpleNamespace(resident_source_index=i, eager_host_mirror=True) |
| 573 | ) |
| 574 | for i in range(2) |
| 575 | ] |
| 576 | worker._record_transfer_completion = MagicMock() |
| 577 | stream = torch.cuda.current_stream() if device == "cuda" else MagicMock() |
| 578 | monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) |
| 579 | transfer = SparseKVPageTransfer(7, 2, (1, 3), False, restore=True) |
| 580 | |
| 581 | restore_pages = worker._restore_pages |
| 582 | if device == "cuda": |
| 583 | import vllm.utils.gpu_sync_debug as gpu_sync_debug |
| 584 | |
| 585 | monkeypatch.setattr(gpu_sync_debug, "_SYNC_CHECK_MODE", "error") |
| 586 | monkeypatch.setattr(gpu_sync_debug, "_sync_check_enabled", False) |
| 587 | gpu_sync_debug.enable_gpu_sync_check() |
| 588 | restore_pages = gpu_sync_debug.with_gpu_sync_check(restore_pages) |
| 589 | restore_pages([transfer]) |
| 590 | |
| 591 | # Appending into unused capacity must preserve every imported prompt row. |
| 592 | for cache, host, block in zip( |
| 593 | worker.resident_caches, worker.host_caches, transfer.resident_block_ids |
| 594 | ): |
| 595 | assert torch.equal(cache[block].cpu(), host[8:12]) |
| 596 | cache[block, 2:] = 999 |
| 597 | assert torch.equal(cache[block, :2].cpu(), host[8:10]) |
| 598 | assert torch.all(cache[0] == -1) |
| 599 | worker._record_transfer_completion.assert_called_once_with( |
| 600 | [transfer], stream=stream |
| 601 | ) |
| 602 | |
| 603 | |
| 604 | def test_hisparse_lazy_mirror_copies_transfer_pages(): |
| 605 | worker = _make_hisparse_worker() |
| 606 | worker.cache_handles = [ |
| 607 | SimpleNamespace(runtime=SimpleNamespace(eager_host_mirror=False)) |
| 608 | ] |
| 609 | worker._record_transfer_completion = MagicMock() |
| 610 | worker._enqueue_transfers = MagicMock() |
| 611 | transfers = [SparseKVPageTransfer(7, 2, (1,), after_forward=False)] |
| 612 | |
| 613 | worker._submit_transfers(transfers) |
| 614 | |
| 615 | worker._enqueue_transfers.assert_called_once_with(transfers) |
| 616 | worker._record_transfer_completion.assert_not_called() |
| 617 | |
| 618 | |
| 619 | def test_hisparse_dma_row_mirror_builds_descriptors(monkeypatch): |
| 620 | worker = _make_hisparse_worker() |
| 621 | worker.is_host_writer = True |
| 622 | worker.kernel_block_size = 2 |
| 623 | source = torch.empty((3, 2, 4), dtype=torch.uint8) |
| 624 | destination = torch.empty((6, 4), dtype=torch.uint8) |
| 625 | worker.resident_caches = (source,) |
| 626 | worker.host_caches = (destination,) |
| 627 | worker.cache_handles = [ |
| 628 | SimpleNamespace( |
| 629 | runtime=SimpleNamespace(resident_source_index=0), decode_batch=True |
| 630 | ) |
| 631 | ] |
| 632 | worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) |
| 633 | worker.dma_stream = MagicMock() |
| 634 | worker.host_write_event = MagicMock() |
| 635 | worker._dma_free_descriptors = [] |
| 636 | worker._pending_dma_descriptors = deque() |
| 637 | worker._pending_transfer_events = deque() |
| 638 | worker._enqueued_transfer_ids = [] |
| 639 | worker._dma_submitted = False |
| 640 | current_stream = MagicMock() |
| 641 | event = MagicMock() |
| 642 | swap_blocks_batch = MagicMock() |
| 643 | monkeypatch.setattr( |
| 644 | hisparse_worker_module, "current_stream", lambda: current_stream |
| 645 | ) |
| 646 | monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext()) |
| 647 | monkeypatch.setattr(torch, "Event", lambda: event) |
| 648 | monkeypatch.setattr( |
| 649 | hisparse_worker_module.ops, "swap_blocks_batch", swap_blocks_batch |
| 650 | ) |
| 651 | |
| 652 | worker._set_row_mirrors((SparseKVRowMirror((3,), 4, 1),)) |
| 653 | worker._enqueue_row_dma(range(1)) |
| 654 | |
| 655 | src_ptrs, dst_ptrs, sizes = swap_blocks_batch.call_args.args |
| 656 | assert src_ptrs.tolist() == [source.data_ptr() + 3 * 4] |
| 657 | assert dst_ptrs.tolist() == [destination.data_ptr() + 4 * 4] |
| 658 | assert sizes.tolist() == [4] |
| 659 | worker.dma_stream.wait_stream.assert_called_once_with(current_stream) |
| 660 | worker.host_write_event.record.assert_called_once_with(worker.dma_stream) |
| 661 | event.record.assert_called_once_with(worker.dma_stream) |
| 662 | assert worker._dma_submitted |
| 663 | |
| 664 | |
| 665 | def test_hisparse_row_dma_uses_resident_spans(): |
| 666 | worker = _make_hisparse_worker() |
| 667 | source = torch.empty((2, 2, 4), dtype=torch.uint8) |
| 668 | destination = torch.empty((12, 4), dtype=torch.uint8) |
| 669 | worker.is_host_writer = True |
| 670 | worker.kernel_block_size = 2 |
| 671 | worker.resident_caches = (source,) |
| 672 | worker.host_caches = (destination,) |
| 673 | worker.cache_handles = [ |
| 674 | SimpleNamespace( |
| 675 | runtime=SimpleNamespace(resident_source_index=0), decode_batch=True |
| 676 | ) |
| 677 | ] |
| 678 | worker._set_row_mirrors( |
| 679 | ( |
| 680 | SparseKVRowMirror((0,), 4, 2), |
| 681 | SparseKVRowMirror((2,), 9, 1), |
| 682 | ) |
| 683 | ) |
| 684 | worker._dma_free_descriptors = [] |
| 685 | worker._submit_dma_descriptors = MagicMock() |
| 686 | |
| 687 | worker._enqueue_row_dma(range(1)) |
| 688 | |
| 689 | descriptors, count = worker._submit_dma_descriptors.call_args.args |
| 690 | assert count == 2 |
| 691 | assert descriptors.src[:count].tolist() == [ |
| 692 | source.data_ptr(), |
| 693 | source.data_ptr() + 2 * 4, |
| 694 | ] |
| 695 | assert descriptors.dst[:count].tolist() == [ |
| 696 | destination.data_ptr() + 4 * 4, |
| 697 | destination.data_ptr() + 9 * 4, |
| 698 | ] |
| 699 | assert descriptors.sizes[:count].tolist() == [2 * 4, 4] |
| 700 | |
| 701 | |
| 702 | def test_hisparse_finish_forward_mirrors_all_layers_once(monkeypatch): |
| 703 | dst_slots = torch.tensor([7, 8, 9], dtype=torch.int64) |
| 704 | req_ids = torch.tensor([0, 1, 2], dtype=torch.int32) |
| 705 | leader = SimpleNamespace( |
| 706 | eager_host_mirror=True, |
| 707 | is_group_leader=True, |
| 708 | invalidate_written_slots=MagicMock(), |
| 709 | ) |
| 710 | follower = SimpleNamespace( |
| 711 | eager_host_mirror=True, |
| 712 | is_group_leader=False, |
| 713 | invalidate_written_slots=MagicMock(), |
| 714 | ) |
| 715 | handles = [ |
| 716 | SimpleNamespace( |
| 717 | runtime=runtime, |
| 718 | decode_batch=True, |
| 719 | host_mirror_required=True, |
| 720 | num_actual_tokens=3, |
| 721 | num_decode_tokens=2, |
| 722 | req_id_per_token=req_ids, |
| 723 | mirror_slot_mapping=dst_slots, |
| 724 | ) |
| 725 | for runtime in (leader, follower) |
| 726 | ] |
| 727 | worker = _make_hisparse_worker() |
| 728 | worker.is_host_writer = True |
| 729 | worker.cache_handles = handles |
| 730 | worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) |
| 731 | worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 3),)) |
| 732 | worker._dma_submitted = False |
| 733 | worker._per_layer_mirrored = set() |
| 734 | worker._submitted_mirror_layers = set() |
| 735 | worker._post_forward_transfers = [] |
| 736 | worker._enqueue_row_dma = MagicMock() |
| 737 | worker.host_write_event = MagicMock() |
| 738 | worker._forward_ready_event = MagicMock() |
| 739 | current_stream = MagicMock() |
| 740 | monkeypatch.setattr( |
| 741 | hisparse_worker_module, "current_stream", lambda: current_stream |
| 742 | ) |
| 743 | |
| 744 | worker.finish_forward() |
| 745 | |
| 746 | worker._forward_ready_event.record.assert_called_once_with() |
| 747 | worker._enqueue_row_dma.assert_called_once_with( |
| 748 | (0, 1), ready_event=worker._forward_ready_event |
| 749 | ) |
| 750 | leader.invalidate_written_slots.assert_called_once() |
| 751 | torch.testing.assert_close( |
| 752 | leader.invalidate_written_slots.call_args.args[0], dst_slots |
| 753 | ) |
| 754 | follower.invalidate_written_slots.assert_not_called() |
| 755 | worker.host_write_event.record.assert_called_once_with(current_stream) |
| 756 | |
| 757 | |
| 758 | def test_hisparse_finish_forward_does_not_repeat_per_layer_mirrors(): |
| 759 | slots = torch.tensor([7, 8], dtype=torch.int64) |
| 760 | runtime = SimpleNamespace( |
| 761 | eager_host_mirror=False, |
| 762 | is_group_leader=False, |
| 763 | invalidate_written_slots=MagicMock(), |
| 764 | ) |
| 765 | handles = [ |
| 766 | SimpleNamespace( |
| 767 | runtime=runtime, |
| 768 | decode_batch=False, |
| 769 | host_mirror_required=True, |
| 770 | num_actual_tokens=2, |
| 771 | num_decode_tokens=0, |
| 772 | req_id_per_token=torch.empty(0, dtype=torch.int32), |
| 773 | mirror_slot_mapping=slots, |
| 774 | ) |
| 775 | for _ in range(2) |
| 776 | ] |
| 777 | worker = _make_hisparse_worker() |
| 778 | worker.is_host_writer = True |
| 779 | worker.cache_handles = handles |
| 780 | worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),)) |
| 781 | worker._per_layer_mirrored = {0, 1} |
| 782 | worker._submitted_mirror_layers = {0, 1} |
| 783 | worker._enqueue_row_dma = MagicMock() |
| 784 | |
| 785 | worker._enqueue_host_mirror() |
| 786 | |
| 787 | worker._enqueue_row_dma.assert_not_called() |
| 788 | |
| 789 | |
| 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] |
| 793 | handles = [ |
| 794 | SimpleNamespace( |
| 795 | runtime=SimpleNamespace( |
| 796 | eager_host_mirror=False, |
| 797 | is_group_leader=False, |
| 798 | resident_source_index=source_index, |
| 799 | ), |
| 800 | decode_batch=False, |
| 801 | host_mirror_required=layer_index < 5, |
| 802 | num_actual_tokens=2 if layer_index < 5 else 0, |
| 803 | num_decode_tokens=0, |
| 804 | req_id_per_token=torch.empty(0, dtype=torch.int32), |
| 805 | mirror_slot_mapping=slots, |
| 806 | ) |
| 807 | for layer_index, source_index in enumerate(source_indices) |
| 808 | ] |
| 809 | worker = _make_hisparse_worker() |
| 810 | worker.is_host_writer = True |
| 811 | worker.cache_handles = handles |
| 812 | worker._set_row_mirrors((SparseKVRowMirror((0, 0, 0), 7, 2),)) |
| 813 | worker._per_layer_mirrored = set() |
| 814 | worker._submitted_mirror_layers = set() |
| 815 | worker._layer_ready_events = tuple(MagicMock() for _ in handles) |
| 816 | worker._enqueue_row_dma = MagicMock() |
| 817 | |
| 818 | for layer_index in range(5): |
| 819 | worker._enqueue_layer_mirror(layer_index) |
| 820 | |
| 821 | worker._enqueue_row_dma.assert_called_once_with( |
| 822 | (0, 1), ready_event=worker._layer_ready_events[1] |
| 823 | ) |
| 824 | worker._enqueue_host_mirror(ready_event=worker._layer_ready_events[-1]) |
| 825 | assert worker._enqueue_row_dma.call_args_list[1].args == ((2, 3, 4),) |
| 826 | assert worker._enqueue_row_dma.call_args_list[1].kwargs == { |
| 827 | "ready_event": worker._layer_ready_events[-1] |
| 828 | } |
| 829 | |
| 830 | |
| 831 | def test_hisparse_finish_forward_rejects_partial_per_layer_mirror(): |
| 832 | slots = torch.tensor([7, 8], dtype=torch.int64) |
| 833 | runtime = SimpleNamespace( |
| 834 | eager_host_mirror=False, |
| 835 | is_group_leader=False, |
| 836 | invalidate_written_slots=MagicMock(), |
| 837 | ) |
| 838 | handles = [ |
| 839 | SimpleNamespace( |
| 840 | runtime=runtime, |
| 841 | decode_batch=False, |
| 842 | host_mirror_required=True, |
| 843 | num_actual_tokens=2, |
| 844 | num_decode_tokens=0, |
| 845 | req_id_per_token=torch.empty(0, dtype=torch.int32), |
| 846 | mirror_slot_mapping=slots, |
| 847 | ) |
| 848 | for _ in range(2) |
| 849 | ] |
| 850 | worker = _make_hisparse_worker() |
| 851 | worker.is_host_writer = True |
| 852 | worker.cache_handles = handles |
| 853 | worker._set_row_mirrors((SparseKVRowMirror((0, 0), 7, 2),)) |
| 854 | worker._per_layer_mirrored = {0} |
| 855 | worker._submitted_mirror_layers = set() |
| 856 | worker._enqueue_row_dma = MagicMock() |
| 857 | |
| 858 | with pytest.raises(RuntimeError, match="did not mirror every active layer"): |
| 859 | worker._enqueue_host_mirror() |
| 860 | |
| 861 | |
| 862 | def test_hisparse_finish_forward_orders_next_forward_after_dma(monkeypatch): |
| 863 | current_stream = MagicMock() |
| 864 | worker = _make_hisparse_worker() |
| 865 | worker.hot_backing = SimpleNamespace(device=torch.device("cuda:0")) |
| 866 | worker.is_host_writer = True |
| 867 | worker._dma_submitted = True |
| 868 | worker.host_write_event = MagicMock() |
| 869 | worker.cache_handles = [ |
| 870 | SimpleNamespace( |
| 871 | runtime=SimpleNamespace(eager_host_mirror=True), num_actual_tokens=0 |
| 872 | ) |
| 873 | ] |
| 874 | worker._post_forward_transfers = [] |
| 875 | worker._pending_dma_descriptors = deque() |
| 876 | worker._dma_free_descriptors = [] |
| 877 | worker._forward_ready_event = MagicMock() |
| 878 | worker._finish_mirror_phase = MagicMock() |
| 879 | monkeypatch.setattr( |
| 880 | hisparse_worker_module, "current_stream", lambda: current_stream |
| 881 | ) |
| 882 | |
| 883 | worker.finish_forward() |
| 884 | |
| 885 | current_stream.wait_event.assert_called_once_with(worker.host_write_event) |
| 886 | assert not worker._dma_submitted |
| 887 | |
| 888 | |
| 889 | def test_hisparse_finish_forward_mirrors_standalone_mtp_cache(monkeypatch): |
| 890 | dst_slots = torch.tensor([7, 8], dtype=torch.int64) |
| 891 | runtime = SimpleNamespace( |
| 892 | eager_host_mirror=True, |
| 893 | is_group_leader=False, |
| 894 | invalidate_written_slots=MagicMock(), |
| 895 | ) |
| 896 | inactive = SimpleNamespace( |
| 897 | runtime=runtime, |
| 898 | decode_batch=False, |
| 899 | host_mirror_required=False, |
| 900 | num_actual_tokens=0, |
| 901 | num_decode_tokens=0, |
| 902 | req_id_per_token=None, |
| 903 | mirror_slot_mapping=dst_slots, |
| 904 | ) |
| 905 | mtp = SimpleNamespace( |
| 906 | runtime=runtime, |
| 907 | decode_batch=True, |
| 908 | host_mirror_required=True, |
| 909 | num_actual_tokens=2, |
| 910 | num_decode_tokens=2, |
| 911 | req_id_per_token=torch.tensor([0, 1], dtype=torch.int32), |
| 912 | mirror_slot_mapping=dst_slots, |
| 913 | ) |
| 914 | worker = _make_hisparse_worker() |
| 915 | worker.is_host_writer = True |
| 916 | worker.cache_handles = [inactive, mtp] |
| 917 | worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),)) |
| 918 | worker._per_layer_mirrored = set() |
| 919 | worker._submitted_mirror_layers = set() |
| 920 | worker._enqueue_row_dma = MagicMock() |
| 921 | |
| 922 | worker._enqueue_host_mirror() |
| 923 | |
| 924 | worker._enqueue_row_dma.assert_called_once_with((1,), ready_event=None) |
| 925 | |
| 926 | |
| 927 | def test_hisparse_shared_host_reader_skips_mirror(monkeypatch): |
| 928 | """A non-writer TP rank must not mirror rows into the shared host pool.""" |
| 929 | dst_slots = torch.tensor([7, 8], dtype=torch.int64) |
| 930 | leader = SimpleNamespace( |
| 931 | eager_host_mirror=True, |
| 932 | is_group_leader=True, |
| 933 | invalidate_written_slots=MagicMock(), |
| 934 | ) |
| 935 | handle = SimpleNamespace( |
| 936 | runtime=leader, |
| 937 | decode_batch=True, |
| 938 | host_mirror_required=True, |
| 939 | num_actual_tokens=2, |
| 940 | num_decode_tokens=1, |
| 941 | req_id_per_token=torch.tensor([0], dtype=torch.int32), |
| 942 | mirror_slot_mapping=dst_slots, |
| 943 | ) |
| 944 | worker = _make_hisparse_worker() |
| 945 | worker.is_host_writer = False |
| 946 | worker.cache_handles = [handle] |
| 947 | worker._set_row_mirrors((SparseKVRowMirror((0,), 7, 2),)) |
| 948 | worker._per_layer_mirrored = set() |
| 949 | worker._enqueue_row_dma = MagicMock() |
| 950 | |
| 951 | worker._enqueue_host_mirror() |
| 952 | |
| 953 | worker._enqueue_row_dma.assert_not_called() |
| 954 | leader.invalidate_written_slots.assert_called_once() |
| 955 | |
| 956 | |
| 957 | def test_hisparse_shared_host_reader_skips_transfer_completion(): |
| 958 | """A non-writer TP rank must not acknowledge host writes.""" |
| 959 | worker = _make_hisparse_worker() |
| 960 | worker.is_host_writer = False |
| 961 | worker.kernel_block_size = 1 |
| 962 | worker._enqueued_transfer_ids = [] |
| 963 | worker._pending_transfer_events = [] |
| 964 | |
| 965 | worker._record_transfer_completion([SparseKVPageTransfer(1, 2, (3,), True)]) |
| 966 | |
| 967 | assert worker._enqueued_transfer_ids == [] |
| 968 | assert worker._pending_transfer_events == [] |
| 969 | |
| 970 | |
| 971 | def test_hisparse_writer_records_transfer_completion_after_dma(monkeypatch): |
| 972 | worker = _make_hisparse_worker() |
| 973 | worker.is_host_writer = True |
| 974 | worker.dma_stream = MagicMock() |
| 975 | worker._enqueued_transfer_ids = [] |
| 976 | worker._pending_transfer_events = deque() |
| 977 | event = MagicMock() |
| 978 | monkeypatch.setattr(torch, "Event", lambda: event) |
| 979 | |
| 980 | worker._record_transfer_completion( |
| 981 | [ |
| 982 | SparseKVPageTransfer(3, 2, (1,), False), |
| 983 | SparseKVPageTransfer(7, 4, (5,), True), |
| 984 | ] |
| 985 | ) |
| 986 | |
| 987 | event.record.assert_called_once_with(worker.dma_stream) |
| 988 | assert worker._enqueued_transfer_ids == [3, 7] |
| 989 | assert worker._pending_transfer_events == deque([(event, (3, 7))]) |
| 990 | |
| 991 | |
| 992 | @pytest.mark.parametrize("is_host_writer", [False, True]) |
| 993 | def test_hisparse_step_waits_for_previous_host_write(monkeypatch, is_host_writer): |
| 994 | worker = _make_hisparse_worker() |
| 995 | worker.is_host_writer = is_host_writer |
| 996 | worker.hot_backing = SimpleNamespace(device=torch.device("cuda:1")) |
| 997 | host_write_events = (MagicMock(), MagicMock()) |
| 998 | worker.host_write_events = host_write_events |
| 999 | worker.host_write_event = host_write_events[1] |
| 1000 | worker._next_host_write_event = 0 |
| 1001 | worker.host_caches = () |
| 1002 | worker.host_num_blocks = 1 |
| 1003 | worker._post_forward_transfers = [] |
| 1004 | worker._pending_invalid_block_ids = [] |
| 1005 | worker.cache_handles = [ |
| 1006 | SimpleNamespace( |
| 1007 | runtime=SimpleNamespace(eager_host_mirror=True), |
| 1008 | decode_batch=False, |
| 1009 | host_mirror_required=False, |
| 1010 | num_actual_tokens=0, |
| 1011 | num_decode_tokens=0, |
| 1012 | req_id_per_token=None, |
| 1013 | ) |
| 1014 | ] |
| 1015 | worker._per_layer_mirrored = set() |
| 1016 | worker._submitted_mirror_layers = set() |
| 1017 | worker._layer_mirror_callbacks = () |
| 1018 | stream = MagicMock() |
| 1019 | monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) |
| 1020 | |
| 1021 | worker.start_step( |
| 1022 | SimpleNamespace( |
| 1023 | host_block_copies=[], |
| 1024 | command=None, |
| 1025 | source_block_ids=[], |
| 1026 | row_mirrors={}, |
| 1027 | all_context_pages_resident=True, |
| 1028 | row_mirrors_from_resident=False, |
| 1029 | ), |
| 1030 | None, |
| 1031 | ) |
| 1032 | worker.start_step( |
| 1033 | SimpleNamespace( |
| 1034 | host_block_copies=[], |
| 1035 | command=None, |
| 1036 | source_block_ids=[], |
| 1037 | row_mirrors={}, |
| 1038 | all_context_pages_resident=True, |
| 1039 | row_mirrors_from_resident=False, |
| 1040 | ), |
| 1041 | None, |
| 1042 | ) |
| 1043 | |
| 1044 | assert stream.wait_event.call_args_list == [ |
| 1045 | call(host_write_events[1]), |
| 1046 | call(host_write_events[0]), |
| 1047 | ] |
| 1048 | assert worker.host_write_event is host_write_events[1] |
| 1049 | |
| 1050 | |
| 1051 | @pytest.mark.parametrize("tp_rank, expected_copies", [(0, 1), (1, 0)]) |
| 1052 | def test_hisparse_shared_host_block_copy_has_one_writer( |
| 1053 | monkeypatch, tp_rank, expected_copies |
| 1054 | ): |
| 1055 | worker = _make_hisparse_worker() |
| 1056 | worker.shared_host_region = object() |
| 1057 | worker.host_caches = (torch.empty(4, 1),) |
| 1058 | worker.host_num_blocks = 4 |
| 1059 | previous_event = MagicMock() |
| 1060 | copies = (KVCacheBlockCopy(0, 1),) |
| 1061 | copy_blocks = MagicMock() |
| 1062 | tp_group = MagicMock() |
| 1063 | monkeypatch.setattr( |
| 1064 | hisparse_worker_module, "get_tensor_model_parallel_rank", lambda: tp_rank |
| 1065 | ) |
| 1066 | monkeypatch.setattr(hisparse_worker_module, "get_tp_group", lambda: tp_group) |
| 1067 | monkeypatch.setattr( |
| 1068 | hisparse_worker_module, "copy_kv_cache_blocks_inplace", copy_blocks |
| 1069 | ) |
| 1070 | |
| 1071 | worker._copy_host_blocks(copies, previous_event) |
| 1072 | |
| 1073 | assert copy_blocks.call_count == expected_copies |
| 1074 | tp_group.barrier.assert_called_once_with() |
| 1075 | assert worker.take_completed_host_copies() == [1] |
| 1076 | assert worker.take_completed_host_copies() == [] |
| 1077 | |
| 1078 | |
| 1079 | def test_hisparse_empty_step_does_not_replay_stale_host_mirror(monkeypatch): |
| 1080 | handle = SimpleNamespace( |
| 1081 | runtime=SimpleNamespace(eager_host_mirror=True), |
| 1082 | decode_batch=True, |
| 1083 | num_actual_tokens=2, |
| 1084 | num_decode_tokens=2, |
| 1085 | req_id_per_token=torch.tensor([0, 1]), |
| 1086 | mirror_slot_mapping=torch.tensor([4, 5]), |
| 1087 | ) |
| 1088 | worker = _make_hisparse_worker() |
| 1089 | worker.is_host_writer = True |
| 1090 | worker.hot_backing = SimpleNamespace(device=torch.device("cpu")) |
| 1091 | worker.host_write_event = MagicMock() |
| 1092 | worker._forward_ready_event = MagicMock() |
| 1093 | worker.host_caches = () |
| 1094 | worker.host_num_blocks = 1 |
| 1095 | worker.cache_handles = [handle] |
| 1096 | worker._per_layer_mirrored = set() |
| 1097 | worker._submitted_mirror_layers = set() |
| 1098 | worker._layer_mirror_callbacks = (MagicMock(),) |
| 1099 | worker._post_forward_transfers = [] |
| 1100 | worker._pending_invalid_block_ids = [] |
| 1101 | worker._enqueue_host_mirror = MagicMock(wraps=worker._enqueue_host_mirror) |
| 1102 | stream = MagicMock() |
| 1103 | monkeypatch.setattr(hisparse_worker_module, "current_stream", lambda: stream) |
| 1104 | |
| 1105 | worker.start_step( |
| 1106 | SimpleNamespace( |
| 1107 | host_block_copies=[], |
| 1108 | command=None, |
| 1109 | source_block_ids=[], |
| 1110 | row_mirrors={}, |
| 1111 | all_context_pages_resident=True, |
| 1112 | row_mirrors_from_resident=False, |
| 1113 | ), |
| 1114 | None, |
| 1115 | ) |
| 1116 | worker.finish_forward() |
| 1117 | |
| 1118 | worker._enqueue_host_mirror.assert_called_once_with(worker._forward_ready_event) |
| 1119 | assert handle.num_actual_tokens == 0 |
| 1120 | torch.testing.assert_close(handle.mirror_slot_mapping, torch.tensor([4, 5])) |
| 1121 | |
| 1122 | |
| 1123 | def test_hisparse_cache_handles_join_index_groups_during_construction(monkeypatch): |
| 1124 | """Followers must not allocate duplicate runtime state before profiling.""" |
| 1125 | config = SimpleNamespace( |
| 1126 | scheduler_config=SimpleNamespace( |
| 1127 | max_num_seqs=2, |
| 1128 | max_num_batched_tokens=2, |
| 1129 | async_scheduling=False, |
| 1130 | ), |
| 1131 | speculative_config=None, |
| 1132 | kv_transfer_config=None, |
| 1133 | ) |
| 1134 | resolved = hisparse_runtime_module.ResolvedHiSparseConfig( |
| 1135 | top_k=4, |
| 1136 | device_buffer_size=8, |
| 1137 | max_union_rows=8, |
| 1138 | ) |
| 1139 | monkeypatch.setattr(hisparse_runtime_module, "_has_hisparse_ops", lambda: True) |
| 1140 | monkeypatch.setattr( |
| 1141 | hisparse_runtime_module.ResolvedHiSparseConfig, |
| 1142 | "from_vllm_config", |
| 1143 | classmethod(lambda cls, vllm_config, model_top_k: resolved), |
| 1144 | ) |
| 1145 | shared_states: list[object] = [] |
| 1146 | streams: list[object] = [] |
| 1147 | |
| 1148 | def create_shared_state(_device, _max_rows, _top_k): |
| 1149 | shared_states.append(object()) |
| 1150 | return shared_states[-1] |
| 1151 | |
| 1152 | def create_stream(_device): |
| 1153 | streams.append(object()) |
| 1154 | return streams[-1] |
| 1155 | |
| 1156 | monkeypatch.setattr( |
| 1157 | hisparse_runtime_module, "_create_shared_topk_state", create_shared_state |
| 1158 | ) |
| 1159 | monkeypatch.setattr(hisparse_runtime_module, "_create_copy_stream", create_stream) |
| 1160 | monkeypatch.setattr(index_group_module, "_create_side_stream", lambda _: object()) |
| 1161 | monkeypatch.setattr(index_group_module, "_create_event", lambda: object()) |
| 1162 | index_group_builder = index_group_module.SparseMLAIndexGroupBuilder( |
| 1163 | torch.empty((2, 4), dtype=torch.int32) |
| 1164 | ) |
| 1165 | |
| 1166 | def make_cache_handle(is_leader: bool): |
| 1167 | index_group, _ = index_group_builder.register_layer(is_leader) |
| 1168 | cache_handle = hisparse_runtime_module.create_hisparse_cache_handle( |
| 1169 | config, |
| 1170 | model_top_k=4, |
| 1171 | is_index_group_leader=is_leader, |
| 1172 | row_width=8, |
| 1173 | kv_dtype=torch.float32, |
| 1174 | index_group=index_group, |
| 1175 | device="cpu", |
| 1176 | ) |
| 1177 | assert cache_handle is not None |
| 1178 | return cache_handle |
| 1179 | |
| 1180 | first_leader = make_cache_handle(True) |
| 1181 | first_index_group = index_group_builder.current_group |
| 1182 | first_follower = make_cache_handle(False) |
| 1183 | second_leader = make_cache_handle(True) |
| 1184 | second_index_group = index_group_builder.current_group |
| 1185 | second_follower = make_cache_handle(False) |
| 1186 | |
| 1187 | assert first_index_group is not None |
| 1188 | assert second_index_group is not None |
| 1189 | assert first_follower.runtime.index_group is first_leader.runtime.index_group |
| 1190 | assert second_follower.runtime.index_group is second_leader.runtime.index_group |
| 1191 | assert first_leader.runtime.index_group is not second_leader.runtime.index_group |
| 1192 | assert first_leader.runtime.index_group.copy_stream is first_index_group.side_stream |
| 1193 | assert ( |
| 1194 | first_leader.runtime.index_group.logical_topk_ready |
| 1195 | is first_index_group.logical_topk_ready |
| 1196 | ) |
| 1197 | assert ( |
| 1198 | second_leader.runtime.index_group.copy_stream is second_index_group.side_stream |
| 1199 | ) |
| 1200 | assert ( |
| 1201 | second_leader.runtime.index_group.logical_topk_ready |
| 1202 | is second_index_group.logical_topk_ready |
| 1203 | ) |
| 1204 | assert len(shared_states) == 2 |
| 1205 | assert streams == [] |
| 1206 | |
| 1207 | |
| 1208 | def test_hisparse_worker_shutdown_releases_pinned_state(monkeypatch): |
| 1209 | worker = object.__new__(HiSparseConnectorWorker) |
| 1210 | worker._initialized = True |
| 1211 | worker._slot_mapping_staging = None |
| 1212 | worker.dma_stream = None |
| 1213 | worker.cache_handles = [] |
| 1214 | worker.pinned_host_pools = [] |
| 1215 | worker.shared_host_region = object() |
| 1216 | released = False |
| 1217 | |
| 1218 | def release_pinned_state(runtimes, pinned_host_pools, shared_host_region): |
| 1219 | nonlocal released |
| 1220 | assert runtimes == [] |
| 1221 | assert pinned_host_pools == [] |
| 1222 | assert shared_host_region is worker.shared_host_region |
| 1223 | released = True |
| 1224 | |
| 1225 | monkeypatch.setattr( |
| 1226 | hisparse_worker_module, "release_pinned_state", release_pinned_state |
| 1227 | ) |
| 1228 | |
| 1229 | worker.shutdown() |
| 1230 | |
| 1231 | assert released |
| 1232 | |
| 1233 | |
| 1234 | class _TestReplaySSMMixer(MambaMixer2): |
| 1235 | def __init__(self): |
| 1236 | torch.nn.Module.__init__(self) |
| 1237 | self.use_replayssm = True |
| 1238 | self.mamba_config = MambaConfig(backend=MambaBackendEnum.FLASHINFER) |
| 1239 | self._replayssm_ring_start = torch.empty(0, dtype=torch.int32) |
| 1240 | self._replayssm_prev_num_accepted = torch.empty(0, dtype=torch.int32) |
| 1241 | self._replayssm_prev_query_len = torch.empty(0, dtype=torch.int32) |
| 1242 | self._commits_replayssm_trackers = True |
| 1243 | self._updates_replayssm_trackers = True |
| 1244 | |
| 1245 | def get_state_shape(self) -> tuple[tuple[int, ...], ...]: |
| 1246 | return ((2,), (3,), (4,), (5,), (6,)) |
| 1247 | |
| 1248 | def get_state_dtype(self) -> tuple[torch.dtype, ...]: |
| 1249 | return (torch.float32,) * 5 |
| 1250 | |
| 1251 | |
| 1252 | def _packed_replayssm_cache(num_blocks: int) -> torch.Tensor: |
| 1253 | return torch.full((num_blocks, 1, 1, 80), 0, dtype=torch.int8) |
| 1254 | |
| 1255 | |
| 1256 | @pytest.mark.parametrize("layers_only", [False, True]) |
| 1257 | def test_bind_kv_cache_shares_replayssm_trackers_by_cache_group(layers_only): |
| 1258 | mixers = [_TestReplaySSMMixer() for _ in range(3)] |
| 1259 | layer_names = [f"layers.{i}.mixer" for i in range(3)] |
| 1260 | ctx = dict(zip(layer_names, mixers)) |
| 1261 | # Reverse insertion order: updater must follow layer index, not dict order. |
| 1262 | kv_cache = { |
| 1263 | layer_names[2]: _packed_replayssm_cache(4), |
| 1264 | layer_names[1]: _packed_replayssm_cache(4), |
| 1265 | layer_names[0]: _packed_replayssm_cache(4), |
| 1266 | } |
| 1267 | kv_cache_groups = [ |
| 1268 | SimpleNamespace(layer_names=[layer_names[0], layer_names[2]]), |
| 1269 | SimpleNamespace(layer_names=[layer_names[1]]), |
| 1270 | ] |
| 1271 | |
| 1272 | if layers_only: |
| 1273 | bind_kv_cache_to_layers(kv_cache, ctx, kv_cache_groups=kv_cache_groups) |
| 1274 | else: |
| 1275 | bind_kv_cache(kv_cache, ctx, [], kv_cache_groups=kv_cache_groups) |
| 1276 | |
| 1277 | tracker_names = ( |
| 1278 | "_replayssm_ring_start", |
| 1279 | "_replayssm_prev_num_accepted", |
| 1280 | "_replayssm_prev_query_len", |
| 1281 | ) |
| 1282 | for tracker_name in tracker_names: |
| 1283 | group_tracker = getattr(mixers[0], tracker_name) |
| 1284 | assert group_tracker.data_ptr() == getattr(mixers[2], tracker_name).data_ptr() |
| 1285 | assert group_tracker.data_ptr() != getattr(mixers[1], tracker_name).data_ptr() |
| 1286 | assert group_tracker.shape == (4,) |
| 1287 | assert torch.count_nonzero(group_tracker) == 0 |
| 1288 | |
| 1289 | assert [m._commits_replayssm_trackers for m in mixers] == [True, True, False] |
| 1290 | assert [m._updates_replayssm_trackers for m in mixers] == [False, True, True] |
| 1291 | |
| 1292 | |
| 1293 | def test_bind_kv_cache(default_vllm_config): |
| 1294 | from vllm.model_executor.layers.attention import Attention |
| 1295 | |
| 1296 | ctx = { |
| 1297 | "layers.0.self_attn": Attention(32, 128, 0.1, prefix="layers.0.self_attn"), |
| 1298 | "layers.1.self_attn": Attention(32, 128, 0.1, prefix="layers.1.self_attn"), |
| 1299 | "layers.2.self_attn": Attention(32, 128, 0.1, prefix="layers.2.self_attn"), |
| 1300 | "layers.3.self_attn": Attention(32, 128, 0.1, prefix="layers.3.self_attn"), |
| 1301 | } |
| 1302 | kv_cache = { |
| 1303 | "layers.0.self_attn": torch.zeros((1,)), |
| 1304 | "layers.1.self_attn": torch.zeros((1,)), |
| 1305 | "layers.2.self_attn": torch.zeros((1,)), |
| 1306 | "layers.3.self_attn": torch.zeros((1,)), |
| 1307 | } |
| 1308 | runner_kv_caches: list[torch.Tensor] = [] |
| 1309 | bind_kv_cache(kv_cache, ctx, runner_kv_caches) |
| 1310 | assert ctx["layers.0.self_attn"].kv_cache is kv_cache["layers.0.self_attn"] |
| 1311 | assert ctx["layers.1.self_attn"].kv_cache is kv_cache["layers.1.self_attn"] |
| 1312 | assert ctx["layers.2.self_attn"].kv_cache is kv_cache["layers.2.self_attn"] |
| 1313 | assert ctx["layers.3.self_attn"].kv_cache is kv_cache["layers.3.self_attn"] |
| 1314 | |
| 1315 | assert runner_kv_caches[0] is kv_cache["layers.0.self_attn"] |
| 1316 | assert runner_kv_caches[1] is kv_cache["layers.1.self_attn"] |
| 1317 | assert runner_kv_caches[2] is kv_cache["layers.2.self_attn"] |
| 1318 | assert runner_kv_caches[3] is kv_cache["layers.3.self_attn"] |
| 1319 | |
| 1320 | |
| 1321 | def test_bind_kv_cache_non_attention(default_vllm_config): |
| 1322 | from vllm.model_executor.layers.attention import Attention |
| 1323 | |
| 1324 | # example from Jamba PP=2 |
| 1325 | ctx = { |
| 1326 | "model.layers.20.attn": Attention(32, 128, 0.1, prefix="model.layers.20.attn"), |
| 1327 | "model.layers.28.attn": Attention(32, 128, 0.1, prefix="model.layers.28.attn"), |
| 1328 | } |
| 1329 | kv_cache = { |
| 1330 | "model.layers.20.attn": torch.zeros((1,)), |
| 1331 | "model.layers.28.attn": torch.zeros((1,)), |
| 1332 | } |
| 1333 | |
| 1334 | runner_kv_caches: list[torch.Tensor] = [] |
| 1335 | bind_kv_cache(kv_cache, ctx, runner_kv_caches) |
| 1336 | |
| 1337 | assert ctx["model.layers.20.attn"].kv_cache is kv_cache["model.layers.20.attn"] |
| 1338 | assert ctx["model.layers.28.attn"].kv_cache is kv_cache["model.layers.28.attn"] |
| 1339 | |
| 1340 | assert runner_kv_caches[0] is kv_cache["model.layers.20.attn"] |
| 1341 | assert runner_kv_caches[1] is kv_cache["model.layers.28.attn"] |
| 1342 | |
| 1343 | |
| 1344 | def test_bind_kv_cache_draft_model(default_vllm_config): |
| 1345 | from vllm.model_executor.layers.attention import Attention |
| 1346 | |
| 1347 | layer_names = [ |
| 1348 | "model.layers.0.attn", |
| 1349 | "model.layers.1.attn", |
| 1350 | "draft_model.layers.0.attn", |
| 1351 | "draft_model.layers.1.attn", |
| 1352 | ] |
| 1353 | ctx = { |
| 1354 | layer_name: Attention(32, 128, 0.1, prefix=layer_name) |
| 1355 | for layer_name in layer_names |
| 1356 | } |
| 1357 | kv_cache = {layer_name: torch.zeros((1,)) for layer_name in layer_names} |
| 1358 | runner_kv_caches: list[torch.Tensor] = [] |
| 1359 | bind_kv_cache(kv_cache, ctx, runner_kv_caches) |
| 1360 | |
| 1361 | assert ctx["model.layers.0.attn"].kv_cache is kv_cache["model.layers.0.attn"] |
| 1362 | assert ctx["model.layers.1.attn"].kv_cache is kv_cache["model.layers.1.attn"] |
| 1363 | assert ( |
| 1364 | ctx["draft_model.layers.0.attn"].kv_cache |
| 1365 | is kv_cache["draft_model.layers.0.attn"] |
| 1366 | ) |
| 1367 | assert ( |
| 1368 | ctx["draft_model.layers.1.attn"].kv_cache |
| 1369 | is kv_cache["draft_model.layers.1.attn"] |
| 1370 | ) |
| 1371 | |
| 1372 | # caches are ordered by layer_index, interleaving target and draft model |
| 1373 | assert runner_kv_caches[0] is kv_cache["model.layers.0.attn"] |
| 1374 | assert runner_kv_caches[1] is kv_cache["draft_model.layers.0.attn"] |
| 1375 | assert runner_kv_caches[2] is kv_cache["model.layers.1.attn"] |
| 1376 | assert runner_kv_caches[3] is kv_cache["draft_model.layers.1.attn"] |
| 1377 | |
| 1378 | |
| 1379 | def _memory_snapshot(total_gib: int, free_gib: int) -> MemorySnapshot: |
| 1380 | return MemorySnapshot( |
| 1381 | free_memory=free_gib * GiB_bytes, |
| 1382 | total_memory=total_gib * GiB_bytes, |
| 1383 | device="cpu", |
| 1384 | auto_measure=False, |
| 1385 | ) |
| 1386 | |
| 1387 | |
| 1388 | def test_request_memory_charges_external_weights(): |
| 1389 | """Externally held weights are charged against the utilization budget; |
| 1390 | the engine is granted only the remainder.""" |
| 1391 | cache_config = CacheConfig(gpu_memory_utilization=0.9) |
| 1392 | # 70 GiB of a 100 GiB device is held externally before the worker starts. |
| 1393 | snapshot = _memory_snapshot(total_gib=100, free_gib=30) |
| 1394 | |
| 1395 | with pytest.raises(ValueError, match="less than desired"): |
| 1396 | request_memory(snapshot, cache_config) |
| 1397 | |
| 1398 | # 90 GiB budget - 70 GiB external = 20 GiB for the engine. |
| 1399 | assert request_memory(snapshot, cache_config, 70 * GiB_bytes) == 20 * GiB_bytes |
| 1400 | |
| 1401 | # The external weights alone exceed the utilization budget. |
| 1402 | with pytest.raises(ValueError, match="exceed the desired"): |
| 1403 | request_memory(snapshot, cache_config, 95 * GiB_bytes) |
| 1404 | |
| 1405 | # Other tenants squeeze free memory below the engine's remainder. |
| 1406 | with pytest.raises(ValueError, match="less than the engine's budget"): |
| 1407 | request_memory( |
| 1408 | _memory_snapshot(total_gib=100, free_gib=10), cache_config, 70 * GiB_bytes |
| 1409 | ) |