[CPU] Create Proper Numa topology for s390x (#40714)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Rehan Khan committed
Jul 13, 2026 at 10:28 UTC
9e57de7197f234f9d9187715d96e07e007048c0f
7 files changed
+84
-105
docker/Dockerfile.s390x
+4
-97
@@ -37,37 +37,6 @@ RUN microdnf install -y \
37
python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip python${PYTHON_VERSION}-wheel && \
38
python${PYTHON_VERSION} -m venv $VIRTUAL_ENV && pip install --no-cache -U pip wheel uv && microdnf clean all
39
40
-FROM python-install AS pyarrow
41
-
42
-# Build Apache Arrow
43
-WORKDIR /tmp
44
-RUN --mount=type=cache,target=/root/.cache/uv \
45
- git clone https://github.com/apache/arrow.git -b maint-19.0.1 && \
46
- cd arrow/cpp && \
47
- mkdir release && cd release && \
48
- cmake -DCMAKE_BUILD_TYPE=Release \
49
- -DCMAKE_INSTALL_PREFIX=/usr/local \
50
- -DARROW_PYTHON=ON \
51
- -DARROW_PARQUET=ON \
52
- -DARROW_ORC=ON \
53
- -DARROW_FILESYSTEM=ON \
54
- -DARROW_WITH_LZ4=ON \
55
- -DARROW_WITH_ZSTD=ON \
56
- -DARROW_WITH_SNAPPY=ON \
57
- -DARROW_JSON=ON \
58
- -DARROW_CSV=ON \
59
- -DARROW_DATASET=ON \
60
- -DPROTOBUF_PROTOC_EXECUTABLE=/usr/bin/protoc \
61
- -DARROW_DEPENDENCY_SOURCE=BUNDLED \
62
- .. && \
63
- make -j$(nproc) && \
64
- make install && \
65
- cd ../../python && \
66
- export PYARROW_PARALLEL=4 && \
67
- export ARROW_BUILD_TYPE=release && \
68
- uv pip install -r requirements-build.txt && \
69
- python setup.py build_ext --build-type=$ARROW_BUILD_TYPE --bundle-arrow-cpp bdist_wheel
70
-
40
FROM python-install AS rust
41
ENV CARGO_HOME=/root/.cargo
42
ENV RUSTUP_HOME=/root/.rustup
@@ -121,60 +90,18 @@ RUN --mount=type=cache,target=/root/.cache/uv \
90
FROM python-install AS numba-builder
91
92
ARG MAX_JOBS
124
-ARG NUMBA_VERSION=0.61.2
93
+ARG NUMBA_VERSION=0.65.0
94
95
WORKDIR /tmp
96
97
# Clone all required dependencies
98
RUN --mount=type=cache,target=/root/.cache/uv \
99
microdnf install ninja-build gcc gcc-c++ -y && \
131
- git clone --recursive https://github.com/llvm/llvm-project.git -b llvmorg-15.0.7 && \
132
- git clone --recursive https://github.com/numba/llvmlite.git -b v0.44.0 && \
100
+ git clone --recursive https://github.com/numba/llvmlite.git -b v0.47.0 && \
101
git clone --recursive https://github.com/numba/numba.git -b ${NUMBA_VERSION} && \
134
- cd llvm-project && mkdir build && cd build && \
102
+ cd llvmlite && \
103
uv pip install 'cmake<4' 'setuptools<70' numpy && \
136
- export PREFIX=/usr/local && CMAKE_ARGS="${CMAKE_ARGS} -DLLVM_ENABLE_PROJECTS=lld;libunwind;compiler-rt" \
137
- CFLAGS="$(echo $CFLAGS | sed 's/-fno-plt //g')" \
138
- CXXFLAGS="$(echo $CXXFLAGS | sed 's/-fno-plt //g')" \
139
- CMAKE_ARGS="${CMAKE_ARGS} -DFFI_INCLUDE_DIR=$PREFIX/include" \
140
- CMAKE_ARGS="${CMAKE_ARGS} -DFFI_LIBRARY_DIR=$PREFIX/lib" \
141
- cmake -DCMAKE_INSTALL_PREFIX="${PREFIX}" \
142
- -DCMAKE_BUILD_TYPE=Release \
143
- -DCMAKE_LIBRARY_PATH="${PREFIX}" \
144
- -DLLVM_ENABLE_LIBEDIT=OFF \
145
- -DLLVM_ENABLE_LIBXML2=OFF \
146
- -DLLVM_ENABLE_RTTI=ON \
147
- -DLLVM_ENABLE_TERMINFO=OFF \
148
- -DLLVM_INCLUDE_BENCHMARKS=OFF \
149
- -DLLVM_INCLUDE_DOCS=OFF \
150
- -DLLVM_INCLUDE_EXAMPLES=OFF \
151
- -DLLVM_INCLUDE_GO_TESTS=OFF \
152
- -DLLVM_INCLUDE_TESTS=OFF \
153
- -DLLVM_INCLUDE_UTILS=ON \
154
- -DLLVM_INSTALL_UTILS=ON \
155
- -DLLVM_UTILS_INSTALL_DIR=libexec/llvm \
156
- -DLLVM_BUILD_LLVM_DYLIB=OFF \
157
- -DLLVM_LINK_LLVM_DYLIB=OFF \
158
- -DLLVM_EXPERIMENTAL_TARGETS_TO_BUILD=WebAssembly \
159
- -DLLVM_ENABLE_FFI=ON \
160
- -DLLVM_ENABLE_Z3_SOLVER=OFF \
161
- -DLLVM_OPTIMIZED_TABLEGEN=ON \
162
- -DCMAKE_POLICY_DEFAULT_CMP0111=NEW \
163
- -DCOMPILER_RT_BUILD_BUILTINS=ON \
164
- -DCOMPILER_RT_BUILTINS_HIDE_SYMBOLS=OFF \
165
- -DCOMPILER_RT_BUILD_LIBFUZZER=OFF \
166
- -DCOMPILER_RT_BUILD_CRT=OFF \
167
- -DCOMPILER_RT_BUILD_MEMPROF=OFF \
168
- -DCOMPILER_RT_BUILD_PROFILE=OFF \
169
- -DCOMPILER_RT_BUILD_SANITIZERS=OFF \
170
- -DCOMPILER_RT_BUILD_XRAY=OFF \
171
- -DCOMPILER_RT_BUILD_GWP_ASAN=OFF \
172
- -DCOMPILER_RT_BUILD_ORC=OFF \
173
- -DCOMPILER_RT_INCLUDE_TESTS=OFF \
174
- ${CMAKE_ARGS} -GNinja ../llvm \
175
- && ninja install . && \
176
- # build llvmlite
177
- cd ../../llvmlite && python setup.py bdist_wheel && \
104
+ python setup.py bdist_wheel && \
105
cd ../numba && \
106
if ! grep '#include "dynamic_annotations.h"' numba/_dispatcher.cpp; then \
107
sed -i '/#include "internal\/pycore_atomic.h"/i\#include "dynamic_annotations.h"' numba/_dispatcher.cpp; \
@@ -192,21 +119,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
119
cd opencv-python && \
120
python -m build --wheel --installer=uv --outdir /tmp/opencv-python/dist
121
195
-## Todo(r3hankhan123): Remove guidance-builder stage once vLLM upgrades to new version of llguidance that fixes s390x issues. See https://github.com/guidance-ai/llguidance/issues/330
196
-FROM python-install AS guidance-builder
197
-WORKDIR /tmp
198
-ENV CARGO_HOME=/root/.cargo
199
-ENV RUSTUP_HOME=/root/.rustup
200
-ENV PATH="$CARGO_HOME/bin:$RUSTUP_HOME/bin:$PATH"
201
-RUN --mount=type=cache,target=/root/.cache/uv \
202
- --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
203
- --mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
204
- git clone https://github.com/guidance-ai/llguidance.git && \
205
- cd llguidance && \
206
- git checkout s390x-fix-v2 && \
207
- uv pip install maturin && \
208
- python -m maturin build --release --out dist --compatibility linux
209
-
122
# # Final build stage
123
FROM python-install AS vllm-cpu
124
ARG PYTHON_VERSION
@@ -235,20 +147,16 @@ RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
147
RUN --mount=type=cache,target=/root/.cache/uv \
148
--mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
149
--mount=type=bind,from=rust,source=/root/.rustup,target=/root/.rustup,rw \
238
- --mount=type=bind,from=pyarrow,source=/tmp/arrow/python/dist,target=/tmp/arrow-wheels \
150
--mount=type=bind,from=torch-vision,source=/tmp/vision/dist,target=/tmp/vision-wheels/ \
151
--mount=type=bind,from=hf-xet-builder,source=/tmp/hf-xet/dist,target=/tmp/hf-xet-wheels/ \
152
--mount=type=bind,from=numba-builder,source=/tmp/llvmlite/dist,target=/tmp/llvmlite-wheels/ \
153
--mount=type=bind,from=numba-builder,source=/tmp/numba/dist,target=/tmp/numba-wheels/ \
154
--mount=type=bind,from=opencv-builder,source=/tmp/opencv-python/dist,target=/tmp/opencv-wheels/ \
244
- --mount=type=bind,from=guidance-builder,source=/tmp/llguidance/dist,target=/tmp/guidance-wheels/ \
245
- ARROW_WHL_FILE=$(ls /tmp/arrow-wheels/*.whl) && \
155
VISION_WHL_FILE=$(ls /tmp/vision-wheels/*.whl) && \
156
HF_XET_WHL_FILE=$(ls /tmp/hf-xet-wheels/*.whl) && \
157
LLVM_WHL_FILE=$(ls /tmp/llvmlite-wheels/*.whl) && \
158
NUMBA_WHL_FILE=$(ls /tmp/numba-wheels/*.whl) && \
159
OPENCV_WHL_FILE=$(ls /tmp/opencv-wheels/*.whl) && \
251
- GUIDANCE_WHL_FILE=$(ls /tmp/guidance-wheels/*.whl) && \
160
uv pip install -v \
161
$ARROW_WHL_FILE \
162
$VISION_WHL_FILE \
@@ -256,7 +164,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
164
$LLVM_WHL_FILE \
165
$NUMBA_WHL_FILE \
166
$OPENCV_WHL_FILE \
259
- $GUIDANCE_WHL_FILE \
167
--torch-backend cpu \
168
--index-strategy unsafe-best-match \
169
-r requirements/build/cpu.txt \
requirements/common.txt
+1
-1
@@ -21,7 +21,7 @@ pillow # Required for image processing
21
prometheus-fastapi-instrumentator >= 8.0.0 # v8 unblocks starlette >= 1.0
22
tiktoken >= 0.6.0 # Required for DBRX tokenizer
23
lm-format-enforcer == 0.11.3
24
-llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le"
24
+llguidance >= 1.7.0, < 1.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64" or platform_machine == "ppc64le" or platform_machine == "s390x"
25
outlines_core == 0.2.14
26
# required for outlines backend disk cache
27
diskcache == 5.6.3
vllm/distributed/device_communicators/cpu_communicator.py
+1
@@ -33,6 +33,7 @@ class CpuCommunicator(DeviceCommunicatorBase):
33
current_platform.get_cpu_architecture() == CpuArchEnum.X86
34
or current_platform.get_cpu_architecture() == CpuArchEnum.ARM
35
or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
36
+ or current_platform.get_cpu_architecture() == CpuArchEnum.S390X
37
)
38
and hasattr(torch.ops._C, "init_shm_manager")
39
and (unique_name.startswith("tp") or unique_name.startswith("pp"))
vllm/platforms/cpu.py
+6
-1
@@ -247,7 +247,12 @@ class CpuPlatform(Platform):
247
if (
248
platform.system() == "Linux"
249
and cpu_architecture
250
- in (CpuArchEnum.ARM, CpuArchEnum.POWERPC, CpuArchEnum.X86)
250
+ in (
251
+ CpuArchEnum.ARM,
252
+ CpuArchEnum.POWERPC,
253
+ CpuArchEnum.X86,
254
+ CpuArchEnum.S390X,
255
+ )
256
and not (
257
"libomp" in ld_preload_str
258
or "libgomp" in ld_preload_str
vllm/utils/cpu_resource_utils.py
+49
-4
@@ -3,6 +3,7 @@
3
4
import json
5
import os
6
+import platform
7
import subprocess
8
import sys
9
from dataclasses import dataclass
@@ -216,9 +217,14 @@ def _get_cpu_list() -> list[LogicalCPUInfo]:
217
# For MacOS, no user-level CPU affinity and SMT, return all CPUs
218
return _synthesize_cpu_list()
219
219
- lscpu_output = subprocess.check_output(
220
- "lscpu --json --extended=CPU,CORE,NODE --online", shell=True, text=True
221
- )
220
+ if platform.machine() == "s390x":
221
+ lscpu_output = subprocess.check_output(
222
+ "lscpu -J -e=CPU,CORE,NODE,SOCKET,BOOK", shell=True, text=True
223
+ )
224
+ else:
225
+ lscpu_output = subprocess.check_output(
226
+ "lscpu --json --extended=CPU,CORE,NODE --online", shell=True, text=True
227
+ )
228
229
# For platforms without NUMA, map bare `-` node to 0 so non-NUMA
230
# systems keep the existing behavior from #39781.
@@ -233,11 +239,50 @@ def _get_cpu_list() -> list[LogicalCPUInfo]:
239
lscpu_output,
240
)
241
242
+ if platform.machine() == "s390x":
243
+ # On s390x, use the best topology level as the abstract CPU group
244
+ # key via numa_node. The hierarchy is book > socket > core; prefer
245
+ # book, fall back to socket if only one book exists.
246
+ lscpu_output = re.sub(r'"book":\s*-\s*(,|\n|\})', r'"book": 0\1', lscpu_output)
247
+ lscpu_output = re.sub(
248
+ r'"socket":\s*-\s*(,|\n|\})', r'"socket": 0\1', lscpu_output
249
+ )
250
+
251
+ raw_cpus = json.loads(lscpu_output)["cpus"]
252
+ book_values = set()
253
+ socket_values = set()
254
+ for entry in raw_cpus:
255
+ book_values.add(LogicalCPUInfo._int(str(entry.get("book", "-1"))))
256
+ socket_values.add(LogicalCPUInfo._int(str(entry.get("socket", "-1"))))
257
+
258
+ # Pick the best partitioning level: book first, then socket
259
+ if len(book_values - {-1}) > 1:
260
+ group_key = "book"
261
+ elif len(socket_values - {-1}) > 1:
262
+ group_key = "socket"
263
+ else:
264
+ group_key = None
265
+
266
+ if group_key is not None:
267
+ logical_book_socket_list: list[LogicalCPUInfo] = []
268
+ for entry in raw_cpus:
269
+ cpu_id = LogicalCPUInfo._int(str(entry.get("cpu", "-1")))
270
+ core = LogicalCPUInfo._int(str(entry.get("core", "-1")))
271
+ group = LogicalCPUInfo._int(str(entry.get(group_key, "-1")))
272
+ if -1 not in (cpu_id, core, group):
273
+ logical_book_socket_list.append(
274
+ LogicalCPUInfo(id=cpu_id, physical_core=core, numa_node=group)
275
+ )
276
+
277
+ if not logical_book_socket_list:
278
+ return _synthesize_cpu_list()
279
+ return logical_book_socket_list
280
+
281
logical_cpu_list: list[LogicalCPUInfo] = json.loads(
282
lscpu_output, object_hook=LogicalCPUInfo.json_decoder
283
)["cpus"]
284
240
- # Filter CPUs with invalid attributes
285
+ # Filter CPUs with invalid attributes (only require id, core, node)
286
logical_cpu_list = [
287
x for x in logical_cpu_list if -1 not in (x.id, x.physical_core, x.numa_node)
288
]
vllm/utils/ompmultiprocessing.py
+8
@@ -204,6 +204,14 @@ class OMPProcessManager:
204
logical_cpu_list = cr_utils.get_allowed_cpu_list()
205
206
local_world_size = self.local_world_size
207
+
208
+ # On s390x, numa_node is remapped to book IDs (the best CPU
209
+ # partitioning level). Derive topology domains from the CPU list
210
+ # since they no longer correspond to physical memory nodes.
211
+ cpu_arch = current_platform.get_cpu_architecture()
212
+ if cpu_arch == CpuArchEnum.S390X:
213
+ allowed_numa_nodes = sorted(set(cpu.numa_node for cpu in logical_cpu_list))
214
+
215
assert (
216
len(allowed_numa_nodes) >= local_world_size or self.simulate_multi_node
217
), (
vllm/v1/worker/cpu_worker.py
+15
-2
@@ -55,9 +55,22 @@ class CPUWorker(Worker):
55
allowed_memory_nodes,
56
)
57
58
- torch.ops._C.init_cpu_memory_env([cpu_core.numa_node])
58
+ # On s390x, numa_node may be a synthetic book ID that doesn't
59
+ # correspond to a real memory node. Fall back to first visible node.
60
+ if cpu_core.numa_node in allowed_memory_nodes:
61
+ memory_node = cpu_core.numa_node
62
+ else:
63
+ logger.warning(
64
+ "CPU group key %s is not a valid memory node. "
65
+ "Falling back to memory node %s.",
66
+ cpu_core.numa_node,
67
+ allowed_memory_nodes[0],
68
+ )
69
+ memory_node = allowed_memory_nodes[0]
70
60
- memory_status = get_memory_node_info(cpu_core.numa_node)
71
+ torch.ops._C.init_cpu_memory_env([memory_node])
72
+
73
+ memory_status = get_memory_node_info(memory_node)
74
memory_fraction = vllm_config.cache_config.gpu_memory_utilization
75
self.requested_cpu_memory = math.ceil(
76
memory_status.total_memory * memory_fraction