From 7635a9002baecca64909dbcf8b1d461d26fb879d Mon Sep 17 00:00:00 2001 From: Nick Hill Date: Tue, 4 Aug 2026 10:06:38 -0700 Subject: [PATCH] [Core] Explicitly manage torch CPU threads in workers (#49919) Signed-off-by: Nick Hill Co-authored-by: Claude --- tests/utils_/test_torch_utils.py | 50 +++++++++++ vllm/utils/torch_utils.py | 116 +++++++++++++++++++++++++ vllm/v1/executor/multiproc_executor.py | 60 ++++++++----- vllm/v1/worker/gpu_worker.py | 6 +- 4 files changed, 207 insertions(+), 25 deletions(-) diff --git a/tests/utils_/test_torch_utils.py b/tests/utils_/test_torch_utils.py index f6a9486a12..ddaadfd123 100644 --- a/tests/utils_/test_torch_utils.py +++ b/tests/utils_/test_torch_utils.py @@ -4,9 +4,13 @@ import pytest import torch from vllm.utils.torch_utils import ( + OMP_NUM_THREADS_SET_BY_VLLM, + available_cpu_count, common_broadcastable_dtype, current_stream, is_lossless_cast, + set_torch_threads_for_runtime, + startup_omp_num_threads, ) @@ -114,3 +118,49 @@ def test_current_stream_multithread(): ) _test_stream_thread(main_dedicated_stream) + + +@pytest.fixture +def restore_torch_threads(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("OMP_NUM_THREADS", raising=False) + original = torch.get_num_threads() + yield + torch.set_num_threads(original) + + +def test_startup_omp_num_threads_divides_between_local_workers(): + """Workers share the node's usable CPUs rather than each taking them all.""" + available = available_cpu_count() + if available < 4: + pytest.skip("needs at least 4 usable CPUs") + assert startup_omp_num_threads(1) == available + assert startup_omp_num_threads(2) == available // 2 + # Never zero, however many workers share the node. + assert startup_omp_num_threads(available * 4) == 1 + + +def test_set_torch_threads_for_runtime(restore_torch_threads): + torch.set_num_threads(max(2, available_cpu_count())) + set_torch_threads_for_runtime() + assert torch.get_num_threads() == 1 + + +def test_runtime_threads_respect_user_omp_num_threads( + restore_torch_threads, monkeypatch: pytest.MonkeyPatch +): + """An externally-set OMP_NUM_THREADS is the user's choice; leave it alone.""" + monkeypatch.setenv("OMP_NUM_THREADS", "3") + torch.set_num_threads(3) + set_torch_threads_for_runtime() + assert torch.get_num_threads() == 3 + + +def test_runtime_threads_override_vllm_set_omp_num_threads( + restore_torch_threads, monkeypatch: pytest.MonkeyPatch +): + """The value vLLM picked for worker startup is dropped once serving starts.""" + monkeypatch.setenv("OMP_NUM_THREADS", "3") + monkeypatch.setenv(OMP_NUM_THREADS_SET_BY_VLLM, "1") + torch.set_num_threads(3) + set_torch_threads_for_runtime() + assert torch.get_num_threads() == 1 diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 1472024c2b..16c092c949 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -4,6 +4,7 @@ import contextlib import importlib.metadata import os import random +import sys import threading from collections.abc import Callable, Collection from typing import TYPE_CHECKING, Any, TypeVar @@ -150,6 +151,121 @@ def set_default_torch_dtype(dtype: torch.dtype): torch.set_default_dtype(old_dtype) +def _cgroup_cpu_limit() -> float | None: + """Effective CPU quota of this process's cgroup, None if unlimited. + + Resolves the process's own cgroup from /proc/self/cgroup and takes the + tightest cpu.max (v2) or cfs quota (v1) along the hierarchy. + """ + limit: float | None = None + try: + with open("/proc/self/cgroup") as f: + entries = [line.strip().split(":", 2) for line in f] + + def visit(base: str, rel_path: str, read_quota) -> None: + nonlocal limit + path = rel_path + while path: + quota = read_quota(os.path.join(base, path.lstrip("/"))) + if quota is not None: + limit = quota if limit is None else min(limit, quota) + path = path.rsplit("/", 1)[0] + + def read_v2(cg_dir: str) -> float | None: + try: + with open(os.path.join(cg_dir, "cpu.max")) as f: + quota, period = f.read().split() + return None if quota == "max" else float(quota) / float(period) + except (OSError, ValueError): + return None + + def read_v1(cg_dir: str) -> float | None: + try: + with open(os.path.join(cg_dir, "cpu.cfs_quota_us")) as f: + quota = int(f.read()) + if quota <= 0: + return None + with open(os.path.join(cg_dir, "cpu.cfs_period_us")) as f: + return quota / int(f.read()) + except (OSError, ValueError): + return None + + for entry in entries: + if len(entry) != 3: + continue + _, controllers, rel_path = entry + if controllers == "": # cgroup v2 + visit("/sys/fs/cgroup", rel_path, read_v2) + elif "cpu" in controllers.split(","): # cgroup v1 + visit("/sys/fs/cgroup/cpu", rel_path, read_v1) + except OSError: + pass + return limit + + +def available_cpu_count() -> int: + """CPUs actually usable by this process: scheduling affinity capped by + the cgroup CPU quota (unlike `os.cpu_count()`, which is quota-blind).""" + if sys.platform != "linux": + return os.cpu_count() or 1 + count = len(os.sched_getaffinity(0)) + limit = _cgroup_cpu_limit() + if limit is not None: + count = min(count, int(limit)) + return max(1, count) + + +# Marks OMP_NUM_THREADS as chosen by vLLM for its worker processes rather than +# set by the user, so a worker knows it may drop the value once startup is done. +OMP_NUM_THREADS_SET_BY_VLLM = "VLLM_OMP_NUM_THREADS_SET_BY_VLLM" + + +def startup_omp_num_threads(num_local_procs: int) -> int: + """Thread count for a worker process's startup work (weight loading). + + Weight loading does CPU-parallel work, so workers benefit from more than + one thread, but only a bounded share of the CPUs this node's workers may + actually use: torch's default is the host core count, which ignores both + scheduling affinity and any cgroup CPU quota, and doesn't account for the + other workers sharing the node. + """ + return max(1, available_cpu_count() // max(1, num_local_procs)) + + +def set_torch_threads_for_runtime() -> None: + """Set torch intra-op threads to 1 for steady-state serving. + + Any multi-threaded torch CPU op in the engine hot loop leaves the OMP + workers spin-waiting after each parallel region, stealing cycles from the + step's serial code (and burning cgroup CPU quota in containers). No + steady-state CPU op benefits from intra-op parallelism. + Respects an externally-set OMP_NUM_THREADS. + """ + if ( + omp_num_threads := os.environ.get("OMP_NUM_THREADS") + ) is not None and os.environ.get(OMP_NUM_THREADS_SET_BY_VLLM) != "1": + try: + if int(omp_num_threads) > 1: + logger.warning_once( + "OMP_NUM_THREADS=%s is set; leaving Torch threads at %d " + "for serving. Multi-threaded torch CPU ops during serving " + "can degrade performance through spin-wait contention and " + "cgroup CPU-quota throttling.", + omp_num_threads, + torch.get_num_threads(), + ) + except ValueError: + pass + return + if torch.get_num_threads() != 1: + logger.info_once( + "Reducing Torch threads from %d to 1 for serving. Set " + "OMP_NUM_THREADS in the external environment to override.", + torch.get_num_threads(), + ) + torch.set_num_threads(1) + + @contextlib.contextmanager def set_default_torch_num_threads(num_threads: int | None = None): """ diff --git a/vllm/v1/executor/multiproc_executor.py b/vllm/v1/executor/multiproc_executor.py index b464f2dcf6..276effd282 100644 --- a/vllm/v1/executor/multiproc_executor.py +++ b/vllm/v1/executor/multiproc_executor.py @@ -58,6 +58,11 @@ from vllm.utils.system_utils import ( get_mp_context, set_process_title, ) +from vllm.utils.torch_utils import ( + OMP_NUM_THREADS_SET_BY_VLLM, + set_torch_threads_for_runtime, + startup_omp_num_threads, +) from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.executor.abstract import Executor, FailureCallback from vllm.v1.executor.vllm_net_devices import set_worker_net_device @@ -122,7 +127,7 @@ class MultiprocExecutor(Executor): f"_parallel_size ({pcp_size}). " ) - set_multiprocessing_worker_envs() + set_multiprocessing_worker_envs(self.local_world_size) # use the loopback address get_loopback_ip() for communication. distributed_init_method = get_distributed_init_method( @@ -200,6 +205,12 @@ class MultiprocExecutor(Executor): # Wait for all local workers to be ready. self.workers = WorkerProc.wait_for_ready(unready_workers) + # The workers have inherited their thread count (see + # set_multiprocessing_worker_envs); this process only schedules, so + # it gets no benefit from torch intra-op parallelism, just CPU + # contention with them. + set_torch_threads_for_runtime() + # Start background thread to monitor worker health if not in headless mode. if self.monitor_workers: self.start_worker_monitor() @@ -1058,32 +1069,33 @@ class WorkerProc: decorate_logs(process_name) -def set_multiprocessing_worker_envs(): +def set_multiprocessing_worker_envs(local_world_size: int = 1): """Set up environment variables that should be used when there are workers in a multiprocessing environment. This should be called by the parent process before worker processes are created""" _maybe_force_spawn() - if not current_platform.is_cpu(): - # Configure thread parallelism if OMP_NUM_THREADS isn't set - # - # Helps to avoid CPU contention. The default of spawning a thread per - # core combined with multiprocessing for each GPU can have a negative - # impact on performance. The contention is amplified when running in a - # container where CPU limits can cause throttling. - default_omp_num_threads = 1 - if ( - "OMP_NUM_THREADS" not in os.environ - and (current_parallelism := torch.get_num_threads()) - > default_omp_num_threads - ): - logger.warning_once( - "Reducing Torch parallelism from %d threads to %d to avoid " - "unnecessary CPU contention. Set OMP_NUM_THREADS in the " - "external environment to tune this value as needed.", - current_parallelism, - default_omp_num_threads, - ) - os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads) - torch.set_num_threads(default_omp_num_threads) + if current_platform.is_cpu() or "OMP_NUM_THREADS" in os.environ: + return + + # Choose the workers' thread count here, before they start, since a worker + # must not set its own: `torch.set_num_threads()` spawns the thread pool + # eagerly, and doing that part way through a worker's startup either races + # the dlopen of shared objects or, in a forked worker, deadlocks (libgomp + # is not fork-safe). + num_threads = startup_omp_num_threads(local_world_size) + os.environ["OMP_NUM_THREADS"] = str(num_threads) + os.environ[OMP_NUM_THREADS_SET_BY_VLLM] = "1" + + # A spawned worker picks the count up from the environment when it imports + # torch. A forked worker instead inherits it from this process, so set it + # here too. This is safe as long as we don't *use* the pool before forking: + # a forked child whose parent had run a parallel region deadlocks, whereas + # one whose parent merely sized the pool does not. + torch.set_num_threads(num_threads) + logger.debug( + "Set OMP_NUM_THREADS=%d for %d worker process(es).", + num_threads, + local_world_size, + ) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index bbab98dbd7..ed11c0def9 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -63,7 +63,7 @@ from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check from vllm.utils.mem_constants import GiB_bytes from vllm.utils.mem_utils import MemorySnapshot, format_gib, memory_profiling -from vllm.utils.torch_utils import set_random_seed +from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec from vllm.v1.outputs import ( @@ -847,6 +847,10 @@ class Worker(WorkerBase): # gate so subsequent `execute_model` / `sample_tokens` calls enforce it. enable_gpu_sync_check() + # Startup is done; steady-state serving gets no benefit from torch + # intra-op parallelism. + set_torch_threads_for_runtime() + return CompilationTimes( language_model=self.compilation_config.compilation_time, encoder=self.compilation_config.encoder_compilation_time,