@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / abebd9323d

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, 14:45 UTC

[CPU] Replace OMP initialization (#36487)

Signed-off-by: Anton Ivanov <anton.ivanov@cambridgegreys.com>

Anton Ivanov committed Apr 3, 2026 at 11:42 UTC abebd9323d7cf8abefd07f976be3be356d873075

7 files changed +321 -426

.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

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

csrc/cpu/torch_bindings.cpp

-3

@@ -8,8 +8,6 @@

8 // libraries use different ISAs.

9 #define TORCH_EXTENSION_NAME _C

11 -std::string init_cpu_threads_env(const std::string& cpu_ids);

12 -

11 void release_dnnl_matmul_handler(int64_t handler);

13 int64_t create_onednn_scaled_mm_handler(const torch::Tensor& b,

@@ -354,7 +352,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

352 "str act, str isa) -> ()");

353 ops.impl("cpu_fused_moe", torch::kCPU, &cpu_fused_moe);

354 #endif

357 - ops.def("init_cpu_threads_env(str cpu_ids) -> str", &init_cpu_threads_env);

355 ops.def(

356 "mla_decode_kvcache("

357 " Tensor! out, Tensor query, Tensor kv_cache,"

csrc/cpu/utils.cpp

-144

@@ -21,150 +21,6 @@ std::string init_cpu_threads_env(const std::string& cpu_ids) {

22 #endif

24 -#ifndef VLLM_NUMA_DISABLED

25 -std::string init_cpu_threads_env(const std::string& cpu_ids) {

26 - bitmask* omp_cpu_mask = numa_parse_cpustring_all(cpu_ids.c_str());

27 - TORCH_CHECK(omp_cpu_mask != nullptr,

28 - "Failed to parse CPU string: " + cpu_ids);

29 - TORCH_CHECK(omp_cpu_mask->size > 0);

30 - std::vector<int> omp_cpu_ids;

31 - omp_cpu_ids.reserve(omp_cpu_mask->size);

32 -

33 - constexpr int group_size = 8 * sizeof(*omp_cpu_mask->maskp);

34 -

35 - for (int offset = 0; offset < omp_cpu_mask->size; offset += group_size) {

36 - unsigned long group_mask = omp_cpu_mask->maskp[offset / group_size];

37 - int i = 0;

38 - while (group_mask) {

39 - if (group_mask & 1) {

40 - omp_cpu_ids.emplace_back(offset + i);

41 - }

42 - ++i;

43 - group_mask >>= 1;

44 - }

45 - }

46 -

47 - // Memory node binding

48 - if (numa_available() != -1) {

49 - std::set<int> node_ids;

50 - for (const auto& cpu_id : omp_cpu_ids) {

51 - int node_id = numa_node_of_cpu(cpu_id);

52 - if (node_id != -1) {

53 - node_ids.insert(node_id);

54 - }

55 - }

56 - // Concatenate all node_ids into a single comma-separated string

57 - if (!node_ids.empty()) {

58 - std::string node_ids_str;

59 - for (const int node_id : node_ids) {

60 - if (!node_ids_str.empty()) {

61 - node_ids_str += ",";

62 - }

63 - node_ids_str += std::to_string(node_id);

64 - }

65 -

66 - bitmask* mask = numa_parse_nodestring(node_ids_str.c_str());

67 - bitmask* src_mask = numa_get_mems_allowed();

68 -

69 - int pid = getpid();

70 -

71 - if (mask && src_mask) {

72 - // move all existing pages to the specified numa node.

73 - *(src_mask->maskp) = *(src_mask->maskp) ^ *(mask->maskp);

74 - int page_num = numa_migrate_pages(pid, src_mask, mask);

75 - if (page_num == -1) {

76 - TORCH_WARN("numa_migrate_pages failed. errno: " +

77 - std::to_string(errno));

78 - }

79 -

80 - // Restrict memory allocation to the selected NUMA node(s).

81 - // Enhances memory locality for the threads bound to those NUMA CPUs.

82 - if (node_ids.size() > 1) {

83 - errno = 0;

84 - numa_set_interleave_mask(mask);

85 - if (errno != 0) {

86 - TORCH_WARN("numa_set_interleave_mask failed. errno: " +

87 - std::to_string(errno));

88 - } else {

89 - TORCH_WARN(

90 - "NUMA binding: Using INTERLEAVE policy for memory "

91 - "allocation across multiple NUMA nodes (nodes: " +

92 - node_ids_str +

93 - "). Memory allocations will be "

94 - "interleaved across the specified NUMA nodes.");

95 - }

96 - } else {

97 - errno = 0;

98 - numa_set_membind(mask);

99 - if (errno != 0) {

100 - TORCH_WARN("numa_set_membind failed. errno: " +

101 - std::to_string(errno));

102 - } else {

103 - TORCH_WARN(

104 - "NUMA binding: Using MEMBIND policy for memory "

105 - "allocation on the NUMA nodes (" +

106 - node_ids_str +

107 - "). Memory allocations will be "

108 - "strictly bound to these NUMA nodes.");

109 - }

110 - }

111 -

112 - numa_set_strict(1);

113 -

114 - numa_free_nodemask(mask);

115 - numa_free_nodemask(src_mask);

116 - } else {

117 - TORCH_WARN(

118 - "numa_parse_nodestring or numa_get_run_node_mask failed. errno: " +

119 - std::to_string(errno));

120 - }

121 - }

122 - }

123 -

124 - // OMP threads binding

125 - omp_set_num_threads((int)omp_cpu_ids.size());

126 - torch::set_num_threads((int)omp_cpu_ids.size());

127 - TORCH_CHECK_EQ(omp_cpu_ids.size(), torch::get_num_threads());

128 - TORCH_CHECK_EQ(omp_cpu_ids.size(), omp_get_max_threads());

129 -

130 - std::vector<std::pair<int, int>> thread_core_mapping;

131 - thread_core_mapping.reserve(omp_cpu_ids.size());

132 - omp_lock_t writelock;

133 - omp_init_lock(&writelock);

134 -

135 - #pragma omp parallel for schedule(static, 1)

136 - for (size_t i = 0; i < omp_cpu_ids.size(); ++i) {

137 - cpu_set_t mask;

138 - CPU_ZERO(&mask);

139 - CPU_SET(omp_cpu_ids[i], &mask);

140 - int ret = sched_setaffinity(0, sizeof(cpu_set_t), &mask);

141 - if (ret == -1) {

142 - TORCH_CHECK(false,

143 - "sched_setaffinity failed. errno: " + std::to_string(errno));

144 - }

145 -

146 - omp_set_lock(&writelock);

147 - thread_core_mapping.emplace_back(gettid(), omp_cpu_ids[i]);

148 - omp_unset_lock(&writelock);

149 - }

150 -

151 - omp_destroy_lock(&writelock);

152 -

153 - numa_free_nodemask(omp_cpu_mask);

154 -

155 - std::stringstream ss;

156 - ss << "OMP threads binding of Process " << getpid() << ":\n";

157 - std::sort(thread_core_mapping.begin(), thread_core_mapping.end(),

158 - [](auto&& a, auto&& b) { return a.second < b.second; });

159 - for (auto&& item : thread_core_mapping) {

160 - ss << "\t"

161 - << "OMP tid: " << item.first << ", core " << item.second << "\n";

162 - }

163 -

164 - return ss.str();

165 -}

166 -#endif // VLLM_NUMA_DISABLED

167 -

24 namespace cpu_utils {

25 ScratchPadManager::ScratchPadManager() : size_(0), ptr_(nullptr) {

26 this->realloc(allocation_unit * 128);

vllm/platforms/cpu.py

+80 -100

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import glob

5 -import json

5 import os

6 import platform

7 import subprocess

@@ -11,11 +10,11 @@ from dataclasses import dataclass

10 from typing import TYPE_CHECKING

12 import psutil

14 -import regex as re

13 import torch

15 from vllm import envs

16 from vllm.logger import init_logger

17 +from vllm.utils.ompmultiprocessing import OMPProcessManager

18 from vllm.utils.torch_utils import is_quantized_kv_cache

19 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -76,6 +75,10 @@ class CpuPlatform(Platform):

75 dispatch_key: str = "CPU"

76 dist_backend: str = "gloo"

77 device_control_env_var = "CPU_VISIBLE_MEMORY_NODES"

78 + omp_process_manager = None

79 + smt = 1 # SMT level for OMP - 4 threads on PowerPC, 1 on others

80 + global_cpu_mask = None

81 + simulate_numa = int(os.environ.get("_SIM_MULTI_NUMA", 0))

83 @property

84 def supported_dtypes(self) -> list[torch.dtype]:

@@ -191,26 +194,10 @@ class CpuPlatform(Platform):

194

195 cache_config.cpu_kvcache_space_bytes = CpuPlatform.get_device_total_memory()

196

194 - # reserve at least one core for nixl_connector under p/d case

195 - if vllm_config.kv_transfer_config and (

196 - envs.VLLM_CPU_NUM_OF_RESERVED_CPU == 0

197 - or envs.VLLM_CPU_NUM_OF_RESERVED_CPU is None

198 - ):

199 - os.environ["VLLM_CPU_NUM_OF_RESERVED_CPU"] = "1"

200 -

197 parallel_config = vllm_config.parallel_config

202 - if (

203 - parallel_config.world_size > 1

204 - and parallel_config.distributed_executor_backend is not None

205 - and parallel_config.distributed_executor_backend != "mp"

206 - ):

207 - logger.warning(

208 - (

209 - "%s is not supported on CPU, fallback to mp "

210 - "distributed executor backend."

211 - ),

212 - parallel_config.distributed_executor_backend,

213 - )

198 + # OMP requires the MP executor to function correctly, UniProc is not

199 + # supported as it is not possible to set the OMP environment correctly

200 + if parallel_config.distributed_executor_backend == "uni":

201 parallel_config.distributed_executor_backend = "mp"

202 if parallel_config.worker_cls == "auto":

203 parallel_config.worker_cls = "vllm.v1.worker.cpu_worker.CPUWorker"

@@ -267,14 +254,6 @@ class CpuPlatform(Platform):

254 # variable "NUMEXPR_MAX_THREADS" (64)'.

255 os.environ["NUMEXPR_MAX_THREADS"] = str(get_max_threads())

256

270 - if envs.VLLM_CPU_OMP_THREADS_BIND != "nobind":

271 - # Set default threads num for OpenMP parallel

272 - os.environ["OMP_NUM_THREADS"] = str(torch.get_num_threads())

273 - else:

274 - # In this case, setting the OpenMP configuration via

275 - # OMP_NUM_THREADS is up to the user.

276 - logger.info("Disabling binding processes to CPU cores...")

277 -

257 # Disable torch async compiling which won't work with daemonic processes

258 os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"

259

@@ -286,8 +265,8 @@ class CpuPlatform(Platform):

265

266 ld_preload_str = os.getenv("LD_PRELOAD", "")

267

289 - # Intel OpenMP setting

290 - if "libiomp5.so" in ld_preload_str:

268 + # Intel and CLANG OpenMP setting

269 + if "libiomp5.so" in ld_preload_str or "libomp5" in ld_preload_str:

270 # The time(milliseconds) that a thread should wait after

271 # completing the execution of a parallel region, before sleeping.

272 os.environ["KMP_BLOCKTIME"] = "1"

@@ -324,37 +303,6 @@ class CpuPlatform(Platform):

303 ld_preload_str = tcmalloc_so

304 os.environ["LD_PRELOAD"] = ld_preload_str

305

327 - if (

328 - platform.system() == "Linux"

329 - and cpu_architecture in (CpuArchEnum.ARM, CpuArchEnum.POWERPC)

330 - and not ("libomp" in ld_preload_str or "libgomp" in ld_preload_str)

331 - ):

332 - # We need to LD_PRELOAD PyTorch's libgomp, otherwise only

333 - # one core will be properly utilized when we thread-bind

334 - # See: https://github.com/vllm-project/vllm/issues/27369

335 - # TODO: Remove once:

336 - # https://github.com/pytorch/pytorch/issues/166087 is fixed

337 -

338 - # We need to find the location of PyTorch's libgomp

339 - torch_pkg = os.path.dirname(torch.__file__)

340 - site_root = os.path.dirname(torch_pkg)

341 - # Search both torch.libs and torch/lib - See: https://github.com/vllm-project/vllm/issues/30470

342 - torch_libs_paths = [

343 - os.path.join(site_root, "torch.libs"),

344 - os.path.join(torch_pkg, "lib"),

345 - ]

346 - pytorch_libgomp_so_candidates = []

347 - for torch_libs in torch_libs_paths:

348 - pytorch_libgomp_so_candidates.extend(

349 - glob.glob(os.path.join(torch_libs, "libgomp*.so*"))

350 - )

351 - if pytorch_libgomp_so_candidates:

352 - pytorch_libgomp_so = pytorch_libgomp_so_candidates[0]

353 - if ld_preload_str:

354 - ld_preload_str += ":"

355 - ld_preload_str += pytorch_libgomp_so

356 - os.environ["LD_PRELOAD"] = ld_preload_str

357 -

306 os.environ["LOCAL_WORLD_SIZE"] = str(

307 vllm_config.parallel_config.tensor_parallel_size

308 )

@@ -369,6 +317,13 @@ class CpuPlatform(Platform):

317 vllm_config.model_config.max_model_len,

318 vllm_config.scheduler_config.DEFAULT_MAX_NUM_BATCHED_TOKENS,

319 )

320 + # CI specific "quick" NUMA simulation - split all available CPUs

321 + # into a fake NUMA topology

322 + if os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", None) is not None:

323 + os.environ["_SIM_MULTI_NUMA"] = str(

324 + vllm_config.parallel_config.world_size

325 + * vllm_config.parallel_config._api_process_count

326 + )

327

328 @classmethod

329 def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None:

@@ -377,46 +332,71 @@ class CpuPlatform(Platform):

332 pass

333

334 @classmethod

380 - def get_allowed_cpu_core_node_list(cls) -> tuple[list[int], list[LogicalCPUInfo]]:

381 - assert platform.system() == "Linux"

335 + def get_omp_manager(cls) -> OMPProcessManager:

336 + # initialise the OMP resource management if need be and return the manager

337 + if cls.omp_process_manager is None:

338 + if cls.get_cpu_architecture() == CpuArchEnum.POWERPC:

339 + cls.smt = 4

340 + cls.omp_process_manager = OMPProcessManager(

341 + affinity=cls.get_global_cpu_mask(), smt=cls.smt

342 + )

343 + # we need to fix up the topology returned by the OMP Manager for

344 + # simulated NUMA environments in CI

345 + if cls.simulate_numa > 0:

346 + logger.info(

347 + "Adjusting numa topology to resemble at least %d nodes",

348 + int(cls.simulate_numa),

349 + )

350 + om = cls.omp_process_manager

351 + while len(om.omp_places) < cls.simulate_numa:

352 + new_omp_places = []

353 + touched = False

354 + for omp_place in om.omp_places:

355 + if len(omp_place["mask"]) > 1:

356 + touched = True

357 + cpu_list = sorted(list(omp_place["mask"]))

358 + new_omp_places.append(

359 + {

360 + "mask": set(cpu_list[0 : int(len(cpu_list) / 2)]),

361 + "available": True,

362 + }

363 + )

364 + new_omp_places.append(

365 + {

366 + "mask": set(cpu_list[int(len(cpu_list) / 2) :]),

367 + "available": True,

368 + }

369 + )

370 + if touched:

371 + om.omp_places = new_omp_places

372 + else:

373 + raise ValueError(

374 + "Cannot split the existing NUMA topology to match "

375 + "simulation requirements"

376 + )

377

383 - # Init LogicalCPUInfo from lscpu

384 - lscpu_output = subprocess.check_output(

385 - "lscpu -J -e=CPU,CORE,NODE", shell=True, text=True

378 + return cls.omp_process_manager

379 +

380 + @classmethod

381 + def get_global_cpu_mask(cls) -> set[int]:

382 + # get global cpu mask

383 + if cls.global_cpu_mask is None:

384 + cls.global_cpu_mask = os.sched_getaffinity(0)

385 + return cls.global_cpu_mask

386 +

387 + @classmethod

388 + def reserve_cpus(cls, reserve: set[int]) -> bool:

389 + # remove CPUs from global mask, for now there is no "release" mechanism

390 + if cls.omp_process_manager is not None:

391 + for place in cls.omp_process_manager.omp_places:

392 + if not place["available"]:

393 + return False

394 + cls.global_cpu_mask = cls.get_global_cpu_mask() - reserve

395 + # reinitialize OMP resource management

396 + cls.omp_process_manager = OMPProcessManager(

397 + affinity=cls.global_cpu_mask, smt=cls.smt

398 )

387 - lscpu_output = re.sub(r'"node":\s*-\s*(,|\n)', r'"node": 0\1', lscpu_output)

388 - logical_cpu_list: list[LogicalCPUInfo] = json.loads(

389 - lscpu_output, object_hook=LogicalCPUInfo.json_decoder

390 - )["cpus"]

391 -

392 - # Filter CPUs with invalid attributes

393 - logical_cpu_list = [

394 - x

395 - for x in logical_cpu_list

396 - if -1 not in (x.id, x.physical_core, x.numa_node)

397 - ]

398 -

399 - # Filter allowed CPUs

400 - if hasattr(os, "sched_getaffinity"):

401 - allowed_cpu_id_list = os.sched_getaffinity(0)

402 - else:

403 - raise NotImplementedError("Unsupported OS")

404 - logical_cpu_list = [x for x in logical_cpu_list if x.id in allowed_cpu_id_list]

405 -

406 - # Get allowed NUMA nodes

407 - allowed_numa_nodes = set()

408 - for x in logical_cpu_list:

409 - allowed_numa_nodes.add(x.numa_node) # type: ignore

410 - allowed_numa_nodes_list = sorted(allowed_numa_nodes)

411 -

412 - env_key = CpuPlatform.device_control_env_var

413 - if env_key in os.environ and os.environ[env_key] != "":

414 - visible_nodes = [int(s) for s in os.environ[env_key].split(",")]

415 - allowed_numa_nodes_list = [

416 - x for x in sorted(list(set(visible_nodes))) if x in allowed_numa_nodes

417 - ]

418 -

419 - return allowed_numa_nodes_list, logical_cpu_list

399 + return True

400

401 @classmethod

402 def discover_numa_topology(cls) -> list[list[int]]:

vllm/utils/ompmultiprocessing.py new

+174

@@ -0,0 +1,174 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""OMP Aware Multiprocessing manager for running multiprocessing.Process()

6 +"""

7 +

8 +import json

9 +import os

10 +import subprocess

11 +

12 +

13 +def _int(arg):

14 + """Relaxed parsing of ints which handles a - instead of a number.

15 + The lscpu json may contain that for nodes in some cases. If that

16 + is the case we parse it to zero

17 + """

18 + try:

19 + if int(arg) >= 0:

20 + return int(arg)

21 + except ValueError:

22 + pass

23 + return 0

24 +

25 +

26 +def parse_mask(mask):

27 + """Expand a X-Y,Z list"""

28 + result = []

29 + for token in mask.split(","):

30 + try:

31 + start, finish = token.split("-")

32 + if int(start) > int(finish):

33 + raise IndexError("Invalid Indexes for cpu ranges")

34 + for cpu in range(int(start), int(finish) + 1):

35 + result.append(cpu)

36 + except ValueError:

37 + result.append(int(token))

38 + return set(result)

39 +

40 +

41 +def enumerate_resources(resource_map, mask=None, allowed=None):

42 + """Enumerate system resources"""

43 + if allowed is None:

44 + allowed = os.sched_getaffinity(0)

45 + if mask is not None:

46 + allowed = allowed & mask

47 +

48 + try:

49 + allowed_nodes = parse_mask(os.environ["CPU_VISIBLE_MEMORY_NODES"])

50 + except KeyError:

51 + allowed_nodes = None

52 +

53 + lscpu: dict[str, dict] = {"cpus": {}, "cores": {}, "nodes": {}}

54 + for cpu in resource_map["cpus"]:

55 + cpunum = int(cpu["cpu"])

56 + if (

57 + cpunum in allowed

58 + and cpunum >= 0

59 + and (allowed_nodes is None or _int(cpu["node"]) in allowed_nodes)

60 + ):

61 + lscpu["cpus"][cpunum] = [cpu]

62 + core = _int(cpu["core"])

63 + if lscpu["cores"].get(core, None) is None:

64 + lscpu["cores"][core] = [cpu]

65 + else:

66 + lscpu["cores"][core].append(cpu)

67 + node = _int(cpu["node"])

68 + if lscpu["nodes"].get(node, None) is None:

69 + lscpu["nodes"][node] = [cpu]

70 + else:

71 + lscpu["nodes"][node].append(cpu)

72 + return lscpu

73 +

74 +

75 +def produce_cpu_list(cpus, smt=1):

76 + """Produce a CPU list with/without SMT pairs - main cpu list case"""

77 + mask: list[int] = []

78 + for key, value in cpus.items():

79 + exists = 0

80 + for cpu in mask:

81 + if cpu == value[0]["core"]:

82 + exists += 1

83 + break

84 + if exists < smt:

85 + mask.append(int(key))

86 + return {"mask": set(mask), "available": True}

87 +

88 +

89 +def produce_cpu_sublist(scpus, smt=1):

90 + """Produce a CPU list with/without SMT pairs - resource leaf case"""

91 + cpu_list: list[dict] = []

92 + for value in scpus:

93 + exists = 0

94 + for cpu in cpu_list:

95 + if int(cpu["core"]) == int(value["core"]):

96 + exists += 1

97 + break

98 + if exists < smt:

99 + cpu_list.append(value)

100 + mask = []

101 + for cpu in cpu_list:

102 + mask.append(int(cpu["cpu"]))

103 +

104 + return {"mask": set(mask), "available": True}

105 +

106 +

107 +def create_omp_places(resources, strategy, smt=True):

108 + """Parse CPU topology and generate possible CPU masks"""

109 + omp_places = []

110 + if strategy == "all":

111 + omp_places.append(produce_cpu_list(resources["cpus"], smt))

112 + elif strategy == "cores":

113 + for value in resources["cores"].values():

114 + omp_places.append(produce_cpu_sublist(value, smt))

115 + elif strategy == "nodes":

116 + for value in resources["nodes"].values():

117 + omp_places.append(produce_cpu_sublist(value, smt))

118 + else:

119 + raise NotImplementedError("Unknown strategy")

120 +

121 + return omp_places

122 +

123 +

124 +# pylint: disable=too-few-public-methods

125 +class OMPProcessManager:

126 + """OMP aware wrapper to run mp Process()"""

127 +

128 + def __init__(self, strategy="nodes", smt=1, mock=None, affinity=None):

129 + self.strategy = strategy

130 + self.smt = smt

131 + self.omp_places = []

132 + vllm_mask = os.environ.get("VLLM_CPU_OMP_THREADS_BIND", None)

133 + self.setup_omp = vllm_mask != "nobind"

134 + if self.setup_omp:

135 + omp_places = []

136 + if vllm_mask is not None:

137 + masks = []

138 + for spec in vllm_mask.split("|"):

139 + masks.append(parse_mask(spec))

140 + else:

141 + masks = [None]

142 + if mock is None:

143 + data = subprocess.run(

144 + ["lscpu", "-Je"], check=True, capture_output=True

145 + ).stdout

146 + else:

147 + with open(mock, mode="rb") as jf:

148 + data = jf.read()

149 + lscpu = json.loads(data)

150 + for mask in masks:

151 + resources = enumerate_resources(lscpu, mask, affinity)

152 + omp_places.extend(create_omp_places(resources, strategy, smt))

153 + self.omp_places = sorted(

154 + omp_places,

155 + key=lambda p: "{:04d}-{:04d}".format(len(p["mask"]), max(p["mask"])),

156 + reverse=True,

157 + )

158 +

159 + def run(self, what, *args, **kwargs):

160 + """Run arg with correct OMP environment"""

161 + if self.setup_omp:

162 + for place in self.omp_places:

163 + if place["available"]:

164 + reserve = int(os.environ.get("VLLM_CPU_NUM_OF_RESERVED_CPU", 0))

165 + place["available"] = False

166 + # pylint: disable=consider-using-f-string

167 + os.environ["OMP_PLACES"] = "{}".format(place["mask"])

168 + os.environ["OMP_NUM_THREADS"] = "{}".format(

169 + len(place["mask"]) - reserve

170 + )

171 + os.environ["OMP_PROC_BIND"] = "TRUE"

172 + return what(*args, **kwargs)

173 + raise IndexError("Out of OMP places")

174 + return what(*args, **kwargs)

vllm/v1/executor/multiproc_executor.py

+48 -32

@@ -119,7 +119,6 @@ class MultiprocExecutor(Executor):

119 f"_parallel_size ({pcp_size}). "

120 )

121

122 - # Set multiprocessing envs

122 set_multiprocessing_worker_envs()

123

124 # use the loopback address get_loopback_ip() for communication.

@@ -172,16 +171,31 @@ class MultiprocExecutor(Executor):

171 for local_rank in range(self.local_world_size):

172 global_rank = global_start_rank + local_rank

173 is_driver_worker = self._is_driver_worker(global_rank)

175 - unready_worker_handle = WorkerProc.make_worker_process(

176 - vllm_config=self.vllm_config,

177 - local_rank=local_rank,

178 - rank=global_rank,

179 - distributed_init_method=distributed_init_method,

180 - input_shm_handle=scheduler_output_handle,

181 - shared_worker_lock=shared_worker_lock,

182 - is_driver_worker=is_driver_worker,

183 - inherited_fds=inherited_fds,

184 - )

174 + if current_platform.is_cpu():

175 + om = current_platform.get_omp_manager()

176 + logger.info("Configured OMP PLACES %s", str(om.omp_places))

177 + unready_worker_handle = om.run(

178 + WorkerProc.make_worker_process,

179 + vllm_config=self.vllm_config,

180 + local_rank=local_rank,

181 + rank=global_rank,

182 + distributed_init_method=distributed_init_method,

183 + input_shm_handle=scheduler_output_handle,

184 + shared_worker_lock=shared_worker_lock,

185 + is_driver_worker=is_driver_worker,

186 + inherited_fds=inherited_fds,

187 + )

188 + else:

189 + unready_worker_handle = WorkerProc.make_worker_process(

190 + vllm_config=self.vllm_config,

191 + local_rank=local_rank,

192 + rank=global_rank,

193 + distributed_init_method=distributed_init_method,

194 + input_shm_handle=scheduler_output_handle,

195 + shared_worker_lock=shared_worker_lock,

196 + is_driver_worker=is_driver_worker,

197 + inherited_fds=inherited_fds,

198 + )

199 unready_workers.append(unready_worker_handle)

200 if inherited_fds is not None:

201 inherited_fds.append(unready_worker_handle.death_writer.fileno())

@@ -1000,24 +1014,26 @@ def set_multiprocessing_worker_envs():

1014

1015 _maybe_force_spawn()

1016

1003 - # Configure thread parallelism if OMP_NUM_THREADS isn't set

1004 - #

1005 - # Helps to avoid CPU contention. The default of spawning a thread per

1006 - # core combined with multiprocessing for each GPU can have a negative

1007 - # impact on performance. The contention is amplified when running in a

1008 - # container where CPU limits can cause throttling.

1009 - default_omp_num_threads = 1

1010 - if (

1011 - "OMP_NUM_THREADS" not in os.environ

1012 - and (current_parallelism := torch.get_num_threads()) > default_omp_num_threads

1013 - ):

1014 - logger.warning_once(

1015 - "Reducing Torch parallelism from %d threads to %d to avoid "

1016 - "unnecessary CPU contention. Set OMP_NUM_THREADS in the "

1017 - "external environment to tune this value as needed.",

1018 - current_parallelism,

1019 - default_omp_num_threads,

1020 - scope="local",

1021 - )

1022 - os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads)

1023 - torch.set_num_threads(default_omp_num_threads)

1017 + if not current_platform.is_cpu():

1018 + # Configure thread parallelism if OMP_NUM_THREADS isn't set

1019 + #

1020 + # Helps to avoid CPU contention. The default of spawning a thread per

1021 + # core combined with multiprocessing for each GPU can have a negative

1022 + # impact on performance. The contention is amplified when running in a

1023 + # container where CPU limits can cause throttling.

1024 + default_omp_num_threads = 1

1025 + if (

1026 + "OMP_NUM_THREADS" not in os.environ

1027 + and (current_parallelism := torch.get_num_threads())

1028 + > default_omp_num_threads

1029 + ):

1030 + logger.warning_once(

1031 + "Reducing Torch parallelism from %d threads to %d to avoid "

1032 + "unnecessary CPU contention. Set OMP_NUM_THREADS in the "

1033 + "external environment to tune this value as needed.",

1034 + current_parallelism,

1035 + default_omp_num_threads,

1036 + scope="local",

1037 + )

1038 + os.environ["OMP_NUM_THREADS"] = str(default_omp_num_threads)

1039 + torch.set_num_threads(default_omp_num_threads)

vllm/v1/worker/cpu_worker.py

-128

@@ -1,18 +1,14 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import os

4 -import platform

4 import sys

6 -from collections.abc import Callable

5 from typing import Any

7 import torch

11 -from vllm import envs

9 from vllm.config import VllmConfig

10 from vllm.logger import init_logger

11 from vllm.platforms import CpuArchEnum, current_platform

15 -from vllm.platforms.cpu import CpuPlatform, LogicalCPUInfo

12 from vllm.profiler.wrapper import TorchProfilerWrapper

13 from vllm.utils.torch_utils import set_random_seed

14 from vllm.v1.worker.cpu_model_runner import CPUModelRunner

@@ -71,44 +67,6 @@ class CPUWorker(Worker):

67 if current_platform.get_cpu_architecture() == CpuArchEnum.X86:

68 check_preloaded_libs("libiomp")

74 - # Setup OpenMP threads affinity.

75 - omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND

76 - # Under numa binding some cores reserved for kv transfer in nixl_connector.py

77 - if omp_cpuids == "auto" and platform.system() == "Linux":

78 - cpu_arch = current_platform.get_cpu_architecture()

79 - if cpu_arch in (CpuArchEnum.POWERPC, CpuArchEnum.S390X):

80 - # For S390X/POWERPC SMT-8/4/2

81 - self.local_omp_cpuid = self._get_autobind_cpu_ids(

82 - lambda cpus: [cpu for cpu in cpus if cpu.id % 8 < 4]

83 - )

84 - elif cpu_arch == CpuArchEnum.X86:

85 - # For x86 SMT-2, use 1 CPU per core

86 - self.local_omp_cpuid = self._get_autobind_cpu_ids(

87 - lambda cpus: cpus[-1:]

88 - )

89 - elif cpu_arch == CpuArchEnum.ARM:

90 - # For AArch64, no SMT

91 - self.local_omp_cpuid = self._get_autobind_cpu_ids(lambda cpus: cpus)

92 - else:

93 - self.local_omp_cpuid = "nobind"

94 - elif omp_cpuids == "nobind":

95 - self.local_omp_cpuid = "nobind"

96 - else:

97 - local_dp_rank = self.parallel_config.data_parallel_rank_local

98 - omp_cpuids_list = omp_cpuids.split("|")

99 - if local_dp_rank is not None:

100 - world_size = self.parallel_config.world_size

101 - omp_cpuids_list = omp_cpuids_list[

102 - local_dp_rank * world_size : (local_dp_rank + 1) * world_size

103 - ]

104 - self.local_omp_cpuid = omp_cpuids_list[self.rank]

105 -

106 - if self.local_omp_cpuid != "nobind":

107 - ret = torch.ops._C.init_cpu_threads_env(self.local_omp_cpuid)

108 - if ret:

109 - logger.info(ret)

110 -

111 - # After the thread binding, changing thread num is not allowed

70 def skip_set_num_threads(x: int):

71 logger.warning(

72 "CPU backend doesn't allow to use "

@@ -153,92 +111,6 @@ class CPUWorker(Worker):

111 self.model_runner.warming_up_model()

112 return self.compilation_config.compilation_time

113

156 - def _get_autobind_cpu_ids(

157 - self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]

158 - ) -> str:

159 - """

160 - Return CPU ids to bind based on NUMA nodes.

161 - Currently for rank N, only CPU ids on the N-th node in available NUMA

162 - node list will be selected.

163 - Args:

164 - cpu_selector: a callable object to select CPUs from a CPU list

165 - of a physical core. The input is a LogicalCPUInfo list, sorted by

166 - the LogicalCPUInfo.id. A selected LogicalCPUInfo list should be

167 - returned.

168 - """

169 - # simulate multiple numa nodes, for testing

170 - sim_multi_numa_nodes = os.environ.get("VLLM_CPU_SIM_MULTI_NUMA", "0") != "0"

171 -

172 - allowed_numa_nodes, logical_cpu_list = (

173 - CpuPlatform.get_allowed_cpu_core_node_list()

174 - )

175 - local_world_size = self.parallel_config.local_world_size

176 - assert len(allowed_numa_nodes) >= local_world_size or sim_multi_numa_nodes, (

177 - f"Not enough allowed NUMA nodes to bind threads of "

178 - f"{local_world_size} local CPUWorkers. "

179 - f"Allowed NUMA nodes are {allowed_numa_nodes}. "

180 - "Please try to bind threads manually."

181 - )

182 -

183 - if not sim_multi_numa_nodes:

184 - # Get CPUs on NUMA node `allowed_numa_nodes[local_rank]`

185 - selected_numa_node = allowed_numa_nodes[self.local_rank] # type: ignore

186 - logical_cpu_list = [

187 - x for x in logical_cpu_list if x.numa_node == selected_numa_node

188 - ]

189 - else:

190 - # This is a bit tricky because the internal DP size

191 - # is always 1 for non-MoE models

192 - world_size_across_dp = (

193 - self.parallel_config.world_size

194 - * self.parallel_config._api_process_count

195 - )

196 - assert len(logical_cpu_list) >= world_size_across_dp

197 - logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.numa_node)

198 - sim_cpu_num_per_node = len(logical_cpu_list) // world_size_across_dp

199 - assert self.parallel_config.data_parallel_rank_local is not None

200 - start_idx = (

201 - self.local_rank

202 - + self.parallel_config.world_size

203 - * self.parallel_config.data_parallel_rank_local

204 - ) * sim_cpu_num_per_node

205 - logical_cpu_list = logical_cpu_list[

206 - start_idx : (start_idx + sim_cpu_num_per_node)

207 - ]

208 -

209 - # Select CPUs from each physical core via cpu_selector

210 - core_to_cpus: dict[int, list[LogicalCPUInfo]] = {}

211 - for cpu_info in logical_cpu_list:

212 - if cpu_info.physical_core not in core_to_cpus:

213 - core_to_cpus[cpu_info.physical_core] = []

214 - core_to_cpus[cpu_info.physical_core].append(cpu_info)

215 - logical_cpu_list = []

216 - for cpu_list in core_to_cpus.values():

217 - cpu_list = sorted(cpu_list, key=lambda x: x.id)

218 - logical_cpu_list.extend(cpu_selector(cpu_list))

219 - logical_cpu_list = sorted(logical_cpu_list, key=lambda x: x.id)

220 -

221 - # Reserve CPUs for other processes

222 - reserve_cpu_num = envs.VLLM_CPU_NUM_OF_RESERVED_CPU

223 - if reserve_cpu_num is None:

224 - need_reserve = (

225 - self.parallel_config.world_size > 1

226 - or self.parallel_config.data_parallel_size_local > 1

227 - )

228 - reserve_cpu_num = 1 if need_reserve else 0

229 - assert len(logical_cpu_list) > reserve_cpu_num, (

230 - f"VLLM_CPU_NUM_OF_RESERVED_CPU ({reserve_cpu_num}) "

231 - f"should less than {len(logical_cpu_list)}."

232 - )

233 - if reserve_cpu_num != 0:

234 - logical_cpu_list = logical_cpu_list[:-reserve_cpu_num]

235 -

236 - logger.info(

237 - "auto thread-binding list (id, physical core): %s",

238 - [(x.id, x.physical_core) for x in logical_cpu_list],

239 - )

240 - return ",".join([str(x.id) for x in logical_cpu_list])

241 -

114 def profile(self, is_start: bool = True, profile_prefix: str | None = None):

115 if self.profiler is None:

116 raise RuntimeError("Profiler is not enabled.")