[Core] Explicitly manage torch CPU threads in workers (#49919)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Nick Hill committed
Aug 4, 2026 at 10:06 UTC
7635a9002baecca64909dbcf8b1d461d26fb879d
4 files changed
+207
-25
tests/utils_/test_torch_utils.py
+50
@@ -4,9 +4,13 @@ import pytest
4
import torch
5
6
from vllm.utils.torch_utils import (
7
+ OMP_NUM_THREADS_SET_BY_VLLM,
8
+ available_cpu_count,
9
common_broadcastable_dtype,
10
current_stream,
11
is_lossless_cast,
12
+ set_torch_threads_for_runtime,
13
+ startup_omp_num_threads,
14
)
15
16
@@ -114,3 +118,49 @@ def test_current_stream_multithread():
118
)
119
120
_test_stream_thread(main_dedicated_stream)
121
+
122
+
123
+@pytest.fixture
124
+def restore_torch_threads(monkeypatch: pytest.MonkeyPatch):
125
+ monkeypatch.delenv("OMP_NUM_THREADS", raising=False)
126
+ original = torch.get_num_threads()
127
+ yield
128
+ torch.set_num_threads(original)
129
+
130
+
131
+def test_startup_omp_num_threads_divides_between_local_workers():
132
+ """Workers share the node's usable CPUs rather than each taking them all."""
133
+ available = available_cpu_count()
134
+ if available < 4:
135
+ pytest.skip("needs at least 4 usable CPUs")
136
+ assert startup_omp_num_threads(1) == available
137
+ assert startup_omp_num_threads(2) == available // 2
138
+ # Never zero, however many workers share the node.
139
+ assert startup_omp_num_threads(available * 4) == 1
140
+
141
+
142
+def test_set_torch_threads_for_runtime(restore_torch_threads):
143
+ torch.set_num_threads(max(2, available_cpu_count()))
144
+ set_torch_threads_for_runtime()
145
+ assert torch.get_num_threads() == 1
146
+
147
+
148
+def test_runtime_threads_respect_user_omp_num_threads(
149
+ restore_torch_threads, monkeypatch: pytest.MonkeyPatch
150
+):
151
+ """An externally-set OMP_NUM_THREADS is the user's choice; leave it alone."""
152
+ monkeypatch.setenv("OMP_NUM_THREADS", "3")
153
+ torch.set_num_threads(3)
154
+ set_torch_threads_for_runtime()
155
+ assert torch.get_num_threads() == 3
156
+
157
+
158
+def test_runtime_threads_override_vllm_set_omp_num_threads(
159
+ restore_torch_threads, monkeypatch: pytest.MonkeyPatch
160
+):
161
+ """The value vLLM picked for worker startup is dropped once serving starts."""
162
+ monkeypatch.setenv("OMP_NUM_THREADS", "3")
163
+ monkeypatch.setenv(OMP_NUM_THREADS_SET_BY_VLLM, "1")
164
+ torch.set_num_threads(3)
165
+ set_torch_threads_for_runtime()
166
+ assert torch.get_num_threads() == 1
vllm/utils/torch_utils.py
+116
@@ -4,6 +4,7 @@ import contextlib
4
import importlib.metadata
5
import os
6
import random
7
+import sys
8
import threading
9
from collections.abc import Callable, Collection
10
from typing import TYPE_CHECKING, Any, TypeVar
@@ -150,6 +151,121 @@ def set_default_torch_dtype(dtype: torch.dtype):
151
torch.set_default_dtype(old_dtype)
152
153
154
+def _cgroup_cpu_limit() -> float | None:
155
+ """Effective CPU quota of this process's cgroup, None if unlimited.
156
+
157
+ Resolves the process's own cgroup from /proc/self/cgroup and takes the
158
+ tightest cpu.max (v2) or cfs quota (v1) along the hierarchy.
159
+ """
160
+ limit: float | None = None
161
+ try:
162
+ with open("/proc/self/cgroup") as f:
163
+ entries = [line.strip().split(":", 2) for line in f]
164
+
165
+ def visit(base: str, rel_path: str, read_quota) -> None:
166
+ nonlocal limit
167
+ path = rel_path
168
+ while path:
169
+ quota = read_quota(os.path.join(base, path.lstrip("/")))
170
+ if quota is not None:
171
+ limit = quota if limit is None else min(limit, quota)
172
+ path = path.rsplit("/", 1)[0]
173
+
174
+ def read_v2(cg_dir: str) -> float | None:
175
+ try:
176
+ with open(os.path.join(cg_dir, "cpu.max")) as f:
177
+ quota, period = f.read().split()
178
+ return None if quota == "max" else float(quota) / float(period)
179
+ except (OSError, ValueError):
180
+ return None
181
+
182
+ def read_v1(cg_dir: str) -> float | None:
183
+ try:
184
+ with open(os.path.join(cg_dir, "cpu.cfs_quota_us")) as f:
185
+ quota = int(f.read())
186
+ if quota <= 0:
187
+ return None
188
+ with open(os.path.join(cg_dir, "cpu.cfs_period_us")) as f:
189
+ return quota / int(f.read())
190
+ except (OSError, ValueError):
191
+ return None
192
+
193
+ for entry in entries:
194
+ if len(entry) != 3:
195
+ continue
196
+ _, controllers, rel_path = entry
197
+ if controllers == "": # cgroup v2
198
+ visit("/sys/fs/cgroup", rel_path, read_v2)
199
+ elif "cpu" in controllers.split(","): # cgroup v1
200
+ visit("/sys/fs/cgroup/cpu", rel_path, read_v1)
201
+ except OSError:
202
+ pass
203
+ return limit
204
+
205
+
206
+def available_cpu_count() -> int:
207
+ """CPUs actually usable by this process: scheduling affinity capped by
208
+ the cgroup CPU quota (unlike `os.cpu_count()`, which is quota-blind)."""
209
+ if sys.platform != "linux":
210
+ return os.cpu_count() or 1
211
+ count = len(os.sched_getaffinity(0))
212
+ limit = _cgroup_cpu_limit()
213
+ if limit is not None:
214
+ count = min(count, int(limit))
215
+ return max(1, count)
216
+
217
+
218
+# Marks OMP_NUM_THREADS as chosen by vLLM for its worker processes rather than
219
+# set by the user, so a worker knows it may drop the value once startup is done.
220
+OMP_NUM_THREADS_SET_BY_VLLM = "VLLM_OMP_NUM_THREADS_SET_BY_VLLM"
221
+
222
+
223
+def startup_omp_num_threads(num_local_procs: int) -> int:
224
+ """Thread count for a worker process's startup work (weight loading).
225
+
226
+ Weight loading does CPU-parallel work, so workers benefit from more than
227
+ one thread, but only a bounded share of the CPUs this node's workers may
228
+ actually use: torch's default is the host core count, which ignores both
229
+ scheduling affinity and any cgroup CPU quota, and doesn't account for the
230
+ other workers sharing the node.
231
+ """
232
+ return max(1, available_cpu_count() // max(1, num_local_procs))
233
+
234
+
235
+def set_torch_threads_for_runtime() -> None:
236
+ """Set torch intra-op threads to 1 for steady-state serving.
237
+
238
+ Any multi-threaded torch CPU op in the engine hot loop leaves the OMP
239
+ workers spin-waiting after each parallel region, stealing cycles from the
240
+ step's serial code (and burning cgroup CPU quota in containers). No
241
+ steady-state CPU op benefits from intra-op parallelism.
242
+ Respects an externally-set OMP_NUM_THREADS.
243
+ """
244
+ if (
245
+ omp_num_threads := os.environ.get("OMP_NUM_THREADS")
246
+ ) is not None and os.environ.get(OMP_NUM_THREADS_SET_BY_VLLM) != "1":
247
+ try:
248
+ if int(omp_num_threads) > 1:
249
+ logger.warning_once(
250
+ "OMP_NUM_THREADS=%s is set; leaving Torch threads at %d "
251
+ "for serving. Multi-threaded torch CPU ops during serving "
252
+ "can degrade performance through spin-wait contention and "
253
+ "cgroup CPU-quota throttling.",
254
+ omp_num_threads,
255
+ torch.get_num_threads(),
256
+ )
257
+ except ValueError:
258
+ pass
259
+ return
260
+ if torch.get_num_threads() != 1:
261
+ logger.info_once(
262
+ "Reducing Torch threads from %d to 1 for serving. Set "
263
+ "OMP_NUM_THREADS in the external environment to override.",
264
+ torch.get_num_threads(),
265
+ )
266
+ torch.set_num_threads(1)
267
+
268
+
269
@contextlib.contextmanager
270
def set_default_torch_num_threads(num_threads: int | None = None):
271
"""
vllm/v1/executor/multiproc_executor.py
+36
-24
@@ -58,6 +58,11 @@ from vllm.utils.system_utils import (
58
get_mp_context,
59
set_process_title,
60
)
61
+from vllm.utils.torch_utils import (
62
+ OMP_NUM_THREADS_SET_BY_VLLM,
63
+ set_torch_threads_for_runtime,
64
+ startup_omp_num_threads,
65
+)
66
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
67
from vllm.v1.executor.abstract import Executor, FailureCallback
68
from vllm.v1.executor.vllm_net_devices import set_worker_net_device
@@ -122,7 +127,7 @@ class MultiprocExecutor(Executor):
127
f"_parallel_size ({pcp_size}). "
128
)
129
125
- set_multiprocessing_worker_envs()
130
+ set_multiprocessing_worker_envs(self.local_world_size)
131
132
# use the loopback address get_loopback_ip() for communication.
133
distributed_init_method = get_distributed_init_method(
@@ -200,6 +205,12 @@ class MultiprocExecutor(Executor):
205
# Wait for all local workers to be ready.
206
self.workers = WorkerProc.wait_for_ready(unready_workers)
207
208
+ # The workers have inherited their thread count (see
209
+ # set_multiprocessing_worker_envs); this process only schedules, so
210
+ # it gets no benefit from torch intra-op parallelism, just CPU
211
+ # contention with them.
212
+ set_torch_threads_for_runtime()
213
+
214
# Start background thread to monitor worker health if not in headless mode.
215
if self.monitor_workers:
216
self.start_worker_monitor()
@@ -1058,32 +1069,33 @@ class WorkerProc:
1069
decorate_logs(process_name)
1070
1071
1061
-def set_multiprocessing_worker_envs():
1072
+def set_multiprocessing_worker_envs(local_world_size: int = 1):
1073
"""Set up environment variables that should be used when there are workers
1074
in a multiprocessing environment. This should be called by the parent
1075
process before worker processes are created"""
1076
1077
_maybe_force_spawn()
1078
1068
- if not current_platform.is_cpu():
1069
- # Configure thread parallelism if OMP_NUM_THREADS isn't set
1070
- #
1071
- # Helps to avoid CPU contention. The default of spawning a thread per
1072
- # core combined with multiprocessing for each GPU can have a negative
1073
- # impact on performance. The contention is amplified when running in a
1074
- # container where CPU limits can cause throttling.
1075
- default_omp_num_threads = 1
1076
- if (
1077
- "OMP_NUM_THREADS" not in os.environ
1078
- and (current_parallelism := torch.get_num_threads())
1079
- > default_omp_num_threads
1080
- ):
1081
- logger.warning_once(
1082
- "Reducing Torch parallelism from %d threads to %d to avoid "
1083
- "unnecessary CPU contention. Set OMP_NUM_THREADS in the "
1084
- "external environment to tune this value as needed.",
1085
- current_parallelism,
1086
- default_omp_num_threads,
1087
- )
1088
- os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads)
1089
- torch.set_num_threads(default_omp_num_threads)
1079
+ if current_platform.is_cpu() or "OMP_NUM_THREADS" in os.environ:
1080
+ return
1081
+
1082
+ # Choose the workers' thread count here, before they start, since a worker
1083
+ # must not set its own: `torch.set_num_threads()` spawns the thread pool
1084
+ # eagerly, and doing that part way through a worker's startup either races
1085
+ # the dlopen of shared objects or, in a forked worker, deadlocks (libgomp
1086
+ # is not fork-safe).
1087
+ num_threads = startup_omp_num_threads(local_world_size)
1088
+ os.environ["OMP_NUM_THREADS"] = str(num_threads)
1089
+ os.environ[OMP_NUM_THREADS_SET_BY_VLLM] = "1"
1090
+
1091
+ # A spawned worker picks the count up from the environment when it imports
1092
+ # torch. A forked worker instead inherits it from this process, so set it
1093
+ # here too. This is safe as long as we don't *use* the pool before forking:
1094
+ # a forked child whose parent had run a parallel region deadlocks, whereas
1095
+ # one whose parent merely sized the pool does not.
1096
+ torch.set_num_threads(num_threads)
1097
+ logger.debug(
1098
+ "Set OMP_NUM_THREADS=%d for %d worker process(es).",
1099
+ num_threads,
1100
+ local_world_size,
1101
+ )
vllm/v1/worker/gpu_worker.py
+5
-1
@@ -63,7 +63,7 @@ from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
63
from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
64
from vllm.utils.mem_constants import GiB_bytes
65
from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling
66
-from vllm.utils.torch_utils import set_random_seed
66
+from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime
67
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
68
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
69
from vllm.v1.outputs import (
@@ -847,6 +847,10 @@ class Worker(WorkerBase):
847
# gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
848
enable_gpu_sync_check()
849
850
+ # Startup is done; steady-state serving gets no benefit from torch
851
+ # intra-op parallelism.
852
+ set_torch_threads_for_runtime()
853
+
854
return CompilationTimes(
855
language_model=self.compilation_config.compilation_time,
856
encoder=self.compilation_config.encoder_compilation_time,