A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 22:30 UTC
[Startup] Parallelize torch/transformers import + weight prefetch + forkserver prewarm (#40331)
Signed-off-by: simon-mo <simon@inferact.ai>
Simon Mo committed
Apr 20, 2026 at 19:49 UTC
8256833fe6f90e508b0264120757c5ea999a044d
3 files changed
+220
-6
vllm/entrypoints/cli/main.py
+71
-1
@@ -5,10 +5,80 @@
5
Note that all future modules must be lazily loaded within main
6
to avoid certain eager import breakage."""
7
8
+import contextlib
9
import importlib.metadata
10
+import os
11
import sys
12
+import threading as _threading
13
11
-from vllm.logger import init_logger
14
+
15
+# [startup] Kick off torch + transformers .so/module loading in a background
16
+# thread before we touch vllm.logger (which pulls vllm/__init__.py ->
17
+# vllm.env_override -> `import torch` on the main thread). Python import
18
+# lock serializes the same-module import across threads, but the .so dlopen
19
+# inside torch's init releases the GIL during file I/O. Main thread's
20
+# non-torch imports (vllm.envs submodules, stdlib, fastapi, etc.) can make
21
+# progress on the CPU while the background thread pays the ~2 s of cuda
22
+# .so loading. `import transformers` is also ~2 s of cold-disk work and
23
+# depends on torch; chain it after torch in the same thread so subsequent
24
+# `from transformers import ...` lines on the main thread hit a warm
25
+# module cache.
26
+def _bg_preload_torch() -> None:
27
+ try:
28
+ import torch # noqa: F401
29
+ except Exception:
30
+ return
31
+ with contextlib.suppress(Exception):
32
+ import transformers # noqa: F401
33
+
34
+
35
+_threading.Thread(
36
+ target=_bg_preload_torch, daemon=True, name="vllm-torch-preload"
37
+).start()
38
+
39
+
40
+# [startup] Pre-spawn EngineCore via forkserver preload, in a background
41
+# thread. Only fires for `vllm serve` (the only subcommand that spawns a
42
+# long-running EngineCore). The forkserver process is forked once and
43
+# preloaded with vllm.v1.engine.async_llm (~3-5 s of imports). When
44
+# AsyncLLM.from_vllm_config later runs, Process.start() forks from the
45
+# already-warm forkserver instead of paying spawn() cost (~5 s in child
46
+# for fresh Python + imports).
47
+#
48
+# Kicking the preload in a BG thread lets the ~3-5 s ensure_running cost
49
+# overlap with APIServer's argparse + config resolution (~5-10 s on cold
50
+# disk). Default cli_env_setup sets spawn; we override to forkserver
51
+# before that runs so the path is consistent.
52
+def _bg_prewarm_forkserver() -> None:
53
+ try:
54
+ import multiprocessing
55
+ import multiprocessing.forkserver as forkserver
56
+
57
+ # set_start_method MUST be called before ensure_running. It also
58
+ # can only be called once per process; any later override by
59
+ # vllm's build_async_engine_client will just see the existing
60
+ # setting.
61
+ multiprocessing.set_start_method("forkserver", force=False)
62
+ multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"])
63
+ forkserver.ensure_running()
64
+ except Exception:
65
+ pass
66
+
67
+
68
+if len(sys.argv) > 1 and sys.argv[1] == "serve":
69
+ os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "forkserver")
70
+ # daemon=True so early CLI exits (bad args, --help, import errors)
71
+ # don't hang waiting for ensure_running(). The forkserver subprocess
72
+ # itself is tracked by module-level state in multiprocessing.forkserver
73
+ # and survives this thread exiting; subsequent spawn() calls reuse it.
74
+ _threading.Thread(
75
+ target=_bg_prewarm_forkserver,
76
+ daemon=True,
77
+ name="vllm-forkserver-prewarm",
78
+ ).start()
79
+
80
+
81
+from vllm.logger import init_logger # noqa: E402
82
83
logger = init_logger(__name__)
84
vllm/entrypoints/openai/api_server.py
+142
-2
@@ -12,7 +12,7 @@ import tempfile
12
import warnings
13
from argparse import Namespace
14
from collections.abc import AsyncIterator
15
-from contextlib import asynccontextmanager
15
+from contextlib import asynccontextmanager, suppress
16
from typing import Any
17
18
import uvloop
@@ -74,6 +74,121 @@ logger = init_logger("vllm.entrypoints.openai.api_server")
74
_FALLBACK_SUPPORTED_TASKS: tuple[SupportedTask, ...] = ("generate",)
75
76
77
+def _startup_prefetch_weights(vllm_config: "VllmConfig") -> None:
78
+ """Kick off reading model weight shards into OS page cache from the
79
+ parent APIServer. EngineCore will read the same files a few seconds
80
+ later from the child; by then the kernel already has them ready.
81
+
82
+ All work (directory resolution, HF/ModelScope cache lookup, globbing,
83
+ and the reads themselves) runs inside the background thread so we do
84
+ not block the asyncio event loop.
85
+
86
+ Best-effort: any failure (unknown model location, permission, etc.) is
87
+ swallowed — vLLM's existing in-child prefetch then runs normally.
88
+ """
89
+ import threading
90
+
91
+ # Capture only the small scalar fields the thread needs. Avoid holding
92
+ # a reference to vllm_config (which contains unpicklable objects) for
93
+ # longer than necessary.
94
+ model_ref = vllm_config.model_config.model
95
+ revision = vllm_config.model_config.revision
96
+ download_dir = vllm_config.load_config.download_dir
97
+
98
+ def _prefetch_worker() -> None:
99
+ import glob
100
+ import os
101
+
102
+ from vllm import envs
103
+
104
+ candidate_dir: str | None = None
105
+
106
+ # 1. Local path?
107
+ if os.path.isdir(model_ref):
108
+ candidate_dir = model_ref
109
+ else:
110
+ # 2. HF / ModelScope repo id — resolve to the local cache
111
+ # snapshot dir using the same revision / cache_dir the weight
112
+ # loader will use, so we prefetch the right files.
113
+ try:
114
+ if envs.VLLM_USE_MODELSCOPE:
115
+ from modelscope.hub.snapshot_download import (
116
+ snapshot_download,
117
+ )
118
+
119
+ candidate_dir = snapshot_download(
120
+ model_id=model_ref,
121
+ revision=revision,
122
+ cache_dir=download_dir,
123
+ local_files_only=True,
124
+ )
125
+ else:
126
+ from huggingface_hub import snapshot_download
127
+
128
+ candidate_dir = snapshot_download(
129
+ repo_id=model_ref,
130
+ revision=revision,
131
+ cache_dir=download_dir,
132
+ allow_patterns=[
133
+ "*.safetensors",
134
+ "*.bin",
135
+ "*.json",
136
+ "*tokenizer*",
137
+ ],
138
+ local_files_only=True,
139
+ )
140
+ except Exception:
141
+ return # not cached yet or not a known repo id
142
+
143
+ if not candidate_dir or not os.path.isdir(candidate_dir):
144
+ return
145
+
146
+ # Weight shards: large files, read into page cache.
147
+ shard_paths = sorted(
148
+ glob.glob(os.path.join(candidate_dir, "*.safetensors"))
149
+ + glob.glob(os.path.join(candidate_dir, "*.bin"))
150
+ )
151
+ # Tokenizer/config sidecars: small, but re-opened in the child and
152
+ # add synchronous open+read latency when the disk is cold.
153
+ sidecar_paths = sorted(
154
+ glob.glob(os.path.join(candidate_dir, "*.json"))
155
+ + glob.glob(os.path.join(candidate_dir, "tokenizer.model"))
156
+ + glob.glob(os.path.join(candidate_dir, "*tokenizer*"))
157
+ )
158
+ shard_paths.extend(sidecar_paths)
159
+ if not shard_paths:
160
+ return
161
+
162
+ logger.debug(
163
+ "Parent-side weight prefetch starting for %d files in %s",
164
+ len(shard_paths),
165
+ candidate_dir,
166
+ )
167
+
168
+ # Match vLLM's in-child prefetch block size + thread count.
169
+ block_size = 16 * 1024 * 1024 # 16 MB
170
+ # Read shards in parallel across 8 worker threads (bounded) to
171
+ # saturate multi-spindle / multi-queue storage without thrashing.
172
+ from concurrent.futures import ThreadPoolExecutor
173
+
174
+ def read_one(p: str) -> None:
175
+ try:
176
+ with open(p, "rb") as f:
177
+ while f.read(block_size):
178
+ pass
179
+ except Exception:
180
+ pass
181
+
182
+ with ThreadPoolExecutor(max_workers=8) as pool:
183
+ list(pool.map(read_one, shard_paths))
184
+
185
+ threading.Thread(
186
+ target=_prefetch_worker,
187
+ daemon=True,
188
+ name="vllm-parent-weight-prefetch",
189
+ ).start()
190
+
191
+
192
@asynccontextmanager
193
async def build_async_engine_client(
194
args: Namespace,
@@ -85,7 +200,10 @@ async def build_async_engine_client(
200
# The executor is expected to be mp.
201
# Pre-import heavy modules in the forkserver process
202
logger.debug("Setup forkserver with pre-imports")
88
- multiprocessing.set_start_method("forkserver")
203
+ # May already have been set by the CLI entry's async prewarm
204
+ # (vllm/entrypoints/cli/main.py); tolerate re-call.
205
+ with suppress(RuntimeError):
206
+ multiprocessing.set_start_method("forkserver", force=False)
207
multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"])
208
forkserver.ensure_running()
209
logger.debug("Forkserver setup complete!")
@@ -123,6 +241,28 @@ async def build_async_engine_client_from_engine_args(
241
# Create the EngineConfig (determines if we can use V1).
242
vllm_config = engine_args.create_engine_config(usage_context=usage_context)
243
244
+ # [startup] Start prefetching model weight shards into the OS page cache
245
+ # in a background thread from the PARENT APIServer process. EngineCore
246
+ # will page-fault on these same files ~10-15 s later (after fork + CUDA
247
+ # context + distributed init + model init). For large-weight cases
248
+ # (tens of GB) this parent-side head start meaningfully shrinks the
249
+ # prefetch+load phase that the engine's in-child prefetch otherwise
250
+ # barely overlaps.
251
+ #
252
+ # Skip in API-only workers that connect to an already-running EngineCore
253
+ # (multi-API-server / disaggregated setups): those processes never load
254
+ # weights, and if we prefetched from all of them we'd contend with the
255
+ # engine's own read. Presence of an `input_address` in client_config is
256
+ # the current marker that this worker is headless.
257
+ #
258
+ # Best-effort: if the model is a local path, glob for safetensors; if
259
+ # it's a repo-id, try to resolve via HF hub (or ModelScope) local cache.
260
+ # Any failure silently falls through to the existing in-child prefetch
261
+ # path. All I/O (incl. directory resolution) runs inside the BG thread
262
+ # so the asyncio event loop is never blocked.
263
+ if not (client_config and client_config.get("input_address")):
264
+ _startup_prefetch_weights(vllm_config)
265
+
266
from vllm.v1.engine.async_llm import AsyncLLM
267
268
async_llm: AsyncLLM | None = None
vllm/envs.py
+7
-3
@@ -62,7 +62,7 @@ if TYPE_CHECKING:
62
VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
63
VLLM_USE_RAY_V2_EXECUTOR_BACKEND: bool = False
64
VLLM_XLA_USE_SPMD: bool = False
65
- VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "fork"
65
+ VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn", "forkserver"] = "fork"
66
VLLM_ASSETS_CACHE: str = os.path.join(VLLM_CACHE_ROOT, "assets")
67
VLLM_ASSETS_CACHE_MODEL_CLEAN: bool = False
68
VLLM_IMAGE_FETCH_TIMEOUT: int = 5
@@ -765,9 +765,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
765
int(os.getenv("VLLM_USE_RAY_V2_EXECUTOR_BACKEND", "0"))
766
),
767
# Use dedicated multiprocess context for workers.
768
- # Both spawn and fork work
768
+ # spawn, fork, and forkserver all work. forkserver is opt-in for fast
769
+ # startup when paired with the CLI's async prewarm (see
770
+ # vllm/entrypoints/cli/main.py) — the forkserver process is preloaded
771
+ # with vllm.v1.engine.async_llm and a subsequent EngineCore Process.start()
772
+ # forks from that warm sibling instead of paying spawn cost.
773
"VLLM_WORKER_MULTIPROC_METHOD": env_with_choices(
770
- "VLLM_WORKER_MULTIPROC_METHOD", "fork", ["spawn", "fork"]
774
+ "VLLM_WORKER_MULTIPROC_METHOD", "fork", ["spawn", "fork", "forkserver"]
775
),
776
# Path to the cache for storing downloaded assets
777
"VLLM_ASSETS_CACHE": lambda: os.path.expanduser(