From 0187f4c88e7ef19cee9ff67b9822f143d142ea2a Mon Sep 17 00:00:00 2001
From: Rehan Khan <Rehan.Khan7@ibm.com>
Date: Wed, 5 Aug 2026 08:25:39 +0530
Subject: [PATCH] [CPU] Enable tcmalloc for s390x (#50841)

Signed-off-by: Rehan Khan <Rehan.Khan7@ibm.com>
Co-authored-by: Li, Jiang <jiang1.li@intel.com>
---
 docker/Dockerfile.s390x                       | 15 +++++++++++++
 .../installation/cpu.s390x.inc.md             | 21 +++++++++++++++++++
 setup.py                                      |  2 +-
 vllm/platforms/cpu.py                         |  3 ++-
 4 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/docker/Dockerfile.s390x b/docker/Dockerfile.s390x
index d645e37e78..9517ff3718 100644
--- a/docker/Dockerfile.s390x
+++ b/docker/Dockerfile.s390x
@@ -59,6 +59,15 @@ RUN curl -LO https://github.com/numactl/numactl/archive/refs/tags/v2.0.19.tar.gz
 # Set include path
 ENV C_INCLUDE_PATH="/usr/local/include:$C_INCLUDE_PATH"
 
+FROM base AS gperftools-build
+WORKDIR /tmp
+ARG GPERFTOOLS_VERSION=2.16
+RUN curl -LO https://github.com/gperftools/gperftools/releases/download/gperftools-${GPERFTOOLS_VERSION}/gperftools-${GPERFTOOLS_VERSION}.tar.gz && \
+    tar -xzf gperftools-${GPERFTOOLS_VERSION}.tar.gz && \
+    cd gperftools-${GPERFTOOLS_VERSION} && \
+    ./configure --enable-minimal && \
+    make -j"$(nproc)"
+
 FROM python-install AS torch-vision
 # Install torchvision
 ARG TORCH_VISION_VERSION=v0.28.0
@@ -168,6 +177,12 @@ WORKDIR /workspace/vllm
 RUN --mount=type=bind,from=numa-build,src=/tmp/numactl-2.0.19,target=/numactl \
     make -C /numactl install
 
+RUN --mount=type=bind,from=gperftools-build,src=/tmp/gperftools-2.16,target=/gperftools \
+    make -C /gperftools install-libLTLIBRARIES && \
+    /sbin/ldconfig
+
+ENV LD_PRELOAD=/usr/local/lib/libtcmalloc_minimal.so.4
+
 # Install dependencies, including PyTorch and Apache Arrow
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,from=rust,source=/root/.cargo,target=/root/.cargo,rw \
diff --git a/docs/getting_started/installation/cpu.s390x.inc.md b/docs/getting_started/installation/cpu.s390x.inc.md
index 0c07c255e2..16fa13525a 100644
--- a/docs/getting_started/installation/cpu.s390x.inc.md
+++ b/docs/getting_started/installation/cpu.s390x.inc.md
@@ -117,6 +117,27 @@ VLLM_TARGET_DEVICE=cpu VLLM_CPU_MOE_PREPACK=0 python setup.py bdist_wheel && \
         pip install dist/*.whl
     ```
 
+!!! warning "set `LD_PRELOAD` for TCMalloc"
+    For best memory allocation performance, build and install
+    [gperftools](https://github.com/gperftools/gperftools) (TCMalloc) from
+    source and add it to `LD_PRELOAD`:
+
+    ```bash
+    # Build and install TCMalloc from source
+    curl -LO https://github.com/gperftools/gperftools/releases/download/gperftools-2.16/gperftools-2.16.tar.gz
+    tar -xzf gperftools-2.16.tar.gz
+    cd gperftools-2.16
+    ./configure --enable-minimal && make -j$(nproc) && sudo make install
+    sudo ldconfig
+    cd ..
+
+    # Add to LD_PRELOAD
+    export LD_PRELOAD="/usr/local/lib/libtcmalloc_minimal.so.4:$LD_PRELOAD"
+    ```
+
+    The Docker image (`Dockerfile.s390x`) already includes TCMalloc and sets
+    `LD_PRELOAD` automatically.
+
 !!! warning "Protobuf workaround for s390x"
     The C++ protobuf extension crashes on s390x. After installation, set the
     following environment variable and remove the C++ extensions:
diff --git a/setup.py b/setup.py
index a8817a9608..587f66dfbe 100644
--- a/setup.py
+++ b/setup.py
@@ -127,7 +127,7 @@ def should_bundle_tcmalloc() -> bool:
     return (
         VLLM_TARGET_DEVICE == "cpu"
         and sys.platform.startswith("linux")
-        and platform.machine() in ("aarch64", "x86_64")
+        and platform.machine() in ("aarch64", "x86_64", "s390x")
     )
 
 
diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py
index 698e83a8ab..837038f7ae 100644
--- a/vllm/platforms/cpu.py
+++ b/vllm/platforms/cpu.py
@@ -290,7 +290,8 @@ class CpuPlatform(Platform):
         # memory allocation overhead
         if (
             platform.system() == "Linux"
-            and cpu_architecture in (CpuArchEnum.ARM, CpuArchEnum.X86)
+            and cpu_architecture
+            in (CpuArchEnum.ARM, CpuArchEnum.X86, CpuArchEnum.S390X)
             and "libtcmalloc" not in ld_preload_str
         ):
             vllm_pkg = os.path.dirname(os.path.dirname(__file__))