| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import time |
| 5 | |
| 6 | import prometheus_client |
| 7 | import pytest |
| 8 | import torch |
| 9 | |
| 10 | from tests.conftest import VllmRunner |
| 11 | from tests.utils import fork_new_process_for_each_test |
| 12 | from vllm.config import AttentionConfig, HiSparseConfig, KVTransferConfig |
| 13 | from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.connector import ( |
| 14 | HiSparseConnector, |
| 15 | ) |
| 16 | from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import ( |
| 17 | HiSparseConnectorWorker, |
| 18 | ) |
| 19 | from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import ( |
| 20 | MultiConnector, |
| 21 | ) |
| 22 | from vllm.platforms import current_platform |
| 23 | from vllm.sampling_params import SamplingParams |
| 24 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 25 | from vllm.v1.hisparse.coordinator import get_hisparse_coordinator |
| 26 | |
| 27 | MODEL = "deepseek-ai/DeepSeek-V3.2" |
| 28 | |
| 29 | HF_OVERRIDES = { |
| 30 | "num_hidden_layers": 8, |
| 31 | "hidden_size": 256, |
| 32 | "intermediate_size": 512, |
| 33 | "num_attention_heads": 8, |
| 34 | "num_key_value_heads": 1, |
| 35 | "n_routed_experts": 8, |
| 36 | "num_experts_per_tok": 2, |
| 37 | "index_topk": 128, |
| 38 | } |
| 39 | |
| 40 | |
| 41 | def _shrink_config(config): |
| 42 | overrides = HF_OVERRIDES.copy() |
| 43 | if any("MTP" in architecture for architecture in config.architectures): |
| 44 | overrides.pop("num_hidden_layers") |
| 45 | config.update(overrides) |
| 46 | return config |
| 47 | |
| 48 | |
| 49 | def _num_hisparse_spills(runner: VllmRunner) -> int: |
| 50 | client = runner.llm.llm_engine.engine_core |
| 51 | coordinator = get_hisparse_coordinator( |
| 52 | client.engine_core.scheduler.kv_cache_manager |
| 53 | ) |
| 54 | return coordinator.next_spill_id |
| 55 | |
| 56 | |
| 57 | def _offload_load_bytes() -> float: |
| 58 | return sum( |
| 59 | sample.value |
| 60 | for metric in prometheus_client.REGISTRY.collect() |
| 61 | if metric.name == "vllm:kv_offload_load_bytes" |
| 62 | for sample in metric.samples |
| 63 | if sample.name == "vllm:kv_offload_load_bytes_total" |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker: |
| 68 | engine_core = runner.llm.llm_engine.engine_core.engine_core |
| 69 | model_runner = engine_core.model_executor.driver_worker.worker.model_runner |
| 70 | connector = model_runner.kv_connector.kv_connector |
| 71 | if isinstance(connector, MultiConnector): |
| 72 | connector = next( |
| 73 | child |
| 74 | for child in connector.sub_connectors |
| 75 | if isinstance(child, HiSparseConnector) |
| 76 | ) |
| 77 | assert isinstance(connector, HiSparseConnector) |
| 78 | assert connector.connector_worker is not None |
| 79 | return connector.connector_worker |
| 80 | |
| 81 | |
| 82 | @pytest.mark.skipif( |
| 83 | not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" |
| 84 | ) |
| 85 | @pytest.mark.parametrize( |
| 86 | "with_offloading", [False, True], ids=["standalone", "offload"] |
| 87 | ) |
| 88 | @fork_new_process_for_each_test |
| 89 | def test_hisparse_spill_and_prefix_restore( |
| 90 | monkeypatch: pytest.MonkeyPatch, |
| 91 | vllm_runner: type[VllmRunner], |
| 92 | with_offloading: bool, |
| 93 | ): |
| 94 | """Spilled prefixes restore and FULL-graph decode writes reach host KV. |
| 95 | |
| 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: |
| 103 | pytest.skip("Sparse MLA requires Hopper or newer") |
| 104 | |
| 105 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 106 | monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") |
| 107 | |
| 108 | target = [1000 + i % 64 for i in range(257)] |
| 109 | pressure = [ |
| 110 | [2000 + request_idx * 128 + i % 64 for i in range(257)] |
| 111 | for request_idx in range(4) |
| 112 | ] |
| 113 | hisparse_connector = { |
| 114 | "kv_connector": "HiSparseConnector", |
| 115 | "kv_role": "kv_both", |
| 116 | "kv_connector_extra_config": {"host_pool_gib": 1}, |
| 117 | } |
| 118 | kv_transfer_config = KVTransferConfig(**hisparse_connector) |
| 119 | if with_offloading: |
| 120 | kv_transfer_config = KVTransferConfig( |
| 121 | kv_connector="MultiConnector", |
| 122 | kv_role="kv_both", |
| 123 | kv_connector_extra_config={ |
| 124 | "connectors": [ |
| 125 | hisparse_connector, |
| 126 | { |
| 127 | "kv_connector": "OffloadingConnector", |
| 128 | "kv_role": "kv_both", |
| 129 | "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30}, |
| 130 | }, |
| 131 | ] |
| 132 | }, |
| 133 | ) |
| 134 | |
| 135 | with vllm_runner( |
| 136 | MODEL, |
| 137 | load_format="dummy", |
| 138 | hf_overrides=_shrink_config, |
| 139 | attention_config=AttentionConfig( |
| 140 | hisparse_config=HiSparseConfig( |
| 141 | device_buffer_size=512, |
| 142 | ) |
| 143 | ), |
| 144 | kv_transfer_config=kv_transfer_config, |
| 145 | block_size=64, |
| 146 | max_model_len=320, |
| 147 | max_num_batched_tokens=1024, |
| 148 | max_num_seqs=4, |
| 149 | num_gpu_blocks_override=128, |
| 150 | disable_log_stats=False, |
| 151 | enable_chunked_prefill=True, |
| 152 | enable_prefix_caching=True, |
| 153 | speculative_config={"method": "mtp", "num_speculative_tokens": 3}, |
| 154 | compilation_config={ |
| 155 | "cudagraph_mode": "FULL_AND_PIECEWISE", |
| 156 | "cudagraph_capture_sizes": [1, 4], |
| 157 | }, |
| 158 | ) as runner: |
| 159 | worker = _get_hisparse_worker(runner) |
| 160 | engine_core = runner.llm.llm_engine.engine_core.engine_core |
| 161 | model_runner = engine_core.model_executor.driver_worker.worker.model_runner |
| 162 | full_graph_calls = 0 |
| 163 | full_replay_pending = False |
| 164 | original_run_fullgraph = model_runner.cudagraph_manager.run_fullgraph |
| 165 | |
| 166 | def record_fullgraph(batch_desc): |
| 167 | nonlocal full_graph_calls, full_replay_pending |
| 168 | full_graph_calls += 1 |
| 169 | full_replay_pending = True |
| 170 | return original_run_fullgraph(batch_desc) |
| 171 | |
| 172 | monkeypatch.setattr( |
| 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 |
| 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 | |
| 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 |
| 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 |
| 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) |
| 243 | |
| 244 | assert _num_hisparse_spills(runner) > 0 |
| 245 | load_bytes = _offload_load_bytes() |
| 246 | actual = runner.generate_greedy([target], max_tokens=8) |
| 247 | runner.generate_greedy([[42]], max_tokens=1) |
| 248 | |
| 249 | assert actual == expected |
| 250 | if with_offloading: |
| 251 | assert _offload_load_bytes() > load_bytes |
| 252 | |
| 253 | |
| 254 | @pytest.mark.skipif( |
| 255 | not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" |
| 256 | ) |
| 257 | @fork_new_process_for_each_test |
| 258 | def test_hisparse_host_exhaustion_defers_requests( |
| 259 | monkeypatch: pytest.MonkeyPatch, |
| 260 | vllm_runner: type[VllmRunner], |
| 261 | ): |
| 262 | """A full host pool defers requests instead of leaving pages GPU-only.""" |
| 263 | capability = current_platform.get_device_capability() |
| 264 | if capability is None or capability.major < 9: |
| 265 | pytest.skip("Sparse MLA requires Hopper or newer") |
| 266 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 267 | monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") |
| 268 | prompts = [ |
| 269 | [1000 + request_idx * 64 + i % 64 for i in range(257)] |
| 270 | for request_idx in range(3) |
| 271 | ] |
| 272 | with vllm_runner( |
| 273 | MODEL, |
| 274 | load_format="dummy", |
| 275 | hf_overrides=_shrink_config, |
| 276 | attention_config=AttentionConfig( |
| 277 | hisparse_config=HiSparseConfig(device_buffer_size=512) |
| 278 | ), |
| 279 | kv_transfer_config=KVTransferConfig( |
| 280 | kv_connector="HiSparseConnector", |
| 281 | kv_role="kv_both", |
| 282 | kv_connector_extra_config={"host_pool_gib": 1}, |
| 283 | ), |
| 284 | block_size=64, |
| 285 | max_model_len=320, |
| 286 | max_num_batched_tokens=512, |
| 287 | max_num_seqs=3, |
| 288 | num_gpu_blocks_override=128, |
| 289 | enable_prefix_caching=False, |
| 290 | enable_chunked_prefill=True, |
| 291 | enforce_eager=True, |
| 292 | ) as runner: |
| 293 | expected = [runner.generate_greedy([p], max_tokens=16)[0] for p in prompts] |
| 294 | engine = runner.llm.llm_engine.engine_core.engine_core |
| 295 | coordinator = get_hisparse_coordinator(engine.scheduler.kv_cache_manager) |
| 296 | host = coordinator.host_manager |
| 297 | assert host is not None |
| 298 | # Leave room for one full request plus part of another. |
| 299 | spare_blocks = 8 |
| 300 | num_pressure = host.block_pool.get_num_free_blocks() - spare_blocks |
| 301 | host.allocate_new_blocks( |
| 302 | "host-pressure", |
| 303 | num_pressure * host.block_size, |
| 304 | num_pressure * host.block_size, |
| 305 | ) |
| 306 | original_plan = coordinator.plan_prefix_materialization |
| 307 | original_count = host.get_num_blocks_to_allocate |
| 308 | checked_steps = 0 |
| 309 | refusals = 0 |
| 310 | |
| 311 | def plan(request_id, num_computed_tokens): |
| 312 | nonlocal checked_steps |
| 313 | assert not any(b.is_null for b in host.req_to_blocks[request_id]) |
| 314 | checked_steps += 1 |
| 315 | original_plan(request_id, num_computed_tokens) |
| 316 | |
| 317 | def count(*args, **kwargs): |
| 318 | nonlocal refusals |
| 319 | required = original_count(*args, **kwargs) |
| 320 | refusals += int(required > 0) |
| 321 | return required |
| 322 | |
| 323 | monkeypatch.setattr(coordinator, "plan_prefix_materialization", plan) |
| 324 | monkeypatch.setattr(host, "get_num_blocks_to_allocate", count) |
| 325 | actual = runner.generate_greedy(prompts, max_tokens=16) |
| 326 | |
| 327 | assert checked_steps > 0 |
| 328 | assert refusals > 0 |
| 329 | assert actual == expected |
| 330 | |
| 331 | |
| 332 | @pytest.mark.skipif( |
| 333 | not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA" |
| 334 | ) |
| 335 | @fork_new_process_for_each_test |
| 336 | def test_hisparse_terminal_prefix_reuse( |
| 337 | monkeypatch: pytest.MonkeyPatch, |
| 338 | vllm_runner: type[VllmRunner], |
| 339 | ): |
| 340 | """Finished requests publish mirrored host KV after a late acknowledgement.""" |
| 341 | capability = current_platform.get_device_capability() |
| 342 | if capability is None or capability.major < 9: |
| 343 | pytest.skip("Sparse MLA requires Hopper or newer") |
| 344 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 345 | monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip") |
| 346 | target = [1000 + i % 64 for i in range(257)] |
| 347 | sampling = SamplingParams(temperature=0, max_tokens=1, ignore_eos=True) |
| 348 | with vllm_runner( |
| 349 | MODEL, |
| 350 | load_format="dummy", |
| 351 | hf_overrides=_shrink_config, |
| 352 | attention_config=AttentionConfig( |
| 353 | hisparse_config=HiSparseConfig( |
| 354 | device_buffer_size=512, eager_host_mirror=True |
| 355 | ) |
| 356 | ), |
| 357 | kv_transfer_config=KVTransferConfig( |
| 358 | kv_connector="HiSparseConnector", |
| 359 | kv_role="kv_both", |
| 360 | kv_connector_extra_config={"host_pool_gib": 0.01}, |
| 361 | ), |
| 362 | block_size=64, |
| 363 | max_model_len=320, |
| 364 | max_num_batched_tokens=512, |
| 365 | max_num_seqs=2, |
| 366 | num_gpu_blocks_override=128, |
| 367 | enable_prefix_caching=True, |
| 368 | enable_chunked_prefill=True, |
| 369 | enforce_eager=True, |
| 370 | ) as runner: |
| 371 | engine = runner.llm.llm_engine.engine_core.engine_core |
| 372 | coordinator = get_hisparse_coordinator(engine.scheduler.kv_cache_manager) |
| 373 | worker = _get_hisparse_worker(runner) |
| 374 | original_updates = worker.take_transfer_updates |
| 375 | original_free = coordinator.free |
| 376 | held_completions: list[int] = [] |
| 377 | update_calls = 0 |
| 378 | terminal_pending_pages = 0 |
| 379 | |
| 380 | def take_updates(): |
| 381 | nonlocal update_calls |
| 382 | with gpu_sync_allowed(): |
| 383 | torch.accelerator.synchronize() |
| 384 | enqueued, completed = original_updates() |
| 385 | held_completions.extend(completed) |
| 386 | update_calls += 1 |
| 387 | if update_calls <= 2: |
| 388 | return enqueued, [] |
| 389 | completed = held_completions.copy() |
| 390 | held_completions.clear() |
| 391 | return enqueued, completed |
| 392 | |
| 393 | def free(request_id): |
| 394 | nonlocal terminal_pending_pages |
| 395 | state = coordinator.request_states.get(request_id) |
| 396 | if state is not None: |
| 397 | terminal_pending_pages += len(state.pending_pages) |
| 398 | original_free(request_id) |
| 399 | |
| 400 | def drain_pending_work(): |
| 401 | deadline = time.monotonic() + 10 |
| 402 | while coordinator.has_pending_work(): |
| 403 | assert time.monotonic() < deadline, "HiSparse transfers did not drain" |
| 404 | runner.llm.llm_engine.step() |
| 405 | |
| 406 | monkeypatch.setattr(worker, "take_transfer_updates", take_updates) |
| 407 | monkeypatch.setattr(coordinator, "free", free) |
| 408 | [first] = runner.llm.generate( |
| 409 | [{"prompt_token_ids": target}], sampling, use_tqdm=False |
| 410 | ) |
| 411 | assert terminal_pending_pages == 4 |
| 412 | assert first.num_cached_tokens == 0 |
| 413 | assert coordinator.has_pending_work() |
| 414 | drain_pending_work() |
| 415 | assert not coordinator.has_pending_work() |
| 416 | assert not held_completions |
| 417 | assert update_calls > 2 |
| 418 | |
| 419 | [repeated] = runner.llm.generate( |
| 420 | [{"prompt_token_ids": target}], sampling, use_tqdm=False |
| 421 | ) |
| 422 | assert repeated.num_cached_tokens == 256 |
| 423 | assert repeated.outputs[0].token_ids == first.outputs[0].token_ids |
| 424 | drain_pending_work() |
| 425 | assert not coordinator.has_pending_work() |