[Fast Start] Cache the MTP draft model in a separate daemon group (#57312)
Signed-off-by: liusy58 <mg21330037@smail.nju.edu.cn> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
siyu committed
Sep 22, 2026 at 11:13 UTC
c723a831a81cb4ff89ea6d61b2d15109306ed1bc
16 files changed
+417
-140
tests/model_executor/model_loader/test_weight_cache.py
+49
-5
@@ -30,10 +30,19 @@ DAEMON_TIMEOUT_S = 600
30
class WeightCacheDaemon:
31
"""Context manager running the real weight cache daemon as a subprocess."""
32
33
- def __init__(self, model: str, tp_size: int, extra_args: list[str] | None = None):
33
+ def __init__(
34
+ self,
35
+ model: str,
36
+ tp_size: int,
37
+ extra_args: list[str] | None = None,
38
+ num_groups: int = 1,
39
+ ):
40
# Short base path: Unix socket paths are limited to ~107 characters.
41
self.socket_dir = tempfile.mkdtemp(prefix="vllm_ipc_")
42
self.tp_size = tp_size
43
+ # Each daemon group (target plus cached drafts) binds one socket per
44
+ # local rank.
45
+ self.num_sockets = tp_size * num_groups
46
self._cmd = [
47
sys.executable,
48
"-m",
@@ -87,7 +96,7 @@ class WeightCacheDaemon:
96
f"Weight cache daemon exited with {self._proc.returncode}:\n"
97
f"{self._logs()}"
98
)
90
- if len(glob.glob(pattern)) >= self.tp_size:
99
+ if len(glob.glob(pattern)) >= self.num_sockets:
100
return
101
time.sleep(1.0)
102
raise TimeoutError(
@@ -112,6 +121,9 @@ class ModelCase:
121
images: list | None = None
122
llm_kwargs: dict[str, Any] = field(default_factory=dict)
123
daemon_args: list[str] = field(default_factory=list)
124
+ # Daemon groups the launcher starts; 2 when a cached draft group joins the
125
+ # target group (MTP/EAGLE speculative decoding).
126
+ daemon_groups: int = 1
127
128
129
def generate(
@@ -172,14 +184,41 @@ K3_CASE = ModelCase(
184
daemon_args=["--trust-remote-code"],
185
)
186
187
+# Qwen3.5-0.8B ships one MTP layer in the target checkpoint, so method="mtp"
188
+# loads the draft from the same model. The daemon must cache it in its draft
189
+# group for the warm runs (fallback=False) to succeed.
190
+QWEN_MTP_CASE = ModelCase(
191
+ model="Qwen/Qwen3.5-0.8B",
192
+ prompts=[
193
+ "Hello, my name is",
194
+ "The capital of France is",
195
+ ],
196
+ llm_kwargs=dict(
197
+ gpu_memory_utilization=0.3,
198
+ enforce_eager=True,
199
+ enable_chunked_prefill=True,
200
+ speculative_config={"method": "mtp", "num_speculative_tokens": 1},
201
+ ),
202
+ daemon_args=[
203
+ "--speculative-config",
204
+ '{"method": "mtp", "num_speculative_tokens": 1}',
205
+ ],
206
+ daemon_groups=2,
207
+)
208
176
-@pytest.mark.parametrize("case", [QWEN_CASE, K3_CASE], ids=["qwen3.5", "kimi-k3"])
209
+
210
+@pytest.mark.parametrize(
211
+ "case",
212
+ [QWEN_CASE, K3_CASE, QWEN_MTP_CASE],
213
+ ids=["qwen3.5", "kimi-k3", "qwen3.5-mtp"],
214
+)
215
def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
216
"""Cold start falls back to disk; warm restarts load weights via CUDA IPC.
217
218
All runs must produce outputs identical to a default-loader baseline. The
219
warm runs disable the disk fallback, so they only pass if the weights
182
- really came from the daemon.
220
+ really came from the daemon — for the MTP case, both the target's and the
221
+ draft's daemon groups.
222
"""
223
if not current_platform.is_cuda_alike():
224
pytest.skip("Weight cache IPC sharing requires CUDA or ROCm")
@@ -199,7 +238,12 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
238
with tempfile.TemporaryDirectory(prefix="vllm_ipc_empty_") as empty_socket_dir:
239
cold_outputs = generate(vllm_runner, case, empty_socket_dir, fallback=True)
240
202
- with WeightCacheDaemon(case.model, tp_size=1, extra_args=case.daemon_args) as d:
241
+ with WeightCacheDaemon(
242
+ case.model,
243
+ tp_size=1,
244
+ extra_args=case.daemon_args,
245
+ num_groups=case.daemon_groups,
246
+ ) as d:
247
warm_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
248
# Warm restart: a second engine lifetime against the same daemon.
249
restart_outputs = generate(vllm_runner, case, d.socket_dir, fallback=False)
tests/model_executor/test_qwen3_omni.py
+2
-1
@@ -369,11 +369,12 @@ def test_dspark_shares_target_embedding_with_smaller_draft_vocabulary():
369
draft_parallel_config=SimpleNamespace(tensor_parallel_size=1),
370
attention_backend=None,
371
kv_cache_dtype=None,
372
+ draft_load_config=None,
373
),
374
parallel_config=ParallelConfig(),
375
attention_config=SimpleNamespace(backend=None),
376
cache_config=SimpleNamespace(),
376
- load_config=SimpleNamespace(),
377
+ load_config=SimpleNamespace(load_format="auto"),
378
model_config=SimpleNamespace(get_vocab_size=Mock(return_value=100)),
379
)
380
tests/v1/spec_decode/test_draft_attention_backend_override.py
+4
-1
@@ -12,7 +12,7 @@ from unittest.mock import patch
12
13
import pytest
14
15
-from vllm.config import LoadConfig
15
+from vllm.config import LoadConfig, SpeculativeConfig
16
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
17
18
@@ -37,6 +37,9 @@ class _SpeculativeConfig:
37
moe_backend: str | None = None
38
kv_cache_dtype: str | None = None
39
draft_model_config: object = None
40
+ draft_load_config: object = None
41
+
42
+ apply_draft_overrides = SpeculativeConfig.apply_draft_overrides
43
44
45
@dataclass
tests/v1/spec_decode/test_draft_moe_backend_override.py
+4
-1
@@ -14,7 +14,7 @@ from unittest.mock import patch
14
15
import pytest
16
17
-from vllm.config import LoadConfig
17
+from vllm.config import LoadConfig, SpeculativeConfig
18
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
19
20
@@ -34,6 +34,9 @@ class _SpeculativeConfig:
34
moe_backend: str | None = None
35
kv_cache_dtype: str | None = None
36
draft_model_config: object = None
37
+ draft_load_config: object = None
38
+
39
+ apply_draft_overrides = SpeculativeConfig.apply_draft_overrides
40
41
42
@dataclass
vllm/config/speculative.py
+26
-1
@@ -14,7 +14,7 @@ from vllm.config.cache import CacheDType
14
from vllm.config.kernel import MoEBackend
15
from vllm.config.model import HfOverrides, ModelConfig
16
from vllm.config.parallel import ParallelConfig
17
-from vllm.config.utils import config
17
+from vllm.config.utils import config, replace
18
from vllm.logger import init_logger
19
from vllm.transformers_utils.config import get_hf_text_config
20
from vllm.utils.hashing import safe_hash
@@ -25,8 +25,10 @@ if TYPE_CHECKING:
25
from transformers import PretrainedConfig
26
27
import vllm.model_executor.layers.quantization as me_quant
28
+ from vllm.config.vllm import VllmConfig
29
else:
30
PretrainedConfig = Any
31
+ VllmConfig = Any
32
33
me_quant = LazyLoader(
34
"model_executor", globals(), "vllm.model_executor.layers.quantization"
@@ -371,6 +373,17 @@ def _validate_qwen3_omni_dspark(
373
)
374
375
376
+# (SpeculativeConfig field, VllmConfig sub-config, overridden field)
377
+_DRAFT_VLLM_CONFIG_OVERRIDES = (
378
+ # Otherwise the draft inherits the target's --moe-backend, which fails
379
+ # when the draft is unquantized and that backend is not.
380
+ ("moe_backend", "kernel_config", "moe_backend"),
381
+ # Only when set, so the draft keeps a KV cache layout the target shares.
382
+ ("attention_backend", "attention_config", "backend"),
383
+ ("kv_cache_dtype", "cache_config", "cache_dtype"),
384
+)
385
+
386
+
387
@config
388
class SpeculativeConfig:
389
"""Configuration for speculative decoding."""
@@ -1771,6 +1784,18 @@ class SpeculativeConfig:
1784
1785
return draft_parallel_config
1786
1787
+ def apply_draft_overrides(self, vllm_config: VllmConfig) -> VllmConfig:
1788
+ """Overlay this config's kernel overrides onto a target VllmConfig.
1789
+
1790
+ Only non-None fields override, so an unset field keeps whatever the
1791
+ target resolved.
1792
+ """
1793
+ for src, config_name, dst in _DRAFT_VLLM_CONFIG_OVERRIDES:
1794
+ if (value := getattr(self, src)) is not None:
1795
+ sub_config = replace(getattr(vllm_config, config_name), **{dst: value})
1796
+ vllm_config = replace(vllm_config, **{config_name: sub_config})
1797
+ return vllm_config
1798
+
1799
@field_validator("attention_backend", mode="before")
1800
@classmethod
1801
def _parse_attention_backend(cls, value: Any) -> Any:
vllm/model_executor/model_loader/utils.py
+37
-1
@@ -12,7 +12,13 @@ from torch import nn
12
from typing_extensions import assert_never
13
14
import vllm.envs as envs
15
-from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
15
+from vllm.config import (
16
+ LoadConfig,
17
+ ModelConfig,
18
+ VllmConfig,
19
+ replace,
20
+ set_current_vllm_config,
21
+)
22
from vllm.logger import init_logger
23
from vllm.model_executor.layers.attention import is_deferred_attention_layer
24
from vllm.model_executor.layers.quantization.base_config import (
@@ -23,6 +29,9 @@ from vllm.model_executor.model_loader.reload import (
29
record_metadata_for_reloading,
30
set_torchao_reload_attrs,
31
)
32
+from vllm.model_executor.model_loader.weight_cache.utils import (
33
+ is_draft_model_cacheable,
34
+)
35
from vllm.model_executor.model_loader.weight_tying import maybe_retie_word_embeddings
36
from vllm.model_executor.models.interfaces import SupportsQuant
37
from vllm.model_executor.utils import is_weights_pre_processed
@@ -34,6 +43,33 @@ from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor
43
logger = init_logger(__name__)
44
45
46
+def get_draft_load_config(vllm_config: VllmConfig) -> LoadConfig:
47
+ """Get load config for the speculative draft model."""
48
+ speculative_config = vllm_config.speculative_config
49
+ if (
50
+ speculative_config is not None
51
+ and speculative_config.draft_load_config is not None
52
+ ):
53
+ return speculative_config.draft_load_config
54
+ load_config = vllm_config.load_config
55
+ if load_config is not None and load_config.load_format != "ipc_cache":
56
+ return load_config
57
+ kwargs = (
58
+ # Route the draft to the daemon's draft group.
59
+ {
60
+ "model_loader_extra_config": {
61
+ **load_config.model_loader_extra_config,
62
+ "is_draft": True,
63
+ }
64
+ }
65
+ if is_draft_model_cacheable(speculative_config)
66
+ # No daemon draft group for this method; load from disk instead of
67
+ # hitting the target daemon with a mismatching fingerprint.
68
+ else {"load_format": "auto", "model_loader_extra_config": {}}
69
+ )
70
+ return replace(load_config, **kwargs)
71
+
72
+
73
@instrument(span_name="Initialize model")
74
def initialize_model(
75
vllm_config: VllmConfig,
vllm/model_executor/model_loader/weight_cache/__init__.py
-23
@@ -1,26 +1,3 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader
5
-from vllm.model_executor.model_loader.weight_cache.protocol import (
6
- CacheConfigMismatchError,
7
- TensorEntry,
8
- UnsupportedPlatformForIPCError,
9
- UnsupportedQuantForIPCError,
10
- WeightCacheKey,
11
- WeightCacheUnavailableError,
12
- check_ipc_platform_support,
13
- check_ipc_quant_support,
14
-)
15
-
16
-__all__ = [
17
- "CacheConfigMismatchError",
18
- "IpcModelLoader",
19
- "TensorEntry",
20
- "UnsupportedPlatformForIPCError",
21
- "UnsupportedQuantForIPCError",
22
- "WeightCacheKey",
23
- "WeightCacheUnavailableError",
24
- "check_ipc_platform_support",
25
- "check_ipc_quant_support",
26
-]
vllm/model_executor/model_loader/weight_cache/daemon.py
+139
-44
@@ -41,6 +41,13 @@ The global TP rank of local GPU ``i`` on node ``r`` is
41
``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank
42
assignment, so each engine worker maps its shard from the daemon on its own
43
node.
44
+
45
+With MTP, EAGLE or EAGLE3 speculative decoding the launcher additionally
46
+starts a draft daemon group that caches the draft model. It uses its own cache
47
+key, Unix sockets (``*_draft.sock``) and rendezvous port
48
+(``--weight-cache-draft-master-port``, default ``--weight-cache-master-port +
49
+1``), so each process serves exactly one model role. Other draft types are not
50
+cached and keep loading from disk in the engine.
51
"""
52
53
import contextlib
@@ -52,10 +59,15 @@ import signal
59
import socket
60
import sys
61
from collections.abc import Callable
62
+from itertools import product
63
64
import torch
65
58
-from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
66
+from vllm.config import (
67
+ ParallelConfig,
68
+ VllmConfig,
69
+ set_current_vllm_config,
70
+)
71
from vllm.distributed import (
72
ensure_model_parallel_initialized,
73
init_distributed_environment,
@@ -77,6 +89,11 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
89
send_msg,
90
verify_peer_is_owner,
91
)
92
+from vllm.model_executor.model_loader.weight_cache.utils import (
93
+ export_model_attrs,
94
+ format_daemon_role,
95
+ is_draft_model_cacheable,
96
+)
97
from vllm.platforms import current_platform
98
from vllm.utils.argparse_utils import FlexibleArgumentParser
99
from vllm.utils.network_utils import get_distributed_init_method, get_open_port
@@ -130,30 +147,6 @@ def export_entries(
147
return entries, aliases
148
149
133
-def get_daemon_model(vllm_config: VllmConfig) -> torch.nn.Module:
134
- """Load the daemon's model, composed from the configured loader.
135
-
136
- Runs the quantization check after model creation but before the slow
137
- weight load, so an unsupported method fails fast. Online quantization
138
- always fails the check, so load_model's finalize step for it is
139
- unnecessary here.
140
- """
141
- model_config = vllm_config.model_config
142
- load_config = vllm_config.load_config
143
- loader = get_model_loader(load_config)
144
- device_config = vllm_config.device_config
145
- target_device = torch.device(
146
- device_config.device if load_config.device is None else load_config.device
147
- )
148
- with set_default_torch_dtype(model_config.dtype):
149
- with target_device:
150
- model = loader.create_model(vllm_config, model_config)
151
- check_ipc_quant_support(model)
152
- loader.load_weights(model, model_config)
153
- process_weights_after_loading(model, model_config, target_device)
154
- return model.eval()
155
-
156
-
150
class WeightCacheDaemon:
151
"""Per-GPU process that loads one TP shard and serves CUDA IPC handles."""
152
@@ -164,19 +157,31 @@ class WeightCacheDaemon:
157
local_rank: int,
158
distributed_init_method: str,
159
socket_dir: str | None = None,
160
+ is_draft: bool = False,
161
):
162
self.vllm_config = vllm_config
163
+ # A draft daemon builds the draft with the target's VllmConfig, like
164
+ # the engine does; only the ModelConfig differs.
165
+ if is_draft:
166
+ spec = vllm_config.speculative_config
167
+ assert spec is not None and spec.draft_model_config is not None
168
+ self.model_config = spec.draft_model_config
169
+ else:
170
+ self.model_config = vllm_config.model_config
171
self.tp_rank = tp_rank
172
self.local_rank = local_rank
173
self.distributed_init_method = distributed_init_method
174
self.socket_dir = socket_dir
175
+ self.is_draft = is_draft
176
+ self.role = format_daemon_role(is_draft)
177
self.model: torch.nn.Module | None = None
178
# Fingerprint before loading: process_weights_after_loading may
179
# mutate hf_config.quantization_config.
180
self.cache_config = WeightCacheKey.from_model_config(
177
- vllm_config.model_config,
181
+ self.model_config,
182
tp_size=vllm_config.parallel_config.tensor_parallel_size,
183
tp_rank=tp_rank,
184
+ is_draft=is_draft,
185
)
186
187
def load_model(self) -> None:
@@ -191,12 +196,37 @@ class WeightCacheDaemon:
196
)
197
with set_current_vllm_config(self.vllm_config):
198
ensure_model_parallel_initialized(tp_size, 1)
194
- self.model = get_daemon_model(self.vllm_config)
199
+ self.model = self.get_model()
200
logger.info(
196
- "Weight cache daemon rank %d loaded model",
201
+ "Weight cache %s daemon rank %d loaded model",
202
+ self.role,
203
self.tp_rank,
204
)
205
206
+ def get_model(self) -> torch.nn.Module:
207
+ """Load the daemon's model, composed from the configured loader.
208
+
209
+ Runs the quantization check after model creation but before the slow
210
+ weight load, so an unsupported method fails fast. Online quantization
211
+ always fails the check, so load_model's finalize step for it is
212
+ unnecessary here.
213
+ """
214
+ vllm_config = self.vllm_config
215
+ model_config = self.model_config
216
+ load_config = vllm_config.load_config
217
+ loader = get_model_loader(load_config)
218
+ device_config = vllm_config.device_config
219
+ target_device = torch.device(
220
+ device_config.device if load_config.device is None else load_config.device
221
+ )
222
+ with set_default_torch_dtype(model_config.dtype):
223
+ with target_device:
224
+ model = loader.create_model(vllm_config, model_config)
225
+ check_ipc_quant_support(model)
226
+ loader.load_weights(model, model_config)
227
+ process_weights_after_loading(model, model_config, target_device)
228
+ return model.eval()
229
+
230
def serve_forever(self, ready_callback: Callable[[], None] | None = None) -> None:
231
"""Serve requests until terminated.
232
@@ -223,7 +253,10 @@ class WeightCacheDaemon:
253
os.chmod(socket_path, 0o600)
254
server.listen()
255
logger.info(
226
- "Weight cache daemon rank %d serving on %s", self.tp_rank, socket_path
256
+ "Weight cache %s daemon rank %d serving on %s",
257
+ self.role,
258
+ self.tp_rank,
259
+ socket_path,
260
)
261
if ready_callback is not None:
262
ready_callback()
@@ -270,7 +303,11 @@ class WeightCacheDaemon:
303
304
@property
305
def _socket_path(self) -> str:
273
- return get_socket_path(get_current_device_uuid(), self.socket_dir)
306
+ return get_socket_path(
307
+ get_current_device_uuid(),
308
+ self.socket_dir,
309
+ is_draft=self.is_draft,
310
+ )
311
312
def _handle_connection(self, conn: socket.socket) -> None:
313
request = recv_msg(conn)
@@ -303,11 +340,13 @@ class WeightCacheDaemon:
340
"status": "ok",
341
"entries": entries,
342
"aliases": aliases,
343
+ "attrs": export_model_attrs(self.model),
344
"gpu_uuid": gpu_uuid,
345
},
346
)
347
logger.info_once(
310
- "Weight cache daemon rank %d sent %d tensors (+%d aliases) to engine",
348
+ "Weight cache %s daemon rank %d sent %d tensors (+%d aliases) to engine",
349
+ self.role,
350
self.tp_rank,
351
len(entries),
352
len(aliases),
@@ -316,7 +355,11 @@ class WeightCacheDaemon:
355
def _handle_release(self, conn: socket.socket) -> None:
356
self.model = None
357
torch.accelerator.empty_cache()
319
- logger.info("Weight cache daemon rank %d released cached weights", self.tp_rank)
358
+ logger.info(
359
+ "Weight cache %s daemon rank %d released cached weights",
360
+ self.role,
361
+ self.tp_rank,
362
+ )
363
send_msg(conn, {"status": "ok"})
364
365
@@ -326,13 +369,27 @@ def _run_daemon(
369
vllm_config: VllmConfig,
370
distributed_init_method: str,
371
socket_dir: str | None,
329
- ready_queue: "multiprocessing.Queue[int]",
372
+ ready_queue: "multiprocessing.Queue[tuple[str, int]]",
373
+ is_draft: bool = False,
374
) -> None:
375
daemon = WeightCacheDaemon(
332
- vllm_config, tp_rank, local_rank, distributed_init_method, socket_dir
376
+ vllm_config,
377
+ tp_rank,
378
+ local_rank,
379
+ distributed_init_method,
380
+ socket_dir,
381
+ is_draft,
382
)
383
daemon.load_model()
335
- daemon.serve_forever(ready_callback=lambda: ready_queue.put(tp_rank))
384
+ daemon.serve_forever(ready_callback=lambda: ready_queue.put((daemon.role, tp_rank)))
385
+
386
+
387
+def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None:
388
+ """VllmConfig for the draft daemon group"""
389
+ if not is_draft_model_cacheable(vllm_config.speculative_config):
390
+ return None
391
+ assert vllm_config.speculative_config is not None
392
+ return vllm_config.speculative_config.apply_draft_overrides(vllm_config)
393
394
395
def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:
@@ -369,6 +426,14 @@ def main() -> None:
426
"serving) and match across nodes. Required when --nnodes > 1; defaults "
427
"to a free port for single-node.",
428
)
429
+ parser.add_argument(
430
+ "--weight-cache-draft-master-port",
431
+ type=int,
432
+ default=None,
433
+ help="Rendezvous port for the MTP draft daemon group. Defaults to "
434
+ "--weight-cache-master-port + 1 for multi-node, or a free port for "
435
+ "single-node.",
436
+ )
437
args = parser.parse_args()
438
engine_args = EngineArgs.from_cli_args(args)
439
vllm_config = engine_args.create_engine_config()
@@ -398,27 +463,56 @@ def main() -> None:
463
master_port = args.weight_cache_master_port or get_open_port()
464
distributed_init_method = get_distributed_init_method(master_addr, master_port)
465
466
+ # (is_draft, vllm_config, rendezvous) per daemon group.
467
+ groups: list[tuple[bool, VllmConfig, str]] = [
468
+ (False, vllm_config, distributed_init_method)
469
+ ]
470
+ draft_vllm_config = get_draft_daemon_config(vllm_config)
471
+ if draft_vllm_config is not None:
472
+ draft_master_port = args.weight_cache_draft_master_port or (
473
+ master_port + 1 if nnodes > 1 else get_open_port()
474
+ )
475
+ if draft_master_port == master_port:
476
+ raise ValueError(
477
+ "--weight-cache-draft-master-port must differ from "
478
+ "--weight-cache-master-port"
479
+ )
480
+ groups.append(
481
+ (
482
+ True,
483
+ draft_vllm_config,
484
+ get_distributed_init_method(master_addr, draft_master_port),
485
+ )
486
+ )
487
+
488
ctx = multiprocessing.get_context("spawn")
402
- ready_queue: multiprocessing.Queue[int] = ctx.Queue()
489
+ ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue()
490
# Global TP rank of local GPU i on this node; local index == device index.
491
global_ranks = [
492
node_rank * local_world_size + local_rank
493
for local_rank in range(local_world_size)
494
]
495
+ expected_ready = {
496
+ (format_daemon_role(is_draft), global_rank)
497
+ for (is_draft, _, _), global_rank in product(groups, global_ranks)
498
+ }
499
procs = [
500
ctx.Process(
501
target=_run_daemon,
502
args=(
503
global_rank,
504
local_rank,
414
- vllm_config,
415
- distributed_init_method,
505
+ config,
506
+ init_method,
507
args.weight_cache_socket_dir,
508
ready_queue,
509
+ is_draft,
510
),
419
- name=f"vllm-weight-cache-daemon-{global_rank}",
511
+ name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-{global_rank}",
512
+ )
513
+ for (is_draft, config, init_method), (local_rank, global_rank) in product(
514
+ groups, enumerate(global_ranks)
515
)
421
- for local_rank, global_rank in enumerate(global_ranks)
516
]
517
for proc in procs:
518
proc.start()
@@ -430,10 +524,10 @@ def main() -> None:
524
signal.signal(signal.SIGINT, _shutdown)
525
signal.signal(signal.SIGTERM, _shutdown)
526
433
- ready_ranks: set[int] = set()
434
- while len(ready_ranks) < local_world_size:
527
+ ready: set[tuple[str, int]] = set()
528
+ while len(ready) < len(expected_ready):
529
try:
436
- ready_ranks.add(ready_queue.get(timeout=1.0))
530
+ ready.add(ready_queue.get(timeout=1.0))
531
except queue.Empty:
532
dead = [p for p in procs if p.exitcode is not None]
533
if dead:
@@ -450,10 +544,11 @@ def main() -> None:
544
socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
545
logger.info_once(
546
"===== Weight cache daemon READY: node %d/%d serving %d local rank(s) "
453
- "in %s =====",
547
+ "x %d role(s) in %s =====",
548
node_rank,
549
nnodes,
550
local_world_size,
551
+ len(groups),
552
socket_dir_msg,
553
)
554
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+53
-24
@@ -25,9 +25,9 @@ from vllm.model_executor.model_loader.utils import (
25
)
26
from vllm.model_executor.model_loader.weight_cache.protocol import (
27
CacheConfigMismatchError,
28
- TensorEntry,
28
UnsupportedQuantForIPCError,
29
WeightCacheKey,
30
+ WeightCacheState,
31
WeightCacheUnavailableError,
32
check_ipc_platform_support,
33
check_ipc_quant_support,
@@ -60,7 +60,7 @@ class IpcModelLoader(BaseModelLoader):
60
Extra config keys (via --model-loader-extra-config):
61
62
- socket_path: explicit daemon socket path. Defaults to a per-GPU path
63
- derived from the physical GPU uuid.
63
+ derived from the physical GPU uuid and the cache role (target/draft).
64
- socket_dir: directory containing the daemon sockets.
65
- mode: "zero_copy" (default) or "copy".
66
- fallback: fall back to disk loading when the daemon is unavailable or
@@ -78,6 +78,15 @@ class IpcModelLoader(BaseModelLoader):
78
extra_config = copy(load_config.model_loader_extra_config or {})
79
self.socket_path: str | None = extra_config.pop("socket_path", None)
80
self.socket_dir: str | None = extra_config.pop("socket_dir", None)
81
+ # Internal: set by the engine when routing a speculative draft to the
82
+ # daemon's draft group.
83
+ self.is_draft = bool(extra_config.pop("is_draft", False))
84
+ if self.is_draft and self.socket_path is not None:
85
+ raise ValueError(
86
+ "socket_path cannot be combined with the draft weight cache role; "
87
+ "use socket_dir so the target and draft sockets are derived "
88
+ "independently"
89
+ )
90
self.mode: str = extra_config.pop("mode", "zero_copy")
91
self.fallback: bool = extra_config.pop("fallback", True)
92
self.connect_timeout_s: float = float(
@@ -108,7 +117,7 @@ class IpcModelLoader(BaseModelLoader):
117
loaded through this loader).
118
"""
119
device_index = torch.accelerator.current_device_index()
111
- entries, _ = self._fetch_entries(model_config)
120
+ entries = self._fetch_entries(model_config).entries
121
params = dict(model.named_parameters())
122
buffers = dict(model.named_buffers())
123
for name, entry in entries.items():
@@ -128,11 +137,22 @@ class IpcModelLoader(BaseModelLoader):
137
check_ipc_platform_support()
138
state_fetched = False
139
try:
131
- entries, aliases = self._fetch_entries(model_config)
140
+ # Cross-check the routing flag against the identity of the model
141
+ # being loaded: a draft load that lost its flag (or a target load
142
+ # that got one) would hit the wrong daemon group and
143
+ # fingerprint-mismatch.
144
+ spec = vllm_config.speculative_config
145
+ inferred = spec is not None and model_config is spec.draft_model_config
146
+ if inferred != self.is_draft:
147
+ raise CacheConfigMismatchError(
148
+ f"Weight cache role mismatch: loading "
149
+ f"{'draft' if inferred else 'target'} model but the loader "
150
+ f"was configured for the "
151
+ f"{'draft' if self.is_draft else 'target'} group"
152
+ )
153
+ state = self._fetch_entries(model_config)
154
state_fetched = True
133
- return self._build_model(
134
- vllm_config, model_config, prefix, entries, aliases
135
- )
155
+ return self._build_model(vllm_config, model_config, prefix, state)
156
except (WeightCacheUnavailableError, CacheConfigMismatchError) as e:
157
if not self.fallback:
158
raise
@@ -163,8 +183,7 @@ class IpcModelLoader(BaseModelLoader):
183
vllm_config: VllmConfig,
184
model_config: ModelConfig,
185
prefix: str,
166
- entries: dict[str, TensorEntry],
167
- aliases: dict[str, str],
186
+ state: WeightCacheState,
187
) -> nn.Module:
188
device_config = vllm_config.device_config
189
load_device = (
@@ -186,7 +205,11 @@ class IpcModelLoader(BaseModelLoader):
205
prefix=prefix,
206
)
207
check_ipc_quant_support(model)
189
- self._apply_entries(model, entries, aliases, device_index)
208
+ self._apply_entries(model, state, device_index)
209
+ # Flags that load_weights would have set (e.g. EAGLE ownership of
210
+ # embed_tokens / lm_head); the daemon ran it, this process did not.
211
+ for name, value in state.attrs.items():
212
+ setattr(model, name, value)
213
# The daemon exports tensors that already went through
214
# process_weights_after_loading; re-run it in pre-processed mode
215
# so quant methods only rebuild Python-side state (e.g. the MoE
@@ -202,7 +225,7 @@ class IpcModelLoader(BaseModelLoader):
225
self._send_release()
226
logger.info(
227
"Mapped %d tensors from the weight cache daemon (%s mode)",
205
- len(entries),
228
+ len(state.entries),
229
self.mode,
230
)
231
return model.eval()
@@ -210,8 +233,7 @@ class IpcModelLoader(BaseModelLoader):
233
def _apply_entries(
234
self,
235
model: nn.Module,
213
- entries: dict[str, TensorEntry],
214
- aliases: dict[str, str],
236
+ state: WeightCacheState,
237
device_index: int,
238
) -> None:
239
# remove_duplicate=False keeps tied module aliases reachable by name:
@@ -242,7 +264,7 @@ class IpcModelLoader(BaseModelLoader):
264
module.register_buffer(leaf, obj)
265
registered[name] = obj
266
245
- for name, entry in entries.items():
267
+ for name, entry in state.entries.items():
268
tensor = entry.rebuild(device_index)
269
if self.mode == "copy":
270
tensor = tensor.clone()
@@ -251,7 +273,7 @@ class IpcModelLoader(BaseModelLoader):
273
# Re-establish tied-weight aliases by registering the *same* object the
274
# canonical name resolved to, so parameter identity (and the tie) is
275
# preserved instead of allocating uninitialized memory.
254
- for alias_name, canonical_name in aliases.items():
276
+ for alias_name, canonical_name in state.aliases.items():
277
obj = registered.get(canonical_name)
278
if obj is None:
279
logger.warning(
@@ -262,19 +284,16 @@ class IpcModelLoader(BaseModelLoader):
284
continue
285
_register(alias_name, obj, isinstance(obj, nn.Parameter))
286
265
- def _fetch_entries(
266
- self, model_config: ModelConfig
267
- ) -> tuple[dict[str, TensorEntry], dict[str, str]]:
287
+ def _fetch_entries(self, model_config: ModelConfig) -> WeightCacheState:
288
cache_config = WeightCacheKey.from_model_config(
289
model_config,
290
tp_size=get_tensor_model_parallel_world_size(),
291
tp_rank=get_tensor_model_parallel_rank(),
292
+ is_draft=self.is_draft,
293
)
294
return self._request_state(cache_config)
295
275
- def _request_state(
276
- self, cache_config: WeightCacheKey
277
- ) -> tuple[dict[str, TensorEntry], dict[str, str]]:
296
+ def _request_state(self, cache_config: WeightCacheKey) -> WeightCacheState:
297
with self._connect(self.state_timeout_s) as conn:
298
send_msg(conn, {"cmd": "get_state", "cache_config": cache_config})
299
response = recv_msg(conn)
@@ -288,7 +307,11 @@ class IpcModelLoader(BaseModelLoader):
307
f"Weight cache daemon error: {response.get('message')}"
308
)
309
self._check_gpu_uuid(response.get("gpu_uuid"))
291
- return response["entries"], response.get("aliases", {})
310
+ return WeightCacheState(
311
+ entries=response["entries"],
312
+ aliases=response.get("aliases", {}),
313
+ attrs=response.get("attrs", {}),
314
+ )
315
316
def _connect(self, timeout: float) -> socket.socket:
317
socket_path = self._resolve_socket_path()
@@ -316,7 +339,11 @@ class IpcModelLoader(BaseModelLoader):
339
def _resolve_socket_path(self) -> str:
340
if self.socket_path is not None:
341
return self.socket_path
319
- return get_socket_path(get_current_device_uuid(), self.socket_dir)
342
+ return get_socket_path(
343
+ get_current_device_uuid(),
344
+ self.socket_dir,
345
+ is_draft=self.is_draft,
346
+ )
347
348
def _check_gpu_uuid(self, daemon_uuid: str | None) -> None:
349
if daemon_uuid is None:
@@ -340,7 +367,9 @@ class IpcModelLoader(BaseModelLoader):
367
# DefaultModelLoader must not see load_format="ipc_cache" or the ipc
368
# extra config keys.
369
return dataclasses.replace(
343
- self.load_config, load_format="auto", model_loader_extra_config={}
370
+ self.load_config,
371
+ load_format="auto",
372
+ model_loader_extra_config={},
373
)
374
375
def _fallback_load(
vllm/model_executor/model_loader/weight_cache/protocol.py
+42
-6
@@ -20,7 +20,7 @@ import stat
20
import struct
21
import tempfile
22
from dataclasses import dataclass, fields
23
-from typing import Any
23
+from typing import Any, NamedTuple
24
25
import torch
26
from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor
@@ -29,13 +29,16 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME
29
import vllm.version
30
from vllm.config import ModelConfig
31
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
32
+from vllm.model_executor.model_loader.weight_cache.utils import (
33
+ format_socket_role_suffix,
34
+)
35
from vllm.model_executor.model_loader.weight_utils import (
36
filter_duplicate_safetensors_files,
37
)
38
from vllm.platforms import current_platform
39
from vllm.utils.hashing import safe_hash
40
38
-SOCKET_NAME_TEMPLATE = "vllm_weight_cache_{gpu_uuid}.sock"
41
+SOCKET_NAME_TEMPLATE = "vllm_weight_cache_{gpu_uuid}{role}.sock"
42
SOCKET_DIR_TEMPLATE = "vllm_weight_cache_{uid}"
43
44
_LEN_STRUCT = struct.Struct("!Q")
@@ -128,9 +131,23 @@ def get_socket_dir(socket_dir: str | None = None) -> str:
131
)
132
133
131
-def get_socket_path(gpu_uuid: str, socket_dir: str | None = None) -> str:
132
- directory = get_socket_dir(socket_dir)
133
- return os.path.join(directory, SOCKET_NAME_TEMPLATE.format(gpu_uuid=gpu_uuid))
134
+def get_socket_path(
135
+ gpu_uuid: str,
136
+ socket_dir: str | None = None,
137
+ *,
138
+ is_draft: bool = False,
139
+) -> str:
140
+ """Socket path of a daemon group; ``is_draft=False`` is the target.
141
+
142
+ The GPU uuid is hashed to keep the name well under the AF_UNIX path
143
+ limit (~108 bytes) even with the draft role suffix.
144
+ """
145
+ gpu_id = safe_hash(gpu_uuid.encode()).hexdigest()
146
+ name = SOCKET_NAME_TEMPLATE.format(
147
+ gpu_uuid=gpu_id,
148
+ role=format_socket_role_suffix(is_draft),
149
+ )
150
+ return os.path.join(get_socket_dir(socket_dir), name)
151
152
153
def ensure_private_socket_dir(directory: str, strict_perms: bool = True) -> None:
@@ -272,10 +289,17 @@ class WeightCacheKey:
289
quant_config_hash: str
290
revision: str | None
291
vllm_version: str
292
+ is_draft: bool = False
293
+ """Daemon group the weights come from; False is the target model."""
294
295
@classmethod
296
def from_model_config(
278
- cls, model_config: ModelConfig, tp_size: int, tp_rank: int
297
+ cls,
298
+ model_config: ModelConfig,
299
+ tp_size: int,
300
+ tp_rank: int,
301
+ *,
302
+ is_draft: bool = False,
303
) -> "WeightCacheKey":
304
"""Build the fingerprint for a model configuration.
305
@@ -302,6 +326,7 @@ class WeightCacheKey:
326
quant_config_hash=_hash_quant_config(quant_config),
327
revision=model_config.revision,
328
vllm_version=vllm.version.__version__,
329
+ is_draft=is_draft,
330
)
331
332
def mismatched_fields(self, other: "WeightCacheKey") -> list[str]:
@@ -345,6 +370,17 @@ class TensorEntry:
370
return rebuild_cuda_tensor(*args)
371
372
373
+class WeightCacheState(NamedTuple):
374
+ """Client-side decode of a daemon's get_state response payload."""
375
+
376
+ entries: dict[str, TensorEntry]
377
+ """Model tensors, exported as CUDA IPC handles or shipped by value."""
378
+ aliases: dict[str, str]
379
+ """Duplicate (tied) weight names aliased to their canonical entry."""
380
+ attrs: dict[str, bool]
381
+ """Python-side flags set by load_weights, e.g. EAGLE ownership flags."""
382
+
383
+
384
def send_msg(sock: socket.socket, obj: Any) -> None:
385
payload = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
386
sock.sendall(_LEN_STRUCT.pack(len(payload)))
vllm/model_executor/model_loader/weight_cache/utils.py
new
+48
@@ -0,0 +1,48 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Helpers shared by the weight cache daemon, the IPC loader and the engine.
4
+
5
+These live outside ``protocol`` so callers that only need to know whether a
6
+draft is cached, or how a daemon group is named, do not have to import the
7
+wire format.
8
+"""
9
+
10
+from typing import Any
11
+
12
+from vllm.config import SpeculativeConfig
13
+from vllm.model_executor.models.interfaces import SupportsEagleBase
14
+
15
+# Speculative methods whose draft model the daemon caches in its own group.
16
+# Other drafts keep loading from disk in the engine.
17
+WEIGHT_CACHE_DRAFT_METHODS = frozenset({"mtp", "eagle", "eagle3"})
18
+
19
+# Python-side flags that weight loading sets on EAGLE-style drafts; the engine
20
+# never runs load_weights for cached models, so the daemon ships them.
21
+EXPORTED_MODEL_ATTRS = tuple(SupportsEagleBase.__annotations__)
22
+
23
+
24
+def export_model_attrs(model: Any) -> dict[str, bool]:
25
+ return {
26
+ name: bool(getattr(model, name))
27
+ for name in EXPORTED_MODEL_ATTRS
28
+ if hasattr(model, name)
29
+ }
30
+
31
+
32
+def is_draft_model_cacheable(speculative_config: SpeculativeConfig | None) -> bool:
33
+ """Whether the daemon serves the speculative draft as a separate role."""
34
+ return (
35
+ speculative_config is not None
36
+ and speculative_config.method in WEIGHT_CACHE_DRAFT_METHODS
37
+ and speculative_config.draft_model_config is not None
38
+ )
39
+
40
+
41
+def format_daemon_role(is_draft: bool) -> str:
42
+ """Name of a daemon group: the target model or the draft."""
43
+ return "draft" if is_draft else "target"
44
+
45
+
46
+def format_socket_role_suffix(is_draft: bool) -> str:
47
+ """Socket-name suffix keeping the draft group distinct from the target."""
48
+ return "_draft" if is_draft else ""
vllm/v1/spec_decode/llm_base_proposer.py
+2
-1
@@ -25,6 +25,7 @@ from vllm.forward_context import set_forward_context
25
from vllm.logger import init_logger
26
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
27
from vllm.model_executor.model_loader import get_model
28
+from vllm.model_executor.model_loader.utils import get_draft_load_config
29
from vllm.model_executor.models import (
30
supports_multimodal,
31
supports_multimodal_embeddings,
@@ -1311,7 +1312,7 @@ class SpecDecodeBaseProposer:
1312
model = get_model(
1313
vllm_config=draft_vllm_config,
1314
model_config=self.speculative_config.draft_model_config,
1314
- load_config=self.speculative_config.draft_load_config,
1315
+ load_config=get_draft_load_config(draft_vllm_config),
1316
)
1317
return model
1318
vllm/v1/worker/gpu/spec_decode/dflash/utils.py
+2
-1
@@ -4,6 +4,7 @@ import torch.nn as nn
4
5
from vllm.config import VllmConfig, replace
6
from vllm.model_executor.model_loader import get_model
7
+from vllm.model_executor.model_loader.utils import get_draft_load_config
8
from vllm.v1.worker.gpu.spec_decode.eagle.utils import (
9
_should_share,
10
get_target_lm_head,
@@ -38,7 +39,7 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
39
if speculative_config.kv_cache_dtype is not None
40
else vllm_config.cache_config
41
),
41
- load_config=get_pp_safe_draft_load_config(vllm_config.load_config),
42
+ load_config=get_pp_safe_draft_load_config(get_draft_load_config(vllm_config)),
43
)
44
with set_model_tag("dflash_head"):
45
dflash_model = get_model(
vllm/v1/worker/gpu/spec_decode/dspark/utils.py
+2
-1
@@ -5,6 +5,7 @@ import torch.nn as nn
5
6
from vllm.config import ModelConfig, ParallelConfig, VllmConfig, replace
7
from vllm.logger import init_logger
8
+from vllm.model_executor.model_loader.utils import get_draft_load_config
9
from vllm.v1.attention.backends.registry import AttentionBackendEnum
10
from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
11
@@ -94,7 +95,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo
95
if speculative_config.kv_cache_dtype is not None
96
else vllm_config.cache_config
97
),
97
- load_config=get_pp_safe_draft_load_config(vllm_config.load_config),
98
+ load_config=get_pp_safe_draft_load_config(get_draft_load_config(vllm_config)),
99
)
100
# VllmConfig post-init restores the target's quant config because the target
101
# config is retained for DSpark's target-layer metadata, so we must override it.
vllm/v1/worker/gpu/spec_decode/eagle/utils.py
+5
-29
@@ -7,6 +7,7 @@ from vllm.config import VllmConfig, replace
7
from vllm.distributed.parallel_state import get_pp_group
8
from vllm.lora.layers.base import BaseLayerWithLoRA
9
from vllm.model_executor.model_loader import get_model
10
+from vllm.model_executor.model_loader.utils import get_draft_load_config
11
from vllm.model_executor.models.utils import PPMissingLayer
12
from vllm.v1.worker.gpu.spec_decode.utils import get_pp_safe_draft_load_config
13
@@ -75,35 +76,10 @@ def load_eagle_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mod
76
speculative_config = vllm_config.speculative_config
77
assert speculative_config is not None
78
draft_model_config = speculative_config.draft_model_config
78
- if speculative_config.moe_backend is not None:
79
- # Otherwise the draft inherits the target's --moe-backend, which
80
- # fails when the draft is unquantized and that backend is not.
81
- vllm_config = replace(
82
- vllm_config,
83
- kernel_config=replace(
84
- vllm_config.kernel_config,
85
- moe_backend=speculative_config.moe_backend,
86
- ),
87
- )
88
- if speculative_config.kv_cache_dtype is not None:
89
- vllm_config = replace(
90
- vllm_config,
91
- cache_config=replace(
92
- vllm_config.cache_config,
93
- cache_dtype=speculative_config.kv_cache_dtype,
94
- ),
95
- )
96
- if speculative_config.attention_backend is not None:
97
- # Before get_model(): the backend is read off the constructed layers.
98
- # Only when set, so the draft keeps a KV cache layout the target shares.
99
- vllm_config = replace(
100
- vllm_config,
101
- attention_config=replace(
102
- vllm_config.attention_config,
103
- backend=speculative_config.attention_backend,
104
- ),
105
- )
106
- draft_load_config = get_pp_safe_draft_load_config(vllm_config.load_config)
79
+ vllm_config = speculative_config.apply_draft_overrides(vllm_config)
80
+ draft_load_config = get_pp_safe_draft_load_config(
81
+ get_draft_load_config(vllm_config)
82
+ )
83
if draft_load_config is not vllm_config.load_config:
84
vllm_config = replace(vllm_config, load_config=draft_load_config)
85
with set_model_tag("eagle_head"):
vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py
+2
-1
@@ -16,6 +16,7 @@ from vllm.config import VllmConfig, replace
16
from vllm.distributed.parallel_state import get_pp_group
17
from vllm.logger import init_logger
18
from vllm.model_executor.model_loader import get_model
19
+from vllm.model_executor.model_loader.utils import get_draft_load_config
20
from vllm.v1.worker.gpu.spec_decode.autoregressive.speculator import (
21
AutoRegressiveSpeculator,
22
)
@@ -40,7 +41,7 @@ class Gemma4Speculator(AutoRegressiveSpeculator):
41
draft_model = get_model(
42
vllm_config=draft_vllm_config,
43
model_config=self.speculative_config.draft_model_config,
43
- load_config=self.speculative_config.draft_load_config,
44
+ load_config=get_draft_load_config(draft_vllm_config),
45
)
46
self._setup_gemma4_kv_sharing(draft_model, target_attn_layer_names)
47
self._share_embeddings(draft_model, target_model)