| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """End-to-end tests for the IPC weight cache loader. |
| 4 | |
| 5 | A cold start (no weight cache daemon, so the loader falls back to disk) and |
| 6 | warm restarts (weights mapped from the daemon via CUDA IPC) must both serve |
| 7 | identical outputs. |
| 8 | """ |
| 9 | |
| 10 | import shutil |
| 11 | import subprocess |
| 12 | import sys |
| 13 | import tempfile |
| 14 | import threading |
| 15 | import urllib.error |
| 16 | import urllib.request |
| 17 | from dataclasses import dataclass, field |
| 18 | from typing import Any |
| 19 | |
| 20 | import pytest |
| 21 | |
| 22 | from vllm import SamplingParams |
| 23 | from vllm.assets.image import ImageAsset |
| 24 | from vllm.platforms import current_platform |
| 25 | from vllm.utils.network_utils import get_open_port |
| 26 | |
| 27 | |
| 28 | class WeightCacheDaemon: |
| 29 | """Context manager running the real weight cache daemon as a subprocess.""" |
| 30 | |
| 31 | def __init__( |
| 32 | self, |
| 33 | model: str, |
| 34 | tp_size: int, |
| 35 | extra_args: list[str] | None = None, |
| 36 | ): |
| 37 | # Short base path: Unix socket paths are limited to ~107 characters. |
| 38 | self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_") |
| 39 | self.health_port = get_open_port() |
| 40 | self._cmd = [ |
| 41 | sys.executable, |
| 42 | "-m", |
| 43 | "vllm.entrypoints.cli.main", |
| 44 | "preload", |
| 45 | "--model", |
| 46 | model, |
| 47 | "--tensor-parallel-size", |
| 48 | str(tp_size), |
| 49 | "--weight-cache-socket-dir", |
| 50 | self.socket_dir, |
| 51 | "--enforce-eager", |
| 52 | "--weight-cache-health-host", |
| 53 | "127.0.0.1", |
| 54 | "--weight-cache-health-port", |
| 55 | str(self.health_port), |
| 56 | *(extra_args or []), |
| 57 | ] |
| 58 | self._proc: subprocess.Popen | None = None |
| 59 | |
| 60 | def __enter__(self) -> "WeightCacheDaemon": |
| 61 | self._proc = subprocess.Popen( |
| 62 | self._cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE, text=True |
| 63 | ) |
| 64 | threading.Thread(target=self._drain_stderr, daemon=True).start() |
| 65 | return self |
| 66 | |
| 67 | def __exit__(self, *exc_info) -> None: |
| 68 | self._stop() |
| 69 | |
| 70 | def _drain_stderr(self) -> None: |
| 71 | assert self._proc is not None and self._proc.stderr is not None |
| 72 | self._proc.stderr.read() |
| 73 | |
| 74 | def check_health(self) -> None: |
| 75 | url = f"http://127.0.0.1:{self.health_port}/health" |
| 76 | try: |
| 77 | with urllib.request.urlopen(url, timeout=5) as response: |
| 78 | assert response.status == 200 |
| 79 | except urllib.error.HTTPError as error: |
| 80 | raise AssertionError( |
| 81 | f"Weight cache daemon is not ready: HTTP {error.code}" |
| 82 | ) from error |
| 83 | |
| 84 | def _stop(self) -> None: |
| 85 | assert self._proc is not None |
| 86 | self._proc.terminate() |
| 87 | try: |
| 88 | self._proc.wait(timeout=60) |
| 89 | except subprocess.TimeoutExpired: |
| 90 | self._proc.kill() |
| 91 | self._proc.wait() |
| 92 | shutil.rmtree(self.socket_dir, ignore_errors=True) |
| 93 | |
| 94 | |
| 95 | @dataclass |
| 96 | class ModelCase: |
| 97 | model: str |
| 98 | prompts: list[str] |
| 99 | images: list | None = None |
| 100 | llm_kwargs: dict[str, Any] = field(default_factory=dict) |
| 101 | daemon_args: list[str] = field(default_factory=list) |
| 102 | |
| 103 | |
| 104 | def generate( |
| 105 | vllm_runner, |
| 106 | case: ModelCase, |
| 107 | socket_dir: str | None, |
| 108 | fallback: bool, |
| 109 | ): |
| 110 | extra_config = ( |
| 111 | {} if socket_dir is None else {"socket_dir": socket_dir, "fallback": fallback} |
| 112 | ) |
| 113 | with vllm_runner( |
| 114 | case.model, |
| 115 | load_format="auto" if socket_dir is None else "ipc_cache", |
| 116 | model_loader_extra_config=extra_config, |
| 117 | # Greedy outputs are compared across engine restarts; cap the batch |
| 118 | # size so scheduling cannot change the results. |
| 119 | max_num_seqs=1, |
| 120 | **case.llm_kwargs, |
| 121 | ) as llm: |
| 122 | sampling_params = SamplingParams(temperature=0, max_tokens=16, ignore_eos=True) |
| 123 | return llm.generate(case.prompts, sampling_params, images=case.images) |
| 124 | |
| 125 | |
| 126 | # Both hybrid-recurrent models need chunked prefill for their mamba-style |
| 127 | # cache mode. |
| 128 | QWEN_CASE = ModelCase( |
| 129 | model="Qwen/Qwen3.5-0.8B", |
| 130 | prompts=[ |
| 131 | "Hello, my name is", |
| 132 | "The capital of France is", |
| 133 | ], |
| 134 | llm_kwargs=dict( |
| 135 | gpu_memory_utilization=0.3, |
| 136 | enforce_eager=True, |
| 137 | enable_chunked_prefill=True, |
| 138 | ), |
| 139 | ) |
| 140 | |
| 141 | # Kimi K3 adds the mxfp4-pack MoE path (post-load kernel setup runs in |
| 142 | # pre-processed mode) and a vision tower, exercised with an image request. |
| 143 | K3_CASE = ModelCase( |
| 144 | model="tiny-random/kimi-k3", |
| 145 | prompts=[ |
| 146 | "The capital of France is", |
| 147 | "def quicksort(arr):", |
| 148 | "<|kimi_image_placeholder|>What is shown in the image?", |
| 149 | ], |
| 150 | images=[None, None, ImageAsset("stop_sign").pil_image], |
| 151 | llm_kwargs=dict( |
| 152 | gpu_memory_utilization=0.3, |
| 153 | enforce_eager=True, |
| 154 | trust_remote_code=True, |
| 155 | enable_chunked_prefill=True, |
| 156 | # The image placeholder expands to >1k tokens for the stop_sign asset. |
| 157 | max_model_len=4096, |
| 158 | ), |
| 159 | daemon_args=["--trust-remote-code"], |
| 160 | ) |
| 161 | |
| 162 | # Qwen3.5-0.8B ships one MTP layer in the target checkpoint, so method="mtp" |
| 163 | # loads the draft from the same model. The daemon must cache it in its draft |
| 164 | # group for the warm runs (fallback=False) to succeed. |
| 165 | QWEN_MTP_CASE = ModelCase( |
| 166 | model="Qwen/Qwen3.5-0.8B", |
| 167 | prompts=[ |
| 168 | "Hello, my name is", |
| 169 | "The capital of France is", |
| 170 | ], |
| 171 | llm_kwargs=dict( |
| 172 | gpu_memory_utilization=0.3, |
| 173 | enforce_eager=True, |
| 174 | enable_chunked_prefill=True, |
| 175 | speculative_config={"method": "mtp", "num_speculative_tokens": 1}, |
| 176 | ), |
| 177 | daemon_args=[ |
| 178 | "--speculative-config", |
| 179 | '{"method": "mtp", "num_speculative_tokens": 1}', |
| 180 | ], |
| 181 | ) |
| 182 | |
| 183 | |
| 184 | @pytest.mark.parametrize( |
| 185 | "case", |
| 186 | [QWEN_CASE, K3_CASE, QWEN_MTP_CASE], |
| 187 | ids=["qwen3.5", "kimi-k3", "qwen3.5-mtp"], |
| 188 | ) |
| 189 | def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): |
| 190 | """Cold start falls back to disk; warm restarts load weights via CUDA IPC. |
| 191 | |
| 192 | All runs must produce outputs identical to a default-loader baseline. The |
| 193 | warm runs disable the disk fallback, so they only pass if the weights |
| 194 | really came from the daemon — for the MTP case, both the target's and the |
| 195 | draft's daemon groups. |
| 196 | """ |
| 197 | if not current_platform.is_cuda_alike(): |
| 198 | pytest.skip("Weight cache IPC sharing requires CUDA or ROCm") |
| 199 | if case is K3_CASE and not current_platform.is_device_capability_family(100): |
| 200 | pytest.skip("Kimi K3 IPC weight cache requires an SM100 MXFP4 backend") |
| 201 | |
| 202 | # Baseline: plain disk loading with the default loader. |
| 203 | baseline_outputs = generate( |
| 204 | vllm_runner, |
| 205 | case, |
| 206 | None, |
| 207 | fallback=True, |
| 208 | ) |
| 209 | assert all(text for _, texts in baseline_outputs for text in texts) |
| 210 | |
| 211 | # Cold start: no daemon is serving, so the loader falls back to disk. |
| 212 | with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir: |
| 213 | cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True) |
| 214 | |
| 215 | with WeightCacheDaemon( |
| 216 | case.model, |
| 217 | tp_size=1, |
| 218 | extra_args=case.daemon_args, |
| 219 | ) as d: |
| 220 | warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) |
| 221 | d.check_health() |
| 222 | # Warm restart: a second engine lifetime against the same daemon. |
| 223 | restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False) |
| 224 | |
| 225 | assert cold_outputs == baseline_outputs |
| 226 | assert warm_outputs == baseline_outputs |
| 227 | assert restart_outputs == baseline_outputs |
| 228 | |
| 229 | |
| 230 | def _parallel(**kw): |
| 231 | from types import SimpleNamespace |
| 232 | |
| 233 | base = dict( |
| 234 | tensor_parallel_size=1, |
| 235 | pipeline_parallel_size=1, |
| 236 | data_parallel_size=1, |
| 237 | data_parallel_size_local=1, |
| 238 | data_parallel_rank=0, |
| 239 | nnodes=1, |
| 240 | node_rank=0, |
| 241 | ) |
| 242 | base.update(kw) |
| 243 | return SimpleNamespace(**base) |
| 244 | |
| 245 | |
| 246 | def test_daemon_places_tp_and_dp_ranks_on_local_gpus(): |
| 247 | """Each node serves a contiguous global-rank block (DP-major), so without |
| 248 | DP local GPU i is TP rank r*local+i, and single-node DP launchers offset |
| 249 | by --data-parallel-start-rank.""" |
| 250 | from vllm.model_executor.model_loader.weight_cache.daemon import plan_local_ranks |
| 251 | |
| 252 | tp_second_node = _parallel(tensor_parallel_size=8, nnodes=2, node_rank=1) |
| 253 | assert plan_local_ranks(tp_second_node) == [(i, 0, 0, 4 + i) for i in range(4)] |
| 254 | |
| 255 | dp_third_node = _parallel( |
| 256 | data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=8 |
| 257 | ) |
| 258 | assert plan_local_ranks(dp_third_node) == [(i, 8 + i, 0, 0) for i in range(4)] |
| 259 | |
| 260 | dp_tp = _parallel( |
| 261 | tensor_parallel_size=2, |
| 262 | data_parallel_size=4, |
| 263 | data_parallel_size_local=2, |
| 264 | data_parallel_rank=2, |
| 265 | ) |
| 266 | assert plan_local_ranks(dp_tp) == [ |
| 267 | (0, 2, 0, 0), |
| 268 | (1, 2, 0, 1), |
| 269 | (2, 3, 0, 0), |
| 270 | (3, 3, 0, 1), |
| 271 | ] |
| 272 | |
| 273 | # DP + nnodes: node-local TP replicas |
| 274 | dp_tp_node0 = _parallel(tensor_parallel_size=8, data_parallel_size=2, nnodes=2) |
| 275 | assert plan_local_ranks(dp_tp_node0) == [(i, 0, 0, i) for i in range(8)] |
| 276 | dp_tp_node1 = _parallel( |
| 277 | tensor_parallel_size=8, data_parallel_size=2, nnodes=2, node_rank=1 |
| 278 | ) |
| 279 | assert plan_local_ranks(dp_tp_node1) == [(i, 1, 0, i) for i in range(8)] |
| 280 | |
| 281 | # DP + nnodes: TP group spans nodes (TP8 x DP2 on 4 nodes) |
| 282 | tp_span_node1 = _parallel( |
| 283 | tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=1 |
| 284 | ) |
| 285 | assert plan_local_ranks(tp_span_node1) == [(i, 0, 0, 4 + i) for i in range(4)] |
| 286 | tp_span_node2 = _parallel( |
| 287 | tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=2 |
| 288 | ) |
| 289 | assert plan_local_ranks(tp_span_node2) == [(i, 1, 0, i) for i in range(4)] |
| 290 | |
| 291 | # DP + nnodes: several DP replicas per node (TP4 x DP4 on 2 nodes) |
| 292 | dp_multi_node0 = _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2) |
| 293 | assert plan_local_ranks(dp_multi_node0) == [(i, i // 4, 0, i % 4) for i in range(8)] |
| 294 | |
| 295 | # PP adds one daemon placement per stage while preserving TP rank order. |
| 296 | pp_tp = _parallel(tensor_parallel_size=2, pipeline_parallel_size=2) |
| 297 | assert plan_local_ranks(pp_tp) == [ |
| 298 | (0, 0, 0, 0), |
| 299 | (1, 0, 0, 1), |
| 300 | (2, 0, 1, 0), |
| 301 | (3, 0, 1, 1), |
| 302 | ] |
| 303 | |
| 304 | |
| 305 | def test_daemon_rejects_unmappable_parallelism(): |
| 306 | from vllm.model_executor.model_loader.weight_cache.daemon import ( |
| 307 | _reject_unsupported_parallelism, |
| 308 | ) |
| 309 | |
| 310 | _reject_unsupported_parallelism( |
| 311 | _parallel( |
| 312 | data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=12 |
| 313 | ) |
| 314 | ) |
| 315 | # DP combines with --nnodes when the world size divides evenly. |
| 316 | _reject_unsupported_parallelism( |
| 317 | _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2) |
| 318 | ) |
| 319 | _reject_unsupported_parallelism(_parallel(pipeline_parallel_size=2)) |
| 320 | with pytest.raises(ValueError, match="evenly divide"): |
| 321 | _reject_unsupported_parallelism(_parallel(tensor_parallel_size=3, nnodes=2)) |
| 322 | with pytest.raises(ValueError, match="evenly divide"): |
| 323 | _reject_unsupported_parallelism( |
| 324 | _parallel(tensor_parallel_size=2, data_parallel_size=3, nnodes=4) |
| 325 | ) |
| 326 | with pytest.raises(ValueError, match="exceeds"): |
| 327 | _reject_unsupported_parallelism( |
| 328 | _parallel( |
| 329 | data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=13 |
| 330 | ) |
| 331 | ) |
| 332 | |
| 333 | |
| 334 | def test_weight_cache_key_distinguishes_dp_ranks(): |
| 335 | from dataclasses import replace |
| 336 | |
| 337 | from vllm.model_executor.model_loader.weight_cache.protocol import WeightCacheKey |
| 338 | |
| 339 | key = WeightCacheKey( |
| 340 | checkpoint="ckpt", |
| 341 | model_arch="Arch", |
| 342 | tp_size=1, |
| 343 | tp_rank=0, |
| 344 | dtype="bf16", |
| 345 | quantization=None, |
| 346 | quant_config_hash="h", |
| 347 | revision=None, |
| 348 | vllm_version="v", |
| 349 | dp_size=16, |
| 350 | dp_rank=3, |
| 351 | pp_size=2, |
| 352 | pp_rank=1, |
| 353 | ) |
| 354 | assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"] |
| 355 | assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"] |
| 356 | assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"] |
| 357 | |
| 358 | |
| 359 | def test_ipc_loader_copy_mode_reports_no_external_weight_memory(): |
| 360 | """Copy mode clones the weights into the engine, so nothing is external |
| 361 | (vllm_config is never touched).""" |
| 362 | from vllm.config import LoadConfig |
| 363 | from vllm.model_executor.model_loader.weight_cache.ipc_loader import ( |
| 364 | IpcModelLoader, |
| 365 | ) |
| 366 | |
| 367 | copy_mode = LoadConfig( |
| 368 | load_format="ipc_cache", model_loader_extra_config={"mode": "copy"} |
| 369 | ) |
| 370 | loader = IpcModelLoader(copy_mode) |
| 371 | assert loader.get_external_weight_memory(None) == 0 |