[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

@@ -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 \

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:

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")
     )
 
 

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__))