[CPU] Enable tcmalloc for s390x (#50841)
Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Rehan Khan committed
Aug 5, 2026 at 08:25 UTC
0187f4c88e7ef19cee9ff67b9822f143d142ea2a
4 files changed
+39
-2
docker/Dockerfile.s390x
+15
@@ -59,6 +59,15 @@ RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz
59
# Set include path
60
ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
61
62
+FROM base AS gperftools-build
63
+WORKDIR /tmp
64
+ARG GPERFTOOLS_VERSION=2.16
65
+RUN curl -LO https://github.com/gperftools/gperftools/releases/download/gperftools-${GPERFTOOLS_VERSION}/gperftools-${GPERFTOOLS_VERSION}.tar.gz && \
66
+ tar -xzf gperftools-${GPERFTOOLS_VERSION}.tar.gz && \
67
+ cd gperftools-${GPERFTOOLS_VERSION} && \
68
+ ./configure --enable-minimal && \
69
+ make -j"$(nproc)"
70
+
71
FROM python-install AS torch-vision
72
# Install torchvision
73
ARG TORCH_VISION_VERSION=v0.28.0
@@ -168,6 +177,12 @@ WORKDIR /workspace/vllm
177
RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
178
make -C /numactl install
179
180
+RUN --mount=type=bind,from=gperftools-build,src=/tmp/gperftools-2.16,target=/gperftools \
181
+ make -C /gperftools install-libLTLIBRARIES && \
182
+ /sbin/ldconfig
183
+
184
+ENV LD_PRELOAD=/usr/local/lib/libtcmalloc_minimal.so.4
185
+
186
# Install dependencies, including PyTorch and Apache Arrow
187
RUN --mount=type=cache,target=/root/.cache/uv \
188
--mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
docs/getting_started/installation/cpu.s390x.inc.md
+21
@@ -117,6 +117,27 @@ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
117
pip install dist/*.whl
118
```
119
120
+!!! warning "set `LD_PRELOAD` for TCMalloc"
121
+ For best memory allocation performance, build and install
122
+ [gperftools](https://github.com/gperftools/gperftools) (TCMalloc) from
123
+ source and add it to `LD_PRELOAD`:
124
+
125
+ ```bash
126
+ # Build and install TCMalloc from source
127
+ curl -LO https://github.com/gperftools/gperftools/releases/download/gperftools-2.16/gperftools-2.16.tar.gz
128
+ tar -xzf gperftools-2.16.tar.gz
129
+ cd gperftools-2.16
130
+ ./configure --enable-minimal && make -j$(nproc) && sudo make install
131
+ sudo ldconfig
132
+ cd ..
133
+
134
+ # Add to LD_PRELOAD
135
+ export LD_PRELOAD="/usr/local/lib/libtcmalloc_minimal.so.4:$LD_PRELOAD"
136
+ ```
137
+
138
+ The Docker image (`Dockerfile.s390x`) already includes TCMalloc and sets
139
+ `LD_PRELOAD` automatically.
140
+
141
!!! warning "Protobuf workaround for s390x"
142
The C++ protobuf extension crashes on s390x. After installation, set the
143
following environment variable and remove the C++ extensions:
setup.py
+1
-1
@@ -127,7 +127,7 @@ def should_bundle_tcmalloc() -> bool:
127
return (
128
VLLM_TARGET_DEVICE == "cpu"
129
and sys.platform.startswith("linux")
130
- and platform.machine() in ("aarch64", "x86_64")
130
+ and platform.machine() in ("aarch64", "x86_64", "s390x")
131
)
132
133
vllm/platforms/cpu.py
+2
-1
@@ -290,7 +290,8 @@ class CpuPlatform(Platform):
290
# memory allocation overhead
291
if (
292
platform.system() == "Linux"
293
- and cpu_architecture in (CpuArchEnum.ARM, CpuArchEnum.X86)
293
+ and cpu_architecture
294
+ in (CpuArchEnum.ARM, CpuArchEnum.X86, CpuArchEnum.S390X)
295
and "libtcmalloc" not in ld_preload_str
296
):
297
vllm_pkg = os.path.dirname(os.path.dirname(__file__))