@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 24928a4d62

[UX][Frontend] Introduce `vllm preload` cli for fast restart (#56680)

Signed-off-by: Isotr0py <Isotr0py@outlook.com>

Isotr0py committed Sep 25, 2026 at 14:56 UTC 24928a4d62d4d317ba070d219ce0ab6c3e58bae4
8 files changed +496 -193
docs/cli/README.md
+28 -1
@@ -9,7 +9,7 @@ vllm --help
9 Available Commands:
10
11 ```bash
12 -vllm {chat,complete,serve,launch,bench,collect-env,run-batch}
12 +vllm {chat,complete,serve,launch,bench,collect-env,run-batch,preload}
13 ```
14
15 ## serve
@@ -212,6 +212,33 @@ vllm run-batch \
212
213 See [vllm run-batch](./run-batch.md) for the full reference of all available arguments.
214
215 +## preload
216 +
217 +Launch weight cache daemons (one per GPU) that hold the post-quantized,
218 +TP-sharded weights in GPU memory and serve CUDA IPC handles to vLLM engines
219 +over a Unix domain socket. Restarting engines then map the weights via
220 +zero-copy IPC instead of reloading from disk, enabling fast engine restarts.
221 +
222 +```bash
223 +# Launch one daemon per GPU
224 +vllm preload --model meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 4
225 +
226 +# Engines then load from the daemons
227 +vllm serve meta-llama/Llama-3.2-1B-Instruct --tensor-parallel-size 4 \
228 + --load-format ipc_cache
229 +```
230 +
231 +The daemon accepts the standard engine arguments (model, dtype, quantization,
232 +tensor-parallel-size, ...) plus `--weight-cache-socket-dir` to override the
233 +directory holding the per-GPU Unix sockets, and `--weight-cache-master-port` /
234 +`--weight-cache-draft-master-port` to pin the daemon rendezvous ports for
235 +multi-node and speculative-decoding setups. Tensor, expert and data
236 +parallelism are supported; pipeline parallelism is rejected at launch.
237 +
238 +See [Preload](../features/preload.md) for how it works, cache modes, and
239 +limitations, and [vllm preload](./preload.md) for the full reference of all
240 +available arguments.
241 +
242 ## More Help
243
244 For detailed options of any subcommand, use:
docs/features/preload.md new
+203
@@ -0,0 +1,203 @@
1 +# Preload
2 +
3 +vLLM's preload feature keeps model artifacts resident in GPU memory across
4 +engine restarts, so a restarting engine reuses them instead of rebuilding them
5 +from scratch. The design is extensible to different kinds of artifacts; it
6 +currently supports **model weights** via the weight cache daemon.
7 +
8 +With weight preloading, a daemon process per GPU holds its rank's
9 +post-quantized, TP-sharded weights and serves CUDA IPC handles to vLLM engines
10 +over a Unix domain socket, so a restarting engine maps the weights via
11 +zero-copy IPC instead of reloading them from disk.
12 +
13 +Key benefits:
14 +
15 +- **Fast engine restarts**: weight loading reduces to mapping existing GPU
16 + memory, cutting restart time from minutes (disk load + quantization
17 + processing) to seconds.
18 +- **Zero-copy sharing**: in the default `zero_copy` mode the engine directly
19 + shares the daemon's GPU memory, so a restart adds no extra GPU memory
20 + footprint.
21 +- **Safe fallback**: if the daemon is unavailable or its cached weights don't
22 + match the engine's configuration, the engine falls back to loading from
23 + disk.
24 +
25 +!!! note
26 + Weight preloading is only supported on CUDA and ROCm platforms. Tensor,
27 + expert and data parallelism are supported, including across nodes;
28 + pipeline parallelism is rejected when launching the daemon.
29 +
30 +## Quick start
31 +
32 +Launch one weight cache daemon per GPU with a single command:
33 +
34 +```bash
35 +vllm preload --model meta-llama/Llama-3.1-8B-Instruct --tensor-parallel-size 4
36 +```
37 +
38 +The daemon process loads and post-processes the weights once, then waits for
39 +engines. Each rank binds its Unix socket only after its shard is fully cached,
40 +so engines connecting before that simply fall back to disk loading.
41 +
42 +Then start (or restart) engines with the `ipc_cache` load format:
43 +
44 +```bash
45 +vllm serve meta-llama/Llama-3.1-8B-Instruct \
46 + --tensor-parallel-size 4 \
47 + --load-format ipc_cache
48 +```
49 +
50 +The daemon itself must load from disk; passing `--load-format ipc_cache` to
51 +`vllm preload` is an error.
52 +
53 +## How it works
54 +
55 +1. `vllm preload` spawns one daemon process per GPU. Each daemon loads its
56 + shard from disk using the configured loader, runs quantization
57 + post-processing, and exports every parameter/buffer as a CUDA IPC handle.
58 +2. An engine started with `--load-format ipc_cache` builds its model on the
59 + meta device (no storage) and requests the tensors from the daemon on its
60 + GPU.
61 +3. Before serving anything, the engine and daemon compare a fingerprint of the
62 + cached weights: checkpoint content (hashed from safetensors metadata, so
63 + identical weights in different directories still match), model
64 + architecture, TP/DP size/rank, dtype, quantization method and config, model
65 + revision, and vLLM version. On any mismatch the engine falls back to disk
66 + loading (unless `fallback` is disabled).
67 +4. The engine also verifies the daemon's GPU UUID matches its own device, so a
68 + stale socket cannot serve weights for the wrong GPU.
69 +
70 +Tied weights (e.g. `lm_head.weight` sharing storage with
71 +`embed_tokens.weight`) are exported once and re-established as aliases in the
72 +engine, preserving parameter identity.
73 +
74 +## Multi-node and data parallelism
75 +
76 +CUDA IPC handles are node-local, so each node serves only its local GPUs'
77 +shards. For multi-node tensor parallelism, run one `vllm preload` launcher per
78 +node with a shared rendezvous: reuse the `--nnodes` / `--node-rank` /
79 +`--master-addr` flags you pass the engine, plus a `--weight-cache-master-port`
80 +distinct from the engine's `--master-port`:
81 +
82 +```bash
83 +# node 0 (8 local GPUs)
84 +vllm preload --model /path/to/model --tensor-parallel-size 16 \
85 + --nnodes 2 --node-rank 0 --master-addr 10.0.0.1 \
86 + --weight-cache-master-port 29600
87 +# node 1 (8 local GPUs)
88 +vllm preload --model /path/to/model --tensor-parallel-size 16 \
89 + --nnodes 2 --node-rank 1 --master-addr 10.0.0.1 \
90 + --weight-cache-master-port 29600
91 +```
92 +
93 +For data parallelism (e.g. a TP1 x DP16 x EP decode fleet), run one launcher
94 +per node with the engine's DP placement flags. Local GPU `i` serves DP rank
95 +`start_rank + i // tp_size` and TP rank `i % tp_size`, and all
96 +`dp_size * tp_size` daemons form one world group on `--data-parallel-address`
97 +/ `--weight-cache-master-port` so the expert shards are laid out exactly as in
98 +the engine:
99 +
100 +```bash
101 +# node r (4 local GPUs)
102 +vllm preload --model /path/to/model --tensor-parallel-size 1 \
103 + --enable-expert-parallel \
104 + --data-parallel-size 16 --data-parallel-size-local 4 \
105 + --data-parallel-start-rank 4r --data-parallel-address 10.0.0.1 \
106 + --weight-cache-master-port 29600
107 +```
108 +
109 +Data parallelism also combines with multi-node tensor parallelism: pass both
110 +flag sets. Each node then serves a contiguous block of the
111 +`dp_size * tp_size` global ranks.
112 +
113 +## Speculative decoding
114 +
115 +With MTP, EAGLE or EAGLE3 speculative decoding, `vllm preload` additionally
116 +starts a draft daemon group that caches the draft model. It uses its own cache
117 +key, Unix sockets (`*_draft.sock`) and rendezvous port
118 +(`--weight-cache-draft-master-port`, default `--weight-cache-master-port + 1`),
119 +so each daemon process serves exactly one model role. Other draft types are
120 +not cached and keep loading from disk in the engine.
121 +
122 +## Cache modes
123 +
124 +The loader supports two modes, selected via `--model-loader-extra-config`:
125 +
126 +- `zero_copy` (default): the engine maps the daemon's CUDA allocations
127 + directly. The daemon must stay alive for the engine's lifetime, and each
128 + GPU carries only one copy of the weights.
129 +- `copy`: the engine clones every tensor into its own GPU memory and then
130 + asks the daemon to release its cache. Use this when the daemon should free
131 + GPU memory after handing off, at the cost of a full copy per restart.
132 +
133 +!!! warning
134 + In `zero_copy` mode the weights live in the daemon's CUDA IPC allocations,
135 + so [sleep mode](sleep_mode.md) (CuMemAllocator weight offloading) must not
136 + be used with this loader.
137 +
138 +## Loader configuration
139 +
140 +The `ipc_cache` loader accepts extra keys via `--model-loader-extra-config`:
141 +
142 +```bash
143 +vllm serve meta-llama/Llama-3.1-8B-Instruct \
144 + --load-format ipc_cache \
145 + --model-loader-extra-config '{"mode": "copy", "fallback": false}'
146 +```
147 +
148 +| Key | Default | Description |
149 +| --- | ------- | ----------- |
150 +| `socket_path` | per-GPU path derived from the GPU UUID | Explicit daemon socket path. |
151 +| `socket_dir` | per-user private dir under the temp dir | Directory containing the daemon sockets. |
152 +| `mode` | `zero_copy` | `zero_copy` or `copy` (see above). |
153 +| `fallback` | `true` | Fall back to disk loading when the daemon is unavailable or the fingerprints mismatch. |
154 +| `connect_timeout_s` | `5.0` | Socket connect timeout in seconds. |
155 +| `state_timeout_s` | `300.0` | Timeout for the weight-transfer request in seconds. |
156 +
157 +The daemon's `--weight-cache-socket-dir` and the loader's `socket_dir` /
158 +`socket_path` must agree when the default per-user directory is not used.
159 +Socket paths are derived from the GPU UUID, so they are stable regardless of
160 +`CUDA_VISIBLE_DEVICES` index remapping.
161 +
162 +## Limitations
163 +
164 +- **Platform**: CUDA and ROCm only; other platforms raise
165 + `UnsupportedPlatformForIPCError` even when `fallback` is enabled, since it
166 + is a permanent misconfiguration rather than a transient daemon outage.
167 +- **Parallelism**: tensor, expert and data parallelism are supported;
168 + launching the daemon with pipeline parallelism is rejected.
169 +- **Quantization**: every quantization method in the model must declare
170 + support for pre-processed weights (the daemon transfers weights *after*
171 + quantization post-processing). Unsupported methods raise
172 + `UnsupportedQuantForIPCError`.
173 +- **Consistency**: the daemon and engine must run the same vLLM version and
174 + agree on model, dtype, quantization, and TP/DP layout, otherwise the
175 + fingerprint mismatch triggers the disk fallback.
176 +- **Localhost only**: the daemon serves over a Unix domain socket; both the
177 + daemon and the engines must run on the same node as the same user.
178 +
179 +## Security
180 +
181 +The socket protocol uses pickle and is intended only for trusted local
182 +processes owned by the same user:
183 +
184 +- Daemon sockets live in a per-user private directory (mode `0700`) and the
185 + socket files are restricted to the owner (`0600`).
186 +- Both sides reject symlinked or non-owned socket paths; the auto-derived
187 + directory is additionally rejected if it is group/world accessible.
188 +- On Linux the daemon verifies the connecting peer's UID via `SO_PEERCRED`.
189 +
190 +See the [security documentation](../usage/security.md) for vLLM's general
191 +threat model.
192 +
193 +## Daemon lifecycle
194 +
195 +- Each GPU's socket path is guarded by an exclusive lock file, so a second
196 + daemon for the same GPU fails fast instead of clobbering the live socket.
197 +- `vllm preload` reports readiness once every rank is serving; if any rank
198 + dies during startup, the remaining ranks are terminated and the command
199 + exits with that rank's exit code.
200 +- `SIGINT`/`SIGTERM` to the `vllm preload` process terminates all daemon
201 + ranks, which unlinks their sockets and releases the GPU memory.
202 +
203 +See [vllm preload](../cli/preload.md) for the full CLI reference.
docs/mkdocs/gen_files/generate_argparse.py
+2
@@ -177,6 +177,7 @@ BenchmarkMMProcessorSubcommand = auto_mock(
177 LaunchSubcommandBase = auto_mock("vllm.entrypoints.cli.launch", "LaunchSubcommandBase")
178 launch_description = auto_mock("vllm.entrypoints.cli.launch", "DESCRIPTION")
179 RenderSubcommand = auto_mock("vllm.entrypoints.cli.launch", "RenderSubcommand")
180 +PreloadSubcommand = auto_mock("vllm.entrypoints.cli.preload", "PreloadSubcommand")
181 sweep_subcommands = auto_mock("vllm.benchmarks.sweep.cli", "SUBCOMMANDS")
182 openai_cli_args = auto_mock("vllm.entrypoints.launchers", "cli_args")
183 run_batch = auto_mock("vllm.entrypoints.launchers", "run_batch")
@@ -330,6 +331,7 @@ pages = {
331 "cli/chat.md": (create_parser(ChatCommand.add_cli_args), False),
332 "cli/complete.md": (create_parser(CompleteCommand.add_cli_args), False),
333 "cli/run-batch.md": (create_parser(run_batch.make_arg_parser), True),
334 + "cli/preload.md": (create_parser(PreloadSubcommand.add_cli_args), True),
335 "cli/launch/render.md": (create_parser(RenderSubcommand.add_cli_args), True),
336 "cli/bench/latency.md": (create_parser(bench_latency.add_cli_args), True),
337 # URL kept as `mm_processor` for back-compat; command name is `mm-processor`
tests/model_executor/model_loader/test_weight_cache.py
+2 -1
@@ -36,7 +36,8 @@ class WeightCacheDaemon:
36 self._cmd = [
37 sys.executable,
38 "-m",
39 - "vllm.model_executor.model_loader.weight_cache.daemon",
39 + "vllm.entrypoints.cli.main",
40 + "preload",
41 "--model",
42 model,
43 "--tensor-parallel-size",
vllm/config/load.py
+2 -3
@@ -39,9 +39,8 @@ class LoadConfig:
39 InstantTensor, which enables distributed loading with pipelined prefetching
40 and fast direct I/O.
41 - "ipc_cache" will map post-quantized weights from a local weight cache
42 - daemon via CUDA IPC for fast engine restarts. See
43 - `vllm/model_executor/model_loader/weight_cache/daemon.py` for how to
44 - launch the daemon.
42 + daemon via CUDA IPC for fast engine restarts. Launch the daemon with
43 + `vllm preload` first.
44 - "npcache" will load the weights in pytorch format and store a numpy cache
45 to speed up the loading.
46 - "dummy" will initialize the weights with random values, which is mainly
vllm/entrypoints/cli/main.py
+2
@@ -44,6 +44,7 @@ def main():
44 import vllm.entrypoints.cli.collect_env
45 import vllm.entrypoints.cli.launch
46 import vllm.entrypoints.cli.openai
47 + import vllm.entrypoints.cli.preload
48 import vllm.entrypoints.cli.run_batch
49 import vllm.entrypoints.cli.serve
50 import vllm.entrypoints.cli.snapshot
@@ -59,6 +60,7 @@ def main():
60 vllm.entrypoints.cli.launch,
61 vllm.entrypoints.cli.benchmark.main,
62 vllm.entrypoints.cli.collect_env,
63 + vllm.entrypoints.cli.preload,
64 vllm.entrypoints.cli.run_batch,
65 vllm.entrypoints.cli.snapshot,
66 ]
vllm/entrypoints/cli/preload.py new
+252
@@ -0,0 +1,252 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""The `preload` subcommand for the vLLM CLI.
4 +
5 +Launches one weight cache daemon per GPU; each daemon holds the
6 +post-quantized, TP-sharded weights of its rank in GPU memory and serves
7 +CUDA IPC handles to vLLM engines over a Unix domain socket, so restarting
8 +engines can map the weights via zero-copy IPC instead of reloading from disk:
9 +
10 + vllm preload --model /path/to/model --tensor-parallel-size 4
11 +
12 +Engines then load from the daemons with:
13 +
14 + vllm serve /path/to/model --tensor-parallel-size 4 --load-format ipc_cache
15 +
16 +Tensor, expert and data parallelism are supported, including across nodes;
17 +see vllm/model_executor/model_loader/weight_cache/daemon.py for the rank
18 +layout and rendezvous details.
19 +"""
20 +
21 +import argparse
22 +import multiprocessing
23 +import queue
24 +import signal
25 +import sys
26 +import typing
27 +from itertools import product
28 +
29 +from vllm.engine.arg_utils import EngineArgs
30 +from vllm.entrypoints.cli.types import CLISubcommand
31 +from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG
32 +from vllm.logger import init_logger
33 +from vllm.model_executor.model_loader.weight_cache.protocol import (
34 + check_ipc_platform_support,
35 +)
36 +from vllm.utils.network_utils import get_distributed_init_method, get_open_port
37 +
38 +if typing.TYPE_CHECKING:
39 + from vllm.config import VllmConfig
40 + from vllm.utils.argparse_utils import FlexibleArgumentParser
41 +else:
42 + FlexibleArgumentParser = argparse.ArgumentParser
43 +
44 +logger = init_logger(__name__)
45 +
46 +
47 +class PreloadSubcommand(CLISubcommand):
48 + """The `preload` subcommand for the vLLM CLI."""
49 +
50 + name = "preload"
51 +
52 + @staticmethod
53 + def add_cli_args(parser: FlexibleArgumentParser) -> None:
54 + EngineArgs.add_cli_args(parser)
55 + parser.add_argument(
56 + "--weight-cache-socket-dir",
57 + type=str,
58 + default=None,
59 + help="Directory for the daemon Unix sockets (default: tempdir).",
60 + )
61 + parser.add_argument(
62 + "--weight-cache-master-port",
63 + type=int,
64 + default=None,
65 + help="Rendezvous port for the daemon's own TP group. Must differ "
66 + "from the engine's --master-port (the daemon holds its group open "
67 + "while serving) and match across nodes. Required when --nnodes > "
68 + "1; defaults to a free port for single-node.",
69 + )
70 + parser.add_argument(
71 + "--weight-cache-draft-master-port",
72 + type=int,
73 + default=None,
74 + help="Rendezvous port for the MTP draft daemon group. Defaults to "
75 + "--weight-cache-master-port + 1 for multi-node, or a free port "
76 + "for single-node.",
77 + )
78 +
79 + @staticmethod
80 + def cmd(args: argparse.Namespace) -> None:
81 + # Imported here so `vllm --help` and friends stay free of torch.
82 + from vllm.model_executor.model_loader.weight_cache.daemon import (
83 + _reject_unsupported_parallelism,
84 + _run_daemon,
85 + get_draft_daemon_config,
86 + plan_local_ranks,
87 + )
88 + from vllm.model_executor.model_loader.weight_cache.utils import (
89 + format_daemon_role,
90 + )
91 +
92 + engine_args = EngineArgs.from_cli_args(args)
93 + vllm_config = engine_args.create_engine_config()
94 + if vllm_config.load_config.load_format == "ipc_cache":
95 + raise ValueError(
96 + "The weight cache daemon itself must load from disk; use the "
97 + "default --load-format"
98 + )
99 + # Config-only so it can fail before any model loading; the quant
100 + # method check needs the created model and runs in
101 + # WeightCacheDaemon.get_model.
102 + check_ipc_platform_support()
103 + parallel_config = vllm_config.parallel_config
104 + _reject_unsupported_parallelism(parallel_config)
105 + tp_size = parallel_config.tensor_parallel_size
106 + dp_size = parallel_config.data_parallel_size
107 +
108 + placements = plan_local_ranks(parallel_config)
109 + local_world_size = len(placements)
110 + if dp_size == 1:
111 + nnodes = parallel_config.nnodes
112 + node_rank = parallel_config.node_rank
113 + master_addr = parallel_config.master_addr
114 + else:
115 + master_addr = parallel_config.data_parallel_master_ip
116 + if parallel_config.nnodes > 1:
117 + nnodes = parallel_config.nnodes
118 + node_rank = parallel_config.node_rank
119 + else:
120 + nnodes = dp_size // parallel_config.data_parallel_size_local
121 + node_rank = parallel_config.data_parallel_rank // (
122 + parallel_config.data_parallel_size_local
123 + )
124 + if nnodes > 1 and master_addr in ("127.0.0.1", "localhost"):
125 + raise ValueError(
126 + "Data parallelism across nodes requires a reachable "
127 + "--data-parallel-address for the daemon rendezvous"
128 + )
129 +
130 + # The daemon forms its own world group and holds it open while
131 + # serving, so it needs a rendezvous port distinct from the engine's.
132 + # All nodes must agree on it; single-node can auto-pick a free port.
133 + # The master address is the engine's --master-addr (TP across nodes)
134 + # or --data-parallel-address.
135 + if nnodes > 1 and args.weight_cache_master_port is None:
136 + raise ValueError(
137 + "--weight-cache-master-port is required when the daemons span nodes"
138 + )
139 + master_port = args.weight_cache_master_port or get_open_port()
140 + distributed_init_method = get_distributed_init_method(master_addr, master_port)
141 +
142 + # (is_draft, vllm_config, rendezvous) per daemon group.
143 + groups: list[tuple[bool, VllmConfig, str]] = [
144 + (False, vllm_config, distributed_init_method)
145 + ]
146 + draft_vllm_config = get_draft_daemon_config(vllm_config)
147 + if draft_vllm_config is not None:
148 + draft_master_port = args.weight_cache_draft_master_port or (
149 + master_port + 1 if nnodes > 1 else get_open_port()
150 + )
151 + if draft_master_port == master_port:
152 + raise ValueError(
153 + "--weight-cache-draft-master-port must differ from "
154 + "--weight-cache-master-port"
155 + )
156 + groups.append(
157 + (
158 + True,
159 + draft_vllm_config,
160 + get_distributed_init_method(master_addr, draft_master_port),
161 + )
162 + )
163 +
164 + ctx = multiprocessing.get_context("spawn")
165 + ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue()
166 + # Local index == device index; global rank enumerates DP then TP.
167 + expected_ready = {
168 + (format_daemon_role(is_draft), dp_rank * tp_size + tp_rank)
169 + for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements)
170 + }
171 + procs = [
172 + ctx.Process(
173 + target=_run_daemon,
174 + args=(
175 + tp_rank,
176 + local_rank,
177 + config,
178 + init_method,
179 + args.weight_cache_socket_dir,
180 + ready_queue,
181 + is_draft,
182 + dp_rank,
183 + ),
184 + name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-"
185 + f"{dp_rank * tp_size + tp_rank}",
186 + )
187 + for (is_draft, config, init_method), (
188 + local_rank,
189 + dp_rank,
190 + tp_rank,
191 + ) in product(groups, placements)
192 + ]
193 + for proc in procs:
194 + proc.start()
195 +
196 + def _shutdown(signum, frame):
197 + for proc in procs:
198 + proc.terminate()
199 +
200 + signal.signal(signal.SIGINT, _shutdown)
201 + signal.signal(signal.SIGTERM, _shutdown)
202 +
203 + ready: set[tuple[str, int]] = set()
204 + while len(ready) < len(expected_ready):
205 + try:
206 + ready.add(ready_queue.get(timeout=1.0))
207 + except queue.Empty:
208 + dead = [p for p in procs if p.exitcode is not None]
209 + if dead:
210 + logger.error(
211 + "Weight cache daemon rank(s) exited during startup "
212 + "(exitcodes=%s); shutting down.",
213 + [p.exitcode for p in dead],
214 + )
215 + for proc in procs:
216 + proc.terminate()
217 + for proc in procs:
218 + proc.join()
219 + sys.exit(max((p.exitcode or 0) for p in procs))
220 + socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
221 + logger.info_once(
222 + "===== Weight cache daemon READY: node %d/%d serving %d local "
223 + "rank(s) x %d role(s) in %s =====",
224 + node_rank,
225 + nnodes,
226 + local_world_size,
227 + len(groups),
228 + socket_dir_msg,
229 + )
230 +
231 + for proc in procs:
232 + proc.join()
233 + sys.exit(max(proc.exitcode or 0 for proc in procs))
234 +
235 + def subparser_init(
236 + self, subparsers: argparse._SubParsersAction
237 + ) -> FlexibleArgumentParser:
238 + preload_parser = subparsers.add_parser(
239 + self.name,
240 + help="Launch weight cache daemons (one per GPU) for fast engine restarts.",
241 + description="Launch weight cache daemons (one per GPU) that "
242 + "keep post-quantized weights in GPU memory and serve them to "
243 + "engines over CUDA IPC.",
244 + usage="vllm preload --model <model> [options]",
245 + )
246 + self.add_cli_args(preload_parser)
247 + preload_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name)
248 + return preload_parser
249 +
250 +
251 +def cmd_init() -> list[CLISubcommand]:
252 + return [PreloadSubcommand()]
vllm/model_executor/model_loader/weight_cache/daemon.py
+5 -188
@@ -9,7 +9,7 @@ reloading from disk.
9
10 Launch one daemon per TP rank with a single command:
11
12 - python -m vllm.model_executor.model_loader.weight_cache.daemon \\
12 + vllm preload \\
13 --model /path/to/model --tensor-parallel-size 4
14
15 Engines then load from the daemons with:
@@ -27,12 +27,12 @@ node-local, so each node serves only its local GPUs' shards). Reuse the same
27 ``--weight-cache-master-port`` distinct from the engine's ``--master-port``:
28
29 # node 0 (8 local GPUs)
30 - python -m vllm.model_executor.model_loader.weight_cache.daemon \\
30 + vllm preload \\
31 --model /path/to/model --tensor-parallel-size 16 \\
32 --nnodes 2 --node-rank 0 --master-addr 10.0.0.1 \\
33 --weight-cache-master-port 29600
34 # node 1 (8 local GPUs)
35 - python -m vllm.model_executor.model_loader.weight_cache.daemon \\
35 + vllm preload \\
36 --model /path/to/model --tensor-parallel-size 16 \\
37 --nnodes 2 --node-rank 1 --master-addr 10.0.0.1 \\
38 --weight-cache-master-port 29600
@@ -50,7 +50,7 @@ node with the engine's DP placement flags. Local GPU ``i`` serves DP rank
50 shards are laid out exactly as in the engine:
51
52 # node r (4 local GPUs)
53 - python -m vllm.model_executor.model_loader.weight_cache.daemon \\
53 + vllm preload \\
54 --model /path/to/model --tensor-parallel-size 1 --enable-expert-parallel \\
55 --data-parallel-size 16 --data-parallel-size-local 4 \\
56 --data-parallel-start-rank 4r --data-parallel-address 10.0.0.1 \\
@@ -62,7 +62,7 @@ a reachable ``--data-parallel-address``). Each node then serves a contiguous
62 block of the ``dp_size * tp_size`` global ranks, e.g. TP8 x DP2 on 4 nodes:
63
64 # node r (4 local GPUs)
65 - python -m vllm.model_executor.model_loader.weight_cache.daemon \\
65 + vllm preload \\
66 --model /path/to/model --tensor-parallel-size 8 --enable-expert-parallel \\
67 --nnodes 4 --node-rank r --master-addr 10.0.0.1 \\
68 --data-parallel-size 2 --data-parallel-address 10.0.0.1 \\
@@ -80,12 +80,8 @@ import contextlib
80 import fcntl
81 import multiprocessing
82 import os
83 -import queue
84 -import signal
83 import socket
86 -import sys
84 from collections.abc import Callable
88 -from itertools import product
85
86 import torch
87
@@ -99,7 +95,6 @@ from vllm.distributed import (
95 ensure_model_parallel_initialized,
96 init_distributed_environment,
97 )
102 -from vllm.engine.arg_utils import EngineArgs
98 from vllm.logger import init_logger
99 from vllm.model_executor.model_loader import get_model_loader
100 from vllm.model_executor.model_loader.utils import process_weights_after_loading
@@ -107,7 +102,6 @@ from vllm.model_executor.model_loader.weight_cache.protocol import (
102 TensorEntry,
103 WeightCacheKey,
104 WeightCacheUnavailableError,
110 - check_ipc_platform_support,
105 check_ipc_quant_support,
106 ensure_private_socket_dir,
107 get_current_device_uuid,
@@ -122,8 +116,6 @@ from vllm.model_executor.model_loader.weight_cache.utils import (
116 is_draft_model_cacheable,
117 )
118 from vllm.platforms import current_platform
125 -from vllm.utils.argparse_utils import FlexibleArgumentParser
126 -from vllm.utils.network_utils import get_distributed_init_method, get_open_port
119 from vllm.utils.torch_utils import set_default_torch_dtype
120 from vllm.v1.worker.workspace import init_workspace_manager
121
@@ -504,178 +496,3 @@ def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:
496 "--data-parallel-start-rank + --data-parallel-size-local "
497 f"({end_rank}) exceeds --data-parallel-size ({dp_size})"
498 )
507 -
508 -
509 -def main() -> None:
510 - parser = FlexibleArgumentParser(
511 - description="Launch weight cache daemons (one per TP rank)."
512 - )
513 - EngineArgs.add_cli_args(parser)
514 - parser.add_argument(
515 - "--weight-cache-socket-dir",
516 - type=str,
517 - default=None,
518 - help="Directory for the daemon Unix sockets (default: tempdir).",
519 - )
520 - parser.add_argument(
521 - "--weight-cache-master-port",
522 - type=int,
523 - default=None,
524 - help="Rendezvous port for the daemon's own TP group. Must differ from "
525 - "the engine's --master-port (the daemon holds its group open while "
526 - "serving) and match across nodes. Required when --nnodes > 1; defaults "
527 - "to a free port for single-node.",
528 - )
529 - parser.add_argument(
530 - "--weight-cache-draft-master-port",
531 - type=int,
532 - default=None,
533 - help="Rendezvous port for the MTP draft daemon group. Defaults to "
534 - "--weight-cache-master-port + 1 for multi-node, or a free port for "
535 - "single-node.",
536 - )
537 - args = parser.parse_args()
538 - engine_args = EngineArgs.from_cli_args(args)
539 - vllm_config = engine_args.create_engine_config()
540 - if vllm_config.load_config.load_format == "ipc_cache":
541 - raise ValueError(
542 - "The weight cache daemon itself must load from disk; use the "
543 - "default --load-format"
544 - )
545 - # Config-only so it can fail before any model loading; the quant method
546 - # check needs the created model and runs in get_daemon_model.
547 - check_ipc_platform_support()
548 - parallel_config = vllm_config.parallel_config
549 - _reject_unsupported_parallelism(parallel_config)
550 - tp_size = parallel_config.tensor_parallel_size
551 - dp_size = parallel_config.data_parallel_size
552 -
553 - placements = plan_local_ranks(parallel_config)
554 - local_world_size = len(placements)
555 - if dp_size == 1:
556 - nnodes = parallel_config.nnodes
557 - node_rank = parallel_config.node_rank
558 - master_addr = parallel_config.master_addr
559 - else:
560 - master_addr = parallel_config.data_parallel_master_ip
561 - if parallel_config.nnodes > 1:
562 - nnodes = parallel_config.nnodes
563 - node_rank = parallel_config.node_rank
564 - else:
565 - nnodes = dp_size // parallel_config.data_parallel_size_local
566 - node_rank = parallel_config.data_parallel_rank // (
567 - parallel_config.data_parallel_size_local
568 - )
569 - if nnodes > 1 and master_addr in ("127.0.0.1", "localhost"):
570 - raise ValueError(
571 - "Data parallelism across nodes requires a reachable "
572 - "--data-parallel-address for the daemon rendezvous"
573 - )
574 -
575 - # The daemon forms its own world group and holds it open while serving, so
576 - # it needs a rendezvous port distinct from the engine's. All nodes must
577 - # agree on it; single-node can auto-pick a free port. The master address is
578 - # the engine's --master-addr (TP across nodes) or --data-parallel-address.
579 - if nnodes > 1 and args.weight_cache_master_port is None:
580 - raise ValueError(
581 - "--weight-cache-master-port is required when the daemons span nodes"
582 - )
583 - master_port = args.weight_cache_master_port or get_open_port()
584 - distributed_init_method = get_distributed_init_method(master_addr, master_port)
585 -
586 - # (is_draft, vllm_config, rendezvous) per daemon group.
587 - groups: list[tuple[bool, VllmConfig, str]] = [
588 - (False, vllm_config, distributed_init_method)
589 - ]
590 - draft_vllm_config = get_draft_daemon_config(vllm_config)
591 - if draft_vllm_config is not None:
592 - draft_master_port = args.weight_cache_draft_master_port or (
593 - master_port + 1 if nnodes > 1 else get_open_port()
594 - )
595 - if draft_master_port == master_port:
596 - raise ValueError(
597 - "--weight-cache-draft-master-port must differ from "
598 - "--weight-cache-master-port"
599 - )
600 - groups.append(
601 - (
602 - True,
603 - draft_vllm_config,
604 - get_distributed_init_method(master_addr, draft_master_port),
605 - )
606 - )
607 -
608 - ctx = multiprocessing.get_context("spawn")
609 - ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue()
610 - # Local index == device index; global rank enumerates DP then TP.
611 - expected_ready = {
612 - (format_daemon_role(is_draft), dp_rank * tp_size + tp_rank)
613 - for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements)
614 - }
615 - procs = [
616 - ctx.Process(
617 - target=_run_daemon,
618 - args=(
619 - tp_rank,
620 - local_rank,
621 - config,
622 - init_method,
623 - args.weight_cache_socket_dir,
624 - ready_queue,
625 - is_draft,
626 - dp_rank,
627 - ),
628 - name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-"
629 - f"{dp_rank * tp_size + tp_rank}",
630 - )
631 - for (is_draft, config, init_method), (
632 - local_rank,
633 - dp_rank,
634 - tp_rank,
635 - ) in product(groups, placements)
636 - ]
637 - for proc in procs:
638 - proc.start()
639 -
640 - def _shutdown(signum, frame):
641 - for proc in procs:
642 - proc.terminate()
643 -
644 - signal.signal(signal.SIGINT, _shutdown)
645 - signal.signal(signal.SIGTERM, _shutdown)
646 -
647 - ready: set[tuple[str, int]] = set()
648 - while len(ready) < len(expected_ready):
649 - try:
650 - ready.add(ready_queue.get(timeout=1.0))
651 - except queue.Empty:
652 - dead = [p for p in procs if p.exitcode is not None]
653 - if dead:
654 - logger.error(
655 - "Weight cache daemon rank(s) exited during startup "
656 - "(exitcodes=%s); shutting down.",
657 - [p.exitcode for p in dead],
658 - )
659 - for proc in procs:
660 - proc.terminate()
661 - for proc in procs:
662 - proc.join()
663 - sys.exit(max((p.exitcode or 0) for p in procs))
664 - socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
665 - logger.info_once(
666 - "===== Weight cache daemon READY: node %d/%d serving %d local rank(s) "
667 - "x %d role(s) in %s =====",
668 - node_rank,
669 - nnodes,
670 - local_world_size,
671 - len(groups),
672 - socket_dir_msg,
673 - )
674 -
675 - for proc in procs:
676 - proc.join()
677 - sys.exit(max(proc.exitcode or 0 for proc in procs))
678 -
679 -
680 -if __name__ == "__main__":
681 - main()