@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / d02421a7db
Fork 0

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 5, 01:00 UTC

[CPU] Refactor CPU affinity and memory management (#39781)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed Apr 17, 2026 at 21:01 UTC d02421a7dbd85eb173cb2620da3dbc16d81135f4
14 files changed +723 -425
.buildkite/scripts/hardware_ci/run-cpu-distributed-smoke-test.sh
+19 -19
@@ -23,22 +23,22 @@ if [ "$failed_req" -ne 0 ]; then
23 exit 1
24 fi
25
26 -#echo "--- DP+TP"
27 -#vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
28 -#server_pid=$!
29 -#timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
30 -#vllm bench serve \
31 -# --backend vllm \
32 -# --dataset-name random \
33 -# --model meta-llama/Llama-3.2-3B-Instruct \
34 -# --num-prompts 20 \
35 -# --result-dir ./test_results \
36 -# --result-filename dp_pp.json \
37 -# --save-result \
38 -# --endpoint /v1/completions
39 -#kill -s SIGTERM $server_pid; wait $server_pid || true
40 -#failed_req=$(jq '.failed' ./test_results/dp_pp.json)
41 -#if [ "$failed_req" -ne 0 ]; then
42 -# echo "Some requests were failed!"
43 -# exit 1
44 -#fi
26 +echo "--- DP+TP"
27 +vllm serve meta-llama/Llama-3.2-3B-Instruct -tp=2 -dp=2 --max-model-len=4096 &
28 +server_pid=$!
29 +timeout 600 bash -c "until curl localhost:8000/v1/models > /dev/null 2>&1; do sleep 1; done" || exit 1
30 +vllm bench serve \
31 + --backend vllm \
32 + --dataset-name random \
33 + --model meta-llama/Llama-3.2-3B-Instruct \
34 + --num-prompts 20 \
35 + --result-dir ./test_results \
36 + --result-filename dp_pp.json \
37 + --save-result \
38 + --endpoint /v1/completions
39 +kill -s SIGTERM $server_pid; wait $server_pid || true
40 +failed_req=$(jq '.failed' ./test_results/dp_pp.json)
41 +if [ "$failed_req" -ne 0 ]; then
42 + echo "Some requests were failed!"
43 + exit 1
44 +fi
.github/workflows/macos-smoke-test.yml
+1
@@ -45,6 +45,7 @@ jobs:
45 - name: Smoke test vllm serve
46 run: |
47 # Start server in background
48 + VLLM_CPU_KVCACHE_SPACE=1 \
49 vllm serve Qwen/Qwen3-0.6B \
50 --max-model-len=2K \
51 --load-format=dummy \
cmake/cpu_extension.cmake
+15 -14
@@ -30,6 +30,21 @@ else()
30 list(APPEND CXX_COMPILE_FLAGS
31 "-fopenmp"
32 "-DVLLM_CPU_EXTENSION")
33 +
34 + # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
35 + # and create a local shim dir with it
36 + vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
37 +
38 + find_library(OPEN_MP
39 + NAMES gomp
40 + PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
41 + NO_DEFAULT_PATH
42 + REQUIRED
43 + )
44 + # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
45 + if (OPEN_MP)
46 + set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
47 + endif()
48 endif()
49
50 if (NOT MACOSX_FOUND)
@@ -175,20 +190,6 @@ if (ENABLE_X86_ISA OR (ASIMD_FOUND AND NOT APPLE_SILICON_FOUND) OR POWER9_FOUND
190 if(NOT NPROC)
191 set(NPROC 4)
192 endif()
178 - # locate PyTorch's libgomp (e.g. site-packages/torch.libs/libgomp-947d5fa1.so.1.0.0)
179 - # and create a local shim dir with it
180 - vllm_prepare_torch_gomp_shim(VLLM_TORCH_GOMP_SHIM_DIR)
181 -
182 - find_library(OPEN_MP
183 - NAMES gomp
184 - PATHS ${VLLM_TORCH_GOMP_SHIM_DIR}
185 - NO_DEFAULT_PATH
186 - REQUIRED
187 - )
188 - # Set LD_LIBRARY_PATH to include the shim dir at build time to use the same libgomp as PyTorch
189 - if (OPEN_MP)
190 - set(ENV{LD_LIBRARY_PATH} "${VLLM_TORCH_GOMP_SHIM_DIR}:$ENV{LD_LIBRARY_PATH}")
191 - endif()
193
194 # Fetch and populate ACL
195 if(DEFINED ENV{ACL_ROOT_DIR} AND IS_DIRECTORY "$ENV{ACL_ROOT_DIR}")
csrc/cpu/torch_bindings.cpp
+4
@@ -141,6 +141,8 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc,
141 torch::Tensor slot_mapping,
142 const int64_t block_size);
143
144 +void init_cpu_memory_env(std::vector<int64_t> node_ids);
145 +
146 namespace cpu_utils {
147 void eagle_prepare_inputs_padded_kernel_impl(
148 const torch::Tensor& cu_num_draft_tokens,
@@ -431,6 +433,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
433 "block_size) -> ()",
434 &compute_slot_mapping_kernel_impl);
435
436 + ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env);
437 +
438 // Speculative decoding kernels
439 ops.def(
440 "eagle_prepare_inputs_padded_kernel_impl(Tensor cu_num_draft_tokens, "
csrc/cpu/utils.cpp
+73 -6
@@ -13,13 +13,80 @@
13 #include "cpu/utils.hpp"
14
15 #ifdef VLLM_NUMA_DISABLED
16 -std::string init_cpu_threads_env(const std::string& cpu_ids) {
17 - return std::string(
18 - "Warning: NUMA is not enabled in this build. `init_cpu_threads_env` has "
19 - "no effect to setup thread affinity.");
20 -}
16 +void init_cpu_memory_env(std::vector<int64_t> node_ids) {}
17 +#else
18 +void init_cpu_memory_env(std::vector<int64_t> node_ids) {
19 + // Memory node binding
20 + if (numa_available() != -1) {
21 + // Concatenate all node_ids into a single comma-separated string
22 + if (!node_ids.empty()) {
23 + std::string node_ids_str;
24 + for (const int node_id : node_ids) {
25 + if (!node_ids_str.empty()) {
26 + node_ids_str += ",";
27 + }
28 + node_ids_str += std::to_string(node_id);
29 + }
30
22 -#endif
31 + bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());
32 + bitmask* src_mask = numa_get_mems_allowed();
33 +
34 + int pid = getpid();
35 +
36 + if (mask && src_mask) {
37 + // move all existing pages to the specified numa node.
38 + *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);
39 + int page_num = numa_migrate_pages(pid, src_mask, mask);
40 + if (page_num == -1) {
41 + TORCH_WARN("numa_migrate_pages failed. errno: " +
42 + std::to_string(errno));
43 + }
44 +
45 + // Restrict memory allocation to the selected NUMA node(s).
46 + // Enhances memory locality for the threads bound to those NUMA CPUs.
47 + if (node_ids.size() > 1) {
48 + errno = 0;
49 + numa_set_interleave_mask(mask);
50 + if (errno != 0) {
51 + TORCH_WARN("numa_set_interleave_mask failed. errno: " +
52 + std::to_string(errno));
53 + } else {
54 + TORCH_WARN(
55 + "NUMA binding: Using INTERLEAVE policy for memory "
56 + "allocation across multiple NUMA nodes (nodes: " +
57 + node_ids_str +
58 + "). Memory allocations will be "
59 + "interleaved across the specified NUMA nodes.");
60 + }
61 + } else {
62 + errno = 0;
63 + numa_set_membind(mask);
64 + if (errno != 0) {
65 + TORCH_WARN("numa_set_membind failed. errno: " +
66 + std::to_string(errno));
67 + } else {
68 + TORCH_WARN(
69 + "NUMA binding: Using MEMBIND policy for memory "
70 + "allocation on the NUMA nodes (" +
71 + node_ids_str +
72 + "). Memory allocations will be "
73 + "strictly bound to these NUMA nodes.");
74 + }
75 + }
76 +
77 + numa_set_strict(1);
78 +
79 + numa_free_nodemask(mask);
80 + numa_free_nodemask(src_mask);
81 + } else {
82 + TORCH_WARN(
83 + "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +
84 + std::to_string(errno));
85 + }
86 + }
87 + }
88 +}
89 +#endif // VLLM_NUMA_DISABLED
90
91 namespace cpu_utils {
92 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {
docker/Dockerfile.cpu
+2 -1
@@ -173,7 +173,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
173 COPY --from=vllm-test-deps /vllm-workspace/requirements/test/cpu.txt requirements/test/cpu.txt
174
175 RUN --mount=type=cache,target=/root/.cache/uv \
176 - uv pip install -r requirements/dev.txt && \
176 + uv pip install -r requirements/lint.txt && \
177 + uv pip install -r requirements/test/cpu.txt && \
178 pre-commit install --hook-type pre-commit --hook-type commit-msg
179
180 ENTRYPOINT ["bash"]
tests/models/language/generation/test_common.py
+1 -6
@@ -46,7 +46,7 @@ AITER_MODEL_LIST = [
46 ),
47 pytest.param(
48 "openai-community/gpt2", # gpt2
49 - marks=[pytest.mark.core_model, pytest.mark.cpu_model],
49 + marks=[pytest.mark.core_model],
50 ),
51 pytest.param("Milos/slovak-gpt-j-405M"), # gptj
52 pytest.param("bigcode/tiny_starcoder_py"), # gpt_bigcode
@@ -143,11 +143,6 @@ def test_models(
143 # in parts of the operators
144 pytest.skip(f"Skipping '{model}' model test with AITER kernel.")
145
146 - if current_platform.is_cpu() and model in ("openai-community/gpt2",):
147 - # These models are sensitive to the rounding error
148 - # Fuse ops to reduce rounding
149 - monkeypatch.setenv("VLLM_CPU_CI_ENV", "0")
150 -
146 with hf_runner(model) as hf_model:
147 hf_outputs = hf_model.generate_greedy_logprobs_limit(
148 example_prompts, max_tokens, num_logprobs
vllm/config/cache.py
-3
@@ -101,8 +101,6 @@ class CacheConfig:
101 kv_cache_dtype_skip_layers: list[str] = field(default_factory=list)
102 """Layer patterns to skip KV cache quantization. Accepts layer indices
103 (e.g., '0', '2', '4') or attention type names (e.g., 'sliding_window')."""
104 - cpu_kvcache_space_bytes: int | None = None
105 - """(CPU backend only) CPU key-value cache space."""
104 mamba_page_size_padded: int | None = None
105 """ Optional override for mamba page size; used by hybrid mamba/attention
106 models to ensure exact alignment with attention page size."""
@@ -183,7 +181,6 @@ class CacheConfig:
181 "num_gpu_blocks_override",
182 "enable_prefix_caching",
183 "prefix_caching_hash_algo",
186 - "cpu_kvcache_space_bytes",
184 "mamba_page_size_padded",
185 "user_specified_block_size",
186 "user_specified_mamba_block_size",
vllm/platforms/cpu.py
+51 -157
@@ -6,15 +6,16 @@ import os
6 import platform
7 import subprocess
8 import sys
9 -from dataclasses import dataclass
9 from typing import TYPE_CHECKING
10
12 -import psutil
11 import torch
12
15 -from vllm import envs
13 from vllm.logger import init_logger
17 -from vllm.utils.ompmultiprocessing import OMPProcessManager
14 +from vllm.utils.cpu_resource_utils import (
15 + DEVICE_CONTROL_ENV_VAR,
16 + get_memory_node_info,
17 +)
18 +from vllm.utils.mem_constants import GiB_bytes
19 from vllm.utils.torch_utils import is_quantized_kv_cache
20 from vllm.v1.attention.backends.registry import AttentionBackendEnum
21
@@ -38,49 +39,13 @@ def get_max_threads(pid=0):
39 raise NotImplementedError("Unsupported OS")
40
41
41 -@dataclass
42 -class LogicalCPUInfo:
43 - id: int = -1
44 - physical_core: int = -1
45 - numa_node: int = -1
46 -
47 - @classmethod
48 - def _int(cls, value: str) -> int:
49 - try:
50 - int_value = int(value)
51 - except Exception:
52 - int_value = -1
53 - return int_value
54 -
55 - @staticmethod
56 - def json_decoder(obj_dict: dict):
57 - id = obj_dict.get("cpu")
58 - physical_core = obj_dict.get("core")
59 - numa_node = obj_dict.get("node")
60 -
61 - if not (id is None or physical_core is None or numa_node is None):
62 - return LogicalCPUInfo(
63 - id=LogicalCPUInfo._int(id),
64 - physical_core=LogicalCPUInfo._int(physical_core),
65 - numa_node=LogicalCPUInfo._int(numa_node),
66 - )
67 - else:
68 - return obj_dict
69 -
70 -
42 class CpuPlatform(Platform):
43 _enum = PlatformEnum.CPU
44 device_name: str = "cpu"
45 device_type: str = "cpu"
46 dispatch_key: str = "CPU"
47 dist_backend: str = "gloo"
77 - device_control_env_var = "CPU_VISIBLE_MEMORY_NODES"
78 - omp_process_manager = None
79 - # Simultaneous Multithreading (SMT) level for OpenMP:
80 - # 4 on PowerPC, 1 on non-PowerPC architectures
81 - smt = 1
82 - global_cpu_mask = None
83 - simulate_numa = int(os.environ.get("_SIM_MULTI_NUMA", 0))
48 + device_control_env_var = DEVICE_CONTROL_ENV_VAR
49
50 @property
51 def supported_dtypes(self) -> list[torch.dtype]:
@@ -123,29 +88,9 @@ class CpuPlatform(Platform):
88
89 @classmethod
90 def get_device_total_memory(cls, device_id: int = 0) -> int:
126 - from vllm.utils.mem_constants import GiB_bytes
127 - from vllm.utils.mem_utils import format_gib
128 -
129 - kv_cache_space = envs.VLLM_CPU_KVCACHE_SPACE
130 - node_dir = "/sys/devices/system/node"
131 - if kv_cache_space is None:
132 - nodes = (
133 - [d for d in os.listdir(node_dir) if d.startswith("node")]
134 - if os.path.exists(node_dir)
135 - else []
136 - )
137 - num_numa_nodes = len(nodes) or 1
138 - free_cpu_memory = psutil.virtual_memory().total // num_numa_nodes
139 - DEFAULT_CPU_MEM_UTILIZATION = 0.5
140 - kv_cache_space = int(free_cpu_memory * DEFAULT_CPU_MEM_UTILIZATION)
141 - logger.warning_once(
142 - "VLLM_CPU_KVCACHE_SPACE not set. Using %s GiB for KV cache.",
143 - format_gib(kv_cache_space),
144 - )
145 - else:
146 - kv_cache_space *= GiB_bytes
91 + meminfo = get_memory_node_info(device_id)
92
148 - return kv_cache_space
93 + return meminfo.total_memory
94
95 @classmethod
96 def set_device(cls, device: torch.device) -> None:
@@ -180,6 +125,12 @@ class CpuPlatform(Platform):
125 "otherwise the performance is not optimized."
126 )
127
128 + # Lagecy setting
129 + env_key = "VLLM_CPU_KVCACHE_SPACE"
130 + if env_key in os.environ and os.environ[env_key] != "":
131 + kv_cache_space = int(os.environ[env_key])
132 + cache_config.kv_cache_memory_bytes = kv_cache_space * GiB_bytes
133 +
134 scheduler_config = vllm_config.scheduler_config
135 # async scheduling is not required on CPU
136 scheduler_config.async_scheduling = False
@@ -198,8 +149,6 @@ class CpuPlatform(Platform):
149 )
150 cache_config.cache_dtype = "auto"
151
201 - cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory()
202 -
152 parallel_config = vllm_config.parallel_config
153 # OMP requires the MP executor to function correctly, UniProc is not
154 # supported as it is not possible to set the OMP environment correctly
@@ -278,21 +227,45 @@ class CpuPlatform(Platform):
227 os.environ["TORCHINDUCTOR_CPP_DYNAMIC_THREADS"] = "1"
228
229 ld_preload_str = os.getenv("LD_PRELOAD", "")
281 -
282 - # Intel and CLANG OpenMP setting
283 - if "libiomp5.so" in ld_preload_str or "libomp5" in ld_preload_str:
284 - # The time(milliseconds) that a thread should wait after
285 - # completing the execution of a parallel region, before sleeping.
286 - os.environ["KMP_BLOCKTIME"] = "1"
287 - # Prevents the CPU to run into low performance state
288 - os.environ["KMP_TPAUSE"] = "0"
289 - # Provides fine granularity parallelism
290 - os.environ["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist"
291 - os.environ["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist"
292 - os.environ["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist"
293 -
230 cpu_architecture = Platform.get_cpu_architecture()
231
232 + if (
233 + platform.system() == "Linux"
234 + and cpu_architecture
235 + in (CpuArchEnum.ARM, CpuArchEnum.POWERPC, CpuArchEnum.X86)
236 + and not (
237 + "libomp" in ld_preload_str
238 + or "libgomp" in ld_preload_str
239 + or "libiomp" in ld_preload_str
240 + )
241 + ):
242 + # We need to LD_PRELOAD PyTorch's libgomp, otherwise only
243 + # one core will be properly utilized when we thread-bind
244 + # See: https://github.com/vllm-project/vllm/issues/27369
245 + # TODO: Remove once:
246 + # https://github.com/pytorch/pytorch/issues/166087 is fixed
247 +
248 + # We need to find the location of PyTorch's libgomp
249 + torch_pkg = os.path.dirname(torch.__file__)
250 + site_root = os.path.dirname(torch_pkg)
251 + # Search both torch.libs and torch/lib - See:
252 + # https://github.com/vllm-project/vllm/issues/30470
253 + torch_libs_paths = [
254 + os.path.join(site_root, "torch.libs"),
255 + os.path.join(torch_pkg, "lib"),
256 + ]
257 + pytorch_libgomp_so_candidates = []
258 + for torch_libs in torch_libs_paths:
259 + pytorch_libgomp_so_candidates.extend(
260 + glob.glob(os.path.join(torch_libs, "libgomp*.so*"))
261 + )
262 + if pytorch_libgomp_so_candidates:
263 + pytorch_libgomp_so = pytorch_libgomp_so_candidates[0]
264 + if ld_preload_str:
265 + ld_preload_str += ":"
266 + ld_preload_str += pytorch_libgomp_so
267 + os.environ["LD_PRELOAD"] = ld_preload_str
268 +
269 # LD_PRELOAD libtcmalloc, bundled under vllm/libs to reduce
270 # memory allocation overhead
271 if (
@@ -331,13 +304,6 @@ class CpuPlatform(Platform):
304 vllm_config.model_config.max_model_len,
305 vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,
306 )
334 - # CI specific "quick" NUMA simulation - split all available CPUs
335 - # into a fake NUMA topology
336 - if os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", None) is not None:
337 - os.environ["_SIM_MULTI_NUMA"] = str(
338 - vllm_config.parallel_config.world_size
339 - * vllm_config.parallel_config._api_process_count
340 - )
307
308 @classmethod
309 def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:
@@ -345,78 +311,6 @@ class CpuPlatform(Platform):
311 # Move that logic here so block_size is chosen by the backend.
312 pass
313
348 - @classmethod
349 - def get_omp_manager(cls) -> OMPProcessManager:
350 - # initialise the OMP resource management if need be and return the manager
351 - if cls.omp_process_manager is None:
352 - if cls.get_cpu_architecture() == CpuArchEnum.POWERPC:
353 - cls.smt = 4
354 - cls.omp_process_manager = OMPProcessManager(
355 - affinity=cls.get_global_cpu_mask(), smt=cls.smt
356 - )
357 - # we need to fix up the topology returned by the OMP Manager for
358 - # simulated NUMA environments in CI
359 - if cls.simulate_numa > 0:
360 - logger.info(
361 - "Adjusting numa topology to resemble at least %d nodes",
362 - int(cls.simulate_numa),
363 - )
364 - om = cls.omp_process_manager
365 - while len(om.omp_places) < cls.simulate_numa:
366 - new_omp_places = []
367 - touched = False
368 - for omp_place in om.omp_places:
369 - if len(omp_place["mask"]) > 1:
370 - touched = True
371 - cpu_list = sorted(list(omp_place["mask"]))
372 - new_omp_places.append(
373 - {
374 - "mask": set(cpu_list[0 : int(len(cpu_list) / 2)]),
375 - "available": True,
376 - }
377 - )
378 - new_omp_places.append(
379 - {
380 - "mask": set(cpu_list[int(len(cpu_list) / 2) :]),
381 - "available": True,
382 - }
383 - )
384 - if touched:
385 - om.omp_places = new_omp_places
386 - else:
387 - raise ValueError(
388 - "Cannot split the existing NUMA topology to match "
389 - "simulation requirements"
390 - )
391 -
392 - return cls.omp_process_manager
393 -
394 - @classmethod
395 - def get_global_cpu_mask(cls) -> set[int]:
396 - # get global cpu mask
397 - if cls.global_cpu_mask is None:
398 - if hasattr(os, "sched_getaffinity"):
399 - cls.global_cpu_mask = os.sched_getaffinity(0)
400 - else:
401 - # macOS does not support sched_getaffinity
402 - cpu_count = os.cpu_count() or 1
403 - cls.global_cpu_mask = set(range(cpu_count))
404 - return cls.global_cpu_mask
405 -
406 - @classmethod
407 - def reserve_cpus(cls, reserve: set[int]) -> bool:
408 - # remove CPUs from global mask, for now there is no "release" mechanism
409 - if cls.omp_process_manager is not None:
410 - for place in cls.omp_process_manager.omp_places:
411 - if not place["available"]:
412 - return False
413 - cls.global_cpu_mask = cls.get_global_cpu_mask() - reserve
414 - # reinitialize OMP resource management
415 - cls.omp_process_manager = OMPProcessManager(
416 - affinity=cls.global_cpu_mask, smt=cls.smt
417 - )
418 - return True
419 -
314 @classmethod
315 def discover_numa_topology(cls) -> list[list[int]]:
316 """
vllm/utils/cpu_resource_utils.py new
+173
@@ -0,0 +1,173 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +import os
6 +import platform
7 +import subprocess
8 +from dataclasses import dataclass
9 +from functools import cache
10 +
11 +import psutil
12 +import regex as re
13 +
14 +DEVICE_CONTROL_ENV_VAR = "CPU_VISIBLE_MEMORY_NODES"
15 +
16 +
17 +@dataclass
18 +class LogicalCPUInfo:
19 + id: int = -1
20 + physical_core: int = -1
21 + numa_node: int = -1
22 +
23 + @classmethod
24 + def _int(cls, value: str) -> int:
25 + try:
26 + int_value = int(value)
27 + except Exception:
28 + int_value = -1
29 + return int_value
30 +
31 + @staticmethod
32 + def json_decoder(obj_dict: dict):
33 + id = obj_dict.get("cpu")
34 + physical_core = obj_dict.get("core")
35 + numa_node = obj_dict.get("node")
36 +
37 + if not (id is None or physical_core is None or numa_node is None):
38 + return LogicalCPUInfo(
39 + id=LogicalCPUInfo._int(id),
40 + physical_core=LogicalCPUInfo._int(physical_core),
41 + numa_node=LogicalCPUInfo._int(numa_node),
42 + )
43 + else:
44 + return obj_dict
45 +
46 +
47 +@dataclass
48 +class MemoryNodeInfo:
49 + total_memory: int = -1
50 + available_memory: int = -1
51 +
52 +
53 +def get_memory_affinity(pid: int = 0) -> list[int]:
54 + pid = os.getpid() if pid == 0 else pid
55 + path = f"/proc/{pid}/status"
56 + with open(path) as f:
57 + for line in f:
58 + if line.startswith("Mems_allowed_list:"):
59 + # Extract the string part (e.g., "0-1,3")
60 + raw_list = line.split(":")[1].strip()
61 + return parse_id_list(raw_list)
62 + return []
63 +
64 +
65 +def parse_id_list(raw_str: str) -> list[int]:
66 + """Parses strings like '0-2,4,7-8' into [0, 1, 2, 4, 7, 8]"""
67 + result: list[int] = []
68 + if not raw_str:
69 + return result
70 +
71 + for part in raw_str.split(","):
72 + if "-" in part:
73 + start, end = map(int, part.split("-"))
74 + result.extend(range(start, end + 1))
75 + else:
76 + result.append(int(part))
77 + return sorted(list(set(result)))
78 +
79 +
80 +def get_memory_node_info(node_id: int = 0) -> MemoryNodeInfo:
81 + if platform.system() == "Darwin":
82 + # MacOS has no memory node
83 + return MemoryNodeInfo(
84 + total_memory=psutil.virtual_memory().total,
85 + available_memory=psutil.virtual_memory().available,
86 + )
87 +
88 + meminfo_path = f"/sys/devices/system/node/node{node_id}/meminfo"
89 + if not os.path.exists(meminfo_path):
90 + raise RuntimeError(f"{meminfo_path} doesn't exit.")
91 +
92 + meminfo = {}
93 + with open(meminfo_path) as f:
94 + for line in f:
95 + # Each line looks like: "Node 0 MemTotal: 97421888 kB"
96 + parts = line.split()
97 + key = parts[2].rstrip(":")
98 + # convert to Bytes
99 + value = int(parts[3]) * 1024
100 + meminfo[key] = value
101 +
102 + total_memory = meminfo["MemTotal"]
103 + free_memory = meminfo["MemFree"]
104 + active_file_memory = meminfo["Active(file)"]
105 + inactive_file_memory = meminfo["Inactive(file)"]
106 + reclaimable_memory = meminfo["SReclaimable"]
107 + available_memory = (
108 + free_memory + active_file_memory + inactive_file_memory + reclaimable_memory
109 + )
110 +
111 + return MemoryNodeInfo(
112 + total_memory=total_memory,
113 + available_memory=available_memory,
114 + )
115 +
116 +
117 +def get_allowed_cpu_list() -> list[LogicalCPUInfo]:
118 + cpu_list = _get_cpu_list()
119 + if platform.system() == "Darwin":
120 + return cpu_list
121 +
122 + global_allowed_cpu_id_list = os.sched_getaffinity(0)
123 + logical_cpu_list = [x for x in cpu_list if x.id in global_allowed_cpu_id_list]
124 +
125 + return logical_cpu_list
126 +
127 +
128 +def get_visible_memory_node() -> list[int]:
129 + if platform.system() == "Darwin":
130 + return [0]
131 +
132 + allowed_memory_node_list = get_memory_affinity()
133 +
134 + env_key = DEVICE_CONTROL_ENV_VAR
135 + if (
136 + ("VLLM_CPU_SIM_MULTI_NUMA" not in os.environ)
137 + and env_key in os.environ
138 + and os.environ[env_key] != ""
139 + ):
140 + visible_nodes = [int(s) for s in os.environ[env_key].split(",")]
141 + visible_nodes = [
142 + node for node in visible_nodes if node in allowed_memory_node_list
143 + ]
144 + return visible_nodes
145 +
146 + return allowed_memory_node_list
147 +
148 +
149 +@cache
150 +def _get_cpu_list() -> list[LogicalCPUInfo]:
151 + if platform.system() == "Darwin":
152 + # For MacOS, no user-level CPU affinity and SMT, return all CPUs
153 + cpu_count = os.cpu_count()
154 + assert cpu_count
155 + return [LogicalCPUInfo(i, i, 0) for i in range(cpu_count)]
156 +
157 + lscpu_output = subprocess.check_output(
158 + "lscpu -J -e=CPU,CORE,NODE", shell=True, text=True
159 + )
160 +
161 + # For platform without NUMA, replace '-' to '0'
162 + lscpu_output = re.sub(r'"node":\s*-\s*(,|\n)', r'"node": 0\1', lscpu_output)
163 +
164 + logical_cpu_list: list[LogicalCPUInfo] = json.loads(
165 + lscpu_output, object_hook=LogicalCPUInfo.json_decoder
166 + )["cpus"]
167 +
168 + # Filter CPUs with invalid attributes
169 + logical_cpu_list = [
170 + x for x in logical_cpu_list if -1 not in (x.id, x.physical_core, x.numa_node)
171 + ]
172 +
173 + return logical_cpu_list
vllm/utils/ompmultiprocessing.py
+271 -187
@@ -5,196 +5,280 @@ Copyright (c) 2026 Red Hat Inc
5 Copyright (c) 2026 Cambridge Greys Ltd
6 """
7
8 -import json
8 import os
10 -import platform
11 -import subprocess
12 -
13 -
14 -def _int(arg):
15 - """Relaxed parsing of ints which handles a - instead of a number.
16 - The lscpu json may contain that for nodes in some cases. If that
17 - is the case we parse it to zero
18 - """
19 - try:
20 - if int(arg) >= 0:
21 - return int(arg)
22 - except ValueError:
23 - pass
24 - return 0
25 -
26 -
27 -def parse_mask(mask):
28 - """Expand a X-Y,Z list"""
29 - result = []
30 - for token in mask.split(","):
31 - try:
32 - start, finish = token.split("-")
33 - if int(start) > int(finish):
34 - raise IndexError("Invalid Indexes for cpu ranges")
35 - for cpu in range(int(start), int(finish) + 1):
36 - result.append(cpu)
37 - except ValueError:
38 - result.append(int(token))
39 - return set(result)
40 -
41 -
42 -def _get_default_affinity() -> set[int]:
43 - """Get the set of CPUs the process is allowed to run on."""
44 - if hasattr(os, "sched_getaffinity"):
45 - return os.sched_getaffinity(0)
46 - # macOS does not support sched_getaffinity; fall back to cpu_count
47 - cpu_count = os.cpu_count() or 1
48 - return set(range(cpu_count))
49 -
50 -
51 -def _get_cpu_topology_json() -> bytes:
52 - """Get CPU topology as JSON.
53 -
54 - On Linux this uses ``lscpu -Je``. On other platforms (e.g. macOS) we
55 - synthesize a simple topology where every logical CPU is its own core
56 - on NUMA node 0, which is sufficient for the OMP place-list builder.
57 - """
58 - if platform.system() == "Linux":
59 - return subprocess.run(["lscpu", "-Je"], check=True, capture_output=True).stdout
60 -
61 - # Fallback for non-Linux (macOS, etc.)
62 - cpu_count = os.cpu_count() or 1
63 - cpus = []
64 - for i in range(cpu_count):
65 - cpus.append({"cpu": str(i), "core": str(i), "node": "0"})
66 - return json.dumps({"cpus": cpus}).encode()
67 -
68 -
69 -def enumerate_resources(resource_map, mask=None, allowed=None):
70 - """Enumerate system resources"""
71 - if allowed is None:
72 - allowed = _get_default_affinity()
73 - if mask is not None:
74 - allowed = allowed & mask
75 -
76 - try:
77 - allowed_nodes = parse_mask(os.environ["CPU_VISIBLE_MEMORY_NODES"])
78 - except KeyError:
79 - allowed_nodes = None
80 -
81 - lscpu: dict[str, dict] = {"cpus": {}, "cores": {}, "nodes": {}}
82 - for cpu in resource_map["cpus"]:
83 - cpunum = int(cpu["cpu"])
84 - if (
85 - cpunum in allowed
86 - and cpunum >= 0
87 - and (allowed_nodes is None or _int(cpu["node"]) in allowed_nodes)
88 - ):
89 - lscpu["cpus"][cpunum] = [cpu]
90 - core = _int(cpu["core"])
91 - if lscpu["cores"].get(core, None) is None:
92 - lscpu["cores"][core] = [cpu]
93 - else:
94 - lscpu["cores"][core].append(cpu)
95 - node = _int(cpu["node"])
96 - if lscpu["nodes"].get(node, None) is None:
97 - lscpu["nodes"][node] = [cpu]
98 - else:
99 - lscpu["nodes"][node].append(cpu)
100 - return lscpu
101 -
102 -
103 -def produce_cpu_list(cpus, smt=1):
104 - """Produce a CPU list with/without SMT pairs - main cpu list case"""
105 - mask: list[int] = []
106 - for key, value in cpus.items():
107 - exists = 0
108 - for cpu in mask:
109 - if cpu == value[0]["core"]:
110 - exists += 1
111 - break
112 - if exists < smt:
113 - mask.append(int(key))
114 - return {"mask": set(mask), "available": True}
115 -
116 -
117 -def produce_cpu_sublist(scpus, smt=1):
118 - """Produce a CPU list with/without SMT pairs - resource leaf case"""
119 - cpu_list: list[dict] = []
120 - for value in scpus:
121 - exists = 0
122 - for cpu in cpu_list:
123 - if int(cpu["core"]) == int(value["core"]):
124 - exists += 1
125 - break
126 - if exists < smt:
127 - cpu_list.append(value)
128 - mask = []
129 - for cpu in cpu_list:
130 - mask.append(int(cpu["cpu"]))
131 -
132 - return {"mask": set(mask), "available": True}
133 -
134 -
135 -def create_omp_places(resources, strategy, smt=True):
136 - """Parse CPU topology and generate possible CPU masks"""
137 - omp_places = []
138 - if strategy == "all":
139 - omp_places.append(produce_cpu_list(resources["cpus"], smt))
140 - elif strategy == "cores":
141 - for value in resources["cores"].values():
142 - omp_places.append(produce_cpu_sublist(value, smt))
143 - elif strategy == "nodes":
144 - for value in resources["nodes"].values():
145 - omp_places.append(produce_cpu_sublist(value, smt))
146 - else:
147 - raise NotImplementedError("Unknown strategy")
148 -
149 - return omp_places
150 -
151 -
152 -# pylint: disable=too-few-public-methods
9 +from collections.abc import Callable
10 +from contextlib import contextmanager
11 +from typing import TYPE_CHECKING
12 +
13 +import vllm.utils.cpu_resource_utils as cr_utils
14 +from vllm import envs
15 +from vllm.logger import init_logger
16 +from vllm.platforms import CpuArchEnum, current_platform
17 +from vllm.utils.cpu_resource_utils import LogicalCPUInfo
18 +
19 +if TYPE_CHECKING:
20 + from vllm.config import VllmConfig
21 +
22 +logger = init_logger(__name__)
23 +
24 +
25 class OMPProcessManager:
154 - """OMP aware wrapper to run mp Process()"""
155 -
156 - def __init__(self, strategy="nodes", smt=1, mock=None, affinity=None):
157 - self.strategy = strategy
158 - self.smt = smt
159 - self.omp_places = []
160 - vllm_mask = os.environ.get("VLLM_CPU_OMP_THREADS_BIND", None)
161 - self.setup_omp = vllm_mask != "nobind"
162 - if self.setup_omp:
163 - omp_places = []
164 - if vllm_mask is not None:
165 - masks = []
166 - for spec in vllm_mask.split("|"):
167 - masks.append(parse_mask(spec))
26 + def __init__(self, config: "VllmConfig"):
27 + if not current_platform.is_cpu():
28 + return
29 +
30 + self.local_world_size = config.parallel_config.local_world_size
31 + self.local_dp_rank = config.parallel_config.data_parallel_rank_local
32 + # This is a bit tricky because the internal DP size
33 + # is always 1 for non-MoE models
34 + self.internal_dp_size = config.parallel_config._api_process_count
35 +
36 + self.simulate_multi_node = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0"
37 + ld_preload_str = os.getenv("LD_PRELOAD", "")
38 + self.use_iomp = "libiomp" in ld_preload_str or "libomp" in ld_preload_str
39 + self.use_gomp = "libgomp" in ld_preload_str
40 +
41 + assert not (self.use_iomp and self.use_gomp)
42 +
43 + # at least reserve 1/local_world_size(for ARM) core for scheduler
44 + # proc as always use MP executor
45 + # TODO: make scheduler proc sleep when idle
46 + self.reserve_cpu_num = (
47 + self.local_world_size
48 + if current_platform.get_cpu_architecture() == CpuArchEnum.ARM
49 + else 1
50 + )
51 + # reserve at one more core for nixl_connector under p/d case
52 + if config.kv_transfer_config:
53 + self.reserve_cpu_num += 1
54 +
55 + if envs.VLLM_CPU_NUM_OF_RESERVED_CPU is not None:
56 + if self.reserve_cpu_num > envs.VLLM_CPU_NUM_OF_RESERVED_CPU:
57 + msg = (
58 + f"VLLM_CPU_NUM_OF_RESERVED_CPU is less than "
59 + "the minimum requirement"
60 + f": {self.reserve_cpu_num} cores"
61 + )
62 + logger.warning(msg=msg)
63 + self.reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU
64 +
65 + self._parse_omp_threads_bind_env()
66 +
67 + assert not self.simulate_multi_node or self.auto_setup
68 +
69 + @contextmanager
70 + def configure_omp_envs(self, rank: int, local_rank: int):
71 + if not current_platform.is_cpu() or self.skip_setup:
72 + yield
73 + return
74 +
75 + envs_dict = {}
76 + cpu_list = [str(i) for i in self.cpu_lists[local_rank]]
77 + envs_dict["OMP_NUM_THREADS"] = str(len(cpu_list))
78 + if self.use_iomp:
79 + # set IOMP envs
80 + cpu_list_str = ",".join(cpu_list)
81 + envs_dict["KMP_AFFINITY"] = (
82 + f"granularity=fine,explicit,proclist=[{cpu_list_str}]"
83 + )
84 + # The time(milliseconds) that a thread should wait after
85 + # completing the execution of a parallel region, before sleeping.
86 + envs_dict["KMP_BLOCKTIME"] = "1"
87 + # Prevents the CPU to run into low performance state
88 + envs_dict["KMP_TPAUSE"] = "0"
89 + # Provides fine granularity parallelism
90 + envs_dict["KMP_FORKJOIN_BARRIER_PATTERN"] = "dist,dist"
91 + envs_dict["KMP_PLAIN_BARRIER_PATTERN"] = "dist,dist"
92 + envs_dict["KMP_REDUCTION_BARRIER_PATTERN"] = "dist,dist"
93 + elif self.use_gomp:
94 + # set GOMP envs
95 + # likes '0 1 2 ...'
96 + cpu_list_str = " ".join(cpu_list)
97 + envs_dict["GOMP_CPU_AFFINITY"] = cpu_list_str
98 + else:
99 + # set OMP envs
100 + # likes '{0,1,2,...}'
101 + cpu_list_str = ",".join(cpu_list)
102 + envs_dict["OMP_PLACES"] = f"{{{cpu_list_str}}}"
103 + envs_dict["OMP_PROC_BIND"] = "true"
104 +
105 + # backup envs
106 + old_envs_dict = {}
107 + for k in envs_dict:
108 + old_envs_dict[k] = os.environ.get(k)
109 +
110 + try:
111 + # set envs
112 + for k, v in envs_dict.items():
113 + os.environ[k] = v
114 + yield
115 + finally:
116 + # restore old envs
117 + for k, v in old_envs_dict.items(): # type: ignore
118 + if v is None:
119 + os.environ.pop(k, None)
120 + else:
121 + os.environ[k] = v
122 +
123 + def _parse_omp_threads_bind_env(self):
124 + vllm_mask = envs.VLLM_CPU_OMP_THREADS_BIND
125 + self.skip_setup = vllm_mask == "nobind"
126 + self.auto_setup = vllm_mask == "auto"
127 + self.reserved_cpu_list = []
128 + self.cpu_lists = []
129 +
130 + if self.auto_setup:
131 + # auto generate CPU lists
132 + cpu_arch = current_platform.get_cpu_architecture()
133 + if cpu_arch == CpuArchEnum.POWERPC:
134 + # For POWERPC SMT-8/4/2
135 + cpu_list, reserve_list = self._get_autobind_cpu_ids(
136 + lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4]
137 + )
138 + elif cpu_arch in (CpuArchEnum.X86, CpuArchEnum.S390X):
139 + # For x86/S390X SMT-2, use 1 logical CPU per physical core
140 + cpu_list, reserve_list = self._get_autobind_cpu_ids(
141 + lambda cpus: cpus[-1:]
142 + )
143 + elif cpu_arch == CpuArchEnum.ARM:
144 + # For AArch64, no SMT, use all logical CPU
145 + cpu_list, reserve_list = self._get_autobind_cpu_ids(lambda cpus: cpus)
146 else:
169 - masks = [None]
170 - if mock is None:
171 - data = _get_cpu_topology_json()
147 + cpu_list, reserve_list = [], []
148 + raise RuntimeError(f"{cpu_arch} doesn't support auto CPU binding.")
149 +
150 + for item in cpu_list:
151 + self.cpu_lists.append([x.id for x in item])
152 + self.reserved_cpu_list = [x.id for x in reserve_list]
153 + elif not self.skip_setup:
154 + # user defined CPU lists
155 + omp_cpuids_list = vllm_mask.split("|")
156 + if self.local_dp_rank is not None:
157 + local_dp_rank = self.local_dp_rank
158 + world_size = self.local_world_size
159 + # Rank mapping [DP, PP, TP]
160 + omp_cpuids_list = omp_cpuids_list[
161 + local_dp_rank * world_size : (local_dp_rank + 1) * world_size
162 + ]
163 +
164 + assert len(omp_cpuids_list) == self.local_world_size, (
165 + "Given "
166 + f"number of CPU id list {omp_cpuids_list} doesn't match "
167 + f"local world size {self.local_world_size}."
168 + )
169 +
170 + # parse CPU list strings like "5,2-4" to [5, 2, 3, 4]
171 + self.cpu_lists = [cr_utils.parse_id_list(s) for s in omp_cpuids_list]
172 + else:
173 + # skip
174 + self.cpu_lists = []
175 +
176 + msg = "OpenMP thread binding info: \n"
177 + for i in range(self.local_world_size):
178 + msg += f"\tlocal_rank={i}, core ids={self.cpu_lists[i]}\n"
179 + msg += f"\treserved_cpus={self.reserved_cpu_list}"
180 + logger.info(msg)
181 +
182 + def _get_autobind_cpu_ids(
183 + self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]
184 + ) -> tuple[list[list[LogicalCPUInfo]], list[LogicalCPUInfo]]:
185 + """
186 + Return CPU ids to bind based on NUMA nodes, and CPU ids reserved for
187 + other processes.
188 + Currently for rank N, only CPU ids on the N-th node in available NUMA
189 + node list will be selected.
190 + Args:
191 + cpu_selector: a callable object to select CPUs from a CPU list
192 + of a physical core. The input is a LogicalCPUInfo list contains
193 + logical CPUs of a physical CPU, sorted by the LogicalCPUInfo.id.
194 + A selected LogicalCPUInfo list should be returned.
195 + """
196 +
197 + # this memory node list has been sliced for DP offset
198 + allowed_numa_nodes = cr_utils.get_visible_memory_node()
199 + logical_cpu_list = cr_utils.get_allowed_cpu_list()
200 +
201 + local_world_size = self.local_world_size
202 + assert (
203 + len(allowed_numa_nodes) >= local_world_size or self.simulate_multi_node
204 + ), (
205 + f"Not enough allowed NUMA nodes to bind threads of "
206 + f"{local_world_size} local CPUWorkers. "
207 + f"Allowed NUMA nodes are {allowed_numa_nodes}. "
208 + "Please try to bind threads manually or decrease DP/TP/PP."
209 + )
210 +
211 + # Generate OMP CPU list for each rank
212 + cpu_lists_of_ranks = []
213 + reserved_cpu_list = []
214 + total_cpu_num = 0
215 + for local_rank in range(self.local_world_size):
216 + if not self.simulate_multi_node:
217 + selected_numa_node = allowed_numa_nodes[local_rank]
218 + selected_logical_cpu_list = [
219 + x for x in logical_cpu_list if x.numa_node == selected_numa_node
220 + ]
221 else:
173 - with open(mock, mode="rb") as jf:
174 - data = jf.read()
175 - lscpu = json.loads(data)
176 - for mask in masks:
177 - resources = enumerate_resources(lscpu, mask, affinity)
178 - omp_places.extend(create_omp_places(resources, strategy, smt))
179 - self.omp_places = sorted(
180 - omp_places,
181 - key=lambda p: "{:04d}-{:04d}".format(len(p["mask"]), max(p["mask"])),
182 - reverse=True,
222 + world_size_across_dp = self.local_world_size * self.internal_dp_size
223 + assert len(logical_cpu_list) >= world_size_across_dp
224 + selected_logical_cpu_list = sorted(
225 + logical_cpu_list, key=lambda x: x.numa_node
226 + )
227 + sim_cpu_num_per_node = (
228 + len(selected_logical_cpu_list) // world_size_across_dp
229 + )
230 + assert self.local_dp_rank is not None
231 + start_idx = (
232 + local_rank + self.local_world_size * self.local_dp_rank
233 + ) * sim_cpu_num_per_node
234 + selected_logical_cpu_list = selected_logical_cpu_list[
235 + start_idx : (start_idx + sim_cpu_num_per_node)
236 + ]
237 +
238 + # Select logical CPUs on same physical cores via cpu_selector
239 + core_to_cpus: dict[int, list[LogicalCPUInfo]] = {}
240 + for cpu_info in selected_logical_cpu_list:
241 + if cpu_info.physical_core not in core_to_cpus:
242 + core_to_cpus[cpu_info.physical_core] = []
243 + core_to_cpus[cpu_info.physical_core].append(cpu_info)
244 + selected_logical_cpu_list = []
245 + for cpu_list in core_to_cpus.values():
246 + cpu_list = sorted(cpu_list, key=lambda x: x.id)
247 + selected_logical_cpu_list.extend(cpu_selector(cpu_list))
248 +
249 + # sort selected cores based on core id
250 + selected_logical_cpu_list = sorted(
251 + selected_logical_cpu_list, key=lambda x: x.id
252 )
253
185 - def run(self, what, *args, **kwargs):
186 - """Run arg with correct OMP environment"""
187 - if self.setup_omp:
188 - for place in self.omp_places:
189 - if place["available"]:
190 - reserve = int(os.environ.get("VLLM_CPU_NUM_OF_RESERVED_CPU", 0))
191 - place["available"] = False
192 - # pylint: disable=consider-using-f-string
193 - os.environ["OMP_PLACES"] = "{}".format(place["mask"])
194 - os.environ["OMP_NUM_THREADS"] = "{}".format(
195 - len(place["mask"]) - reserve
196 - )
197 - os.environ["OMP_PROC_BIND"] = "TRUE"
198 - return what(*args, **kwargs)
199 - raise IndexError("Out of OMP places")
200 - return what(*args, **kwargs)
254 + cpu_lists_of_ranks.append(selected_logical_cpu_list)
255 + total_cpu_num += len(selected_logical_cpu_list)
256 +
257 + # Reserve CPUs for other processes
258 + if total_cpu_num <= self.reserve_cpu_num:
259 + logger.warning(
260 + "Selected CPU core number (%s) "
261 + "should be greater than reserved CPU core "
262 + "number (%s).",
263 + total_cpu_num,
264 + self.reserve_cpu_num,
265 + )
266 + return cpu_lists_of_ranks, []
267 +
268 + reserve_num_per_rank = [
269 + self.reserve_cpu_num // self.local_world_size
270 + ] * self.local_world_size
271 + # last rank first
272 + for i in range(
273 + self.local_world_size - 1,
274 + self.local_world_size - 1 - self.reserve_cpu_num % self.local_world_size,
275 + -1,
276 + ):
277 + reserve_num_per_rank[i] += 1
278 + for i in range(self.local_world_size):
279 + num = reserve_num_per_rank[i]
280 + if num > 0:
281 + reserved_cpu_list.extend(cpu_lists_of_ranks[i][-num:])
282 + cpu_lists_of_ranks[i] = cpu_lists_of_ranks[i][:-num]
283 +
284 + return cpu_lists_of_ranks, reserved_cpu_list
vllm/v1/executor/multiproc_executor.py
+6 -15
@@ -51,6 +51,7 @@ from vllm.utils.network_utils import (
51 get_loopback_ip,
52 get_open_port,
53 )
54 +from vllm.utils.ompmultiprocessing import OMPProcessManager
55 from vllm.utils.system_utils import (
56 _maybe_force_spawn,
57 decorate_logs,
@@ -169,24 +170,14 @@ class MultiprocExecutor(Executor):
170 [] if context.get_start_method() == "fork" else None
171 )
172
173 + # For CPU backend only, to setup OpenMP threads affinity
174 + cpu_omp_manager = OMPProcessManager(self.vllm_config)
175 for local_rank in range(self.local_world_size):
176 global_rank = global_start_rank + local_rank
177 is_driver_worker = self._is_driver_worker(global_rank)
175 - if current_platform.is_cpu():
176 - om = current_platform.get_omp_manager()
177 - logger.info("Configured OMP PLACES %s", str(om.omp_places))
178 - unready_worker_handle = om.run(
179 - WorkerProc.make_worker_process,
180 - vllm_config=self.vllm_config,
181 - local_rank=local_rank,
182 - rank=global_rank,
183 - distributed_init_method=distributed_init_method,
184 - input_shm_handle=scheduler_output_handle,
185 - shared_worker_lock=shared_worker_lock,
186 - is_driver_worker=is_driver_worker,
187 - inherited_fds=inherited_fds,
188 - )
189 - else:
178 + with cpu_omp_manager.configure_omp_envs(
179 + rank=global_rank, local_rank=local_rank
180 + ):
181 unready_worker_handle = WorkerProc.make_worker_process(
182 vllm_config=self.vllm_config,
183 local_rank=local_rank,
vllm/v1/worker/cpu_model_runner.py
+1 -15
@@ -116,21 +116,7 @@ class CPUModelRunner(GPUModelRunner):
116 logger.info("Warming up model for the compilation...")
117 # Only generate graph for the generic shape
118 with _set_global_compilation_settings(self.vllm_config):
119 - self._dummy_run(
120 - min(
121 - max(16, self.max_num_reqs),
122 - self.scheduler_config.max_num_batched_tokens,
123 - )
124 - )
125 -
126 - # Warm up drafter for speculative decoding
127 - if self.speculative_config and (self.speculative_config.uses_draft_model()):
128 - from vllm.v1.spec_decode.draft_model import DraftModelProposer
129 -
130 - if isinstance(self.drafter, (DraftModelProposer)):
131 - logger.info("Warming up drafter model...")
132 - self.drafter.dummy_run(max(16, self.max_num_reqs))
133 -
119 + self.profile_run()
120 logger.info("Warming up done.")
121
122 def initialize_kv_cache(
vllm/v1/worker/cpu_worker.py
+106 -2
@@ -1,15 +1,23 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import math
4 import os
5 import sys
6 from typing import Any
7
8 +import psutil
9 import torch
10
11 from vllm.config import VllmConfig
12 from vllm.logger import init_logger
13 from vllm.platforms import CpuArchEnum, current_platform
14 from vllm.profiler.wrapper import TorchProfilerWrapper
15 +from vllm.utils.cpu_resource_utils import (
16 + get_allowed_cpu_list,
17 + get_memory_node_info,
18 + get_visible_memory_node,
19 +)
20 +from vllm.utils.mem_utils import format_gib
21 from vllm.utils.torch_utils import set_random_seed
22 from vllm.v1.worker.cpu_model_runner import CPUModelRunner
23 from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
@@ -27,6 +35,46 @@ class CPUWorker(Worker):
35 distributed_init_method: str,
36 is_driver_worker: bool = False,
37 ):
38 + # TODO: use numactl for process setup
39 + # TODO: optimize for `interleaved` policy
40 + # Bind memory node
41 + allowed_memory_nodes = get_visible_memory_node()
42 + allowed_cpu_list = get_allowed_cpu_list()
43 + cpu_core = allowed_cpu_list[0]
44 +
45 + # TODO: some CI hosts are not correctly set, change to assertion
46 + # after fix
47 + if cpu_core.numa_node not in allowed_memory_nodes:
48 + logger.warning(
49 + "Node %s is not in available memory nodes %s.",
50 + cpu_core.numa_node,
51 + allowed_memory_nodes,
52 + )
53 +
54 + torch.ops._C.init_cpu_memory_env([cpu_core.numa_node])
55 +
56 + memory_status = get_memory_node_info(cpu_core.numa_node)
57 + memory_fraction = vllm_config.cache_config.gpu_memory_utilization
58 + self.requested_cpu_memory = math.ceil(
59 + memory_status.total_memory * memory_fraction
60 + )
61 + available_memory = memory_status.available_memory
62 +
63 + if (
64 + vllm_config.cache_config.kv_cache_memory_bytes is None
65 + and self.requested_cpu_memory > available_memory
66 + ):
67 + raise ValueError(
68 + f"Available memory on node {cpu_core.numa_node} "
69 + f"({format_gib(available_memory)}/"
70 + f"{format_gib(memory_status.total_memory)} GiB) on startup "
71 + f"is less than desired CPU memory utilization "
72 + f"({vllm_config.cache_config.gpu_memory_utilization}, "
73 + f"{format_gib(self.requested_cpu_memory)} GiB). "
74 + "Decrease --gpu-memory-utilization"
75 + f" or reduce CPU memory used by other processes."
76 + )
77 +
78 super().__init__(
79 vllm_config,
80 local_rank,
@@ -103,13 +151,69 @@ class CPUWorker(Worker):
151 pass
152
153 def determine_available_memory(self) -> int:
106 - return self.cache_config.cpu_kvcache_space_bytes or 0
154 + self.model_runner.warming_up_model()
155 +
156 + allowed_cpu_list = get_allowed_cpu_list()
157 + cpu_core = allowed_cpu_list[0]
158 +
159 + memory_status = get_memory_node_info(cpu_core.numa_node)
160 + available_memory = memory_status.available_memory
161 + explicit_kv_cache_size = self.cache_config.kv_cache_memory_bytes
162 +
163 + kv_cache_size = None
164 + msg = None
165 + if explicit_kv_cache_size is not None:
166 + if explicit_kv_cache_size > available_memory:
167 + raise ValueError(
168 + f"Available memory on node {cpu_core.numa_node} "
169 + f"({format_gib(available_memory)}/"
170 + f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
171 + f" allocation is less than requested memory for kv "
172 + f"({format_gib(explicit_kv_cache_size)} GiB). "
173 + "Decrease --kv-cache-memory-bytes, VLLM_CPU_KVCACHE_SPACE, "
174 + "or reduce CPU memory used by other processes."
175 + )
176 + kv_cache_size = explicit_kv_cache_size
177 + msg = (
178 + f"Explicitly set ({format_gib(kv_cache_size)}/"
179 + f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
180 + f"on node {cpu_core.numa_node}."
181 + )
182 + else:
183 + consumed_memory = psutil.Process(os.getpid()).memory_info().rss
184 + requested_memory_for_kv = int(self.requested_cpu_memory - consumed_memory)
185 + if (
186 + requested_memory_for_kv <= 0
187 + or requested_memory_for_kv > available_memory
188 + ):
189 + raise ValueError(
190 + f"Available memory on node {cpu_core.numa_node} "
191 + f"({format_gib(available_memory)}/"
192 + f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
193 + f" allocation is less than requested memory for kv "
194 + f"({format_gib(requested_memory_for_kv)}/"
195 + f"{format_gib(self.requested_cpu_memory)} GiB). "
196 + "Reduce CPU memory used by other processes."
197 + )
198 + kv_cache_size = requested_memory_for_kv
199 + msg = (
200 + f"Auto set ({format_gib(kv_cache_size)}/"
201 + f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
202 + f"on node {cpu_core.numa_node}, with "
203 + f"{format_gib(self.requested_cpu_memory)} GiB requested memory"
204 + f" for the worker. {format_gib(consumed_memory)} GiB"
205 + f" memory was consumed by non-kv usages."
206 + )
207 +
208 + logger.info(msg)
209 +
210 + return kv_cache_size
211
212 def compile_or_warm_up_model(self) -> CompilationTimes:
213 # Reset the seed to ensure that the random state is not affected by
214 # the model initialization and profiling.
215 set_random_seed(self.model_config.seed)
112 - self.model_runner.warming_up_model()
216 + # Note: the model has been compiled in determine_available_memory()
217 return CompilationTimes(
218 language_model=self.compilation_config.compilation_time,
219 encoder=self.compilation_config.encoder_compilation_time,