[Core] Explicitly manage torch CPU threads in workers (#49919)

Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

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

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):
     """

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,
+    )

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,