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, 21:15 UTC
Revert "[Startup] Parallelize torch/transformers import + weight prefetch + forkserver prewarm" (#40438)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Apr 21, 2026 at 16:47 UTC
3975eb6de6ea914b9d7b27fd517e0c971ddeb6fc
3 files changed
+6
-220
vllm/entrypoints/cli/main.py
+1
-71
@@ -5,80 +5,10 @@
5
Note that all future modules must be lazily loaded within main
6
to avoid certain eager import breakage."""
7
8
-import contextlib
8
import importlib.metadata
10
-import os
9
import sys
12
-import threading as _threading
10
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
11
+from vllm.logger import init_logger
12
13
logger = init_logger(__name__)
14
vllm/entrypoints/openai/api_server.py
+2
-142
@@ -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, suppress
15
+from contextlib import asynccontextmanager
16
from typing import Any
17
18
import uvloop
@@ -74,121 +74,6 @@ 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
-
77
@asynccontextmanager
78
async def build_async_engine_client(
79
args: Namespace,
@@ -200,10 +85,7 @@ async def build_async_engine_client(
85
# The executor is expected to be mp.
86
# Pre-import heavy modules in the forkserver process
87
logger.debug("Setup forkserver with pre-imports")
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)
88
+ multiprocessing.set_start_method("forkserver")
89
multiprocessing.set_forkserver_preload(["vllm.v1.engine.async_llm"])
90
forkserver.ensure_running()
91
logger.debug("Forkserver setup complete!")
@@ -241,28 +123,6 @@ async def build_async_engine_client_from_engine_args(
123
# Create the EngineConfig (determines if we can use V1).
124
vllm_config = engine_args.create_engine_config(usage_context=usage_context)
125
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
-
126
from vllm.v1.engine.async_llm import AsyncLLM
127
128
async_llm: AsyncLLM | None = None
vllm/envs.py
+3
-7
@@ -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", "forkserver"] = "fork"
65
+ VLLM_WORKER_MULTIPROC_METHOD: Literal["fork", "spawn"] = "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,13 +765,9 @@ 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
- # 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.
768
+ # Both spawn and fork work
769
"VLLM_WORKER_MULTIPROC_METHOD": env_with_choices(
774
- "VLLM_WORKER_MULTIPROC_METHOD", "fork", ["spawn", "fork", "forkserver"]
770
+ "VLLM_WORKER_MULTIPROC_METHOD", "fork", ["spawn", "fork"]
771
),
772
# Path to the cache for storing downloaded assets
773
"VLLM_ASSETS_CACHE": lambda: os.path.expanduser(