@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 396ee94180
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 14, 23:45 UTC

[CI] Unifying Dockerfiles for ARM and X86 Builds (#21343)

Signed-off-by: Kebe <mail@kebe7jun.com>

Kebe committed Jul 25, 2025 at 22:33 UTC 396ee941803d0382603f567bba41116bb3d04dda
5 files changed +29 -67
.github/workflows/lint-and-deploy.yaml
+1 -1
@@ -7,7 +7,7 @@ permissions:
7
8 jobs:
9 lint-and-deploy:
10 - runs-on: ubuntu-latest
10 + runs-on: ubuntu-24.04-arm
11 steps:
12 - name: Checkout
13 uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
docker/Dockerfile.arm deleted
-62
@@ -1,62 +0,0 @@
1 -# This vLLM Dockerfile is used to construct an image that can build and run vLLM on ARM CPU platform.
2 -
3 -FROM ubuntu:22.04 AS cpu-test-arm
4 -
5 -ENV CCACHE_DIR=/root/.cache/ccache
6 -
7 -ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache
8 -
9 -RUN --mount=type=cache,target=/var/cache/apt \
10 - apt-get update -y \
11 - && apt-get install -y curl ccache git wget vim numactl gcc-12 g++-12 python3 python3-pip libtcmalloc-minimal4 libnuma-dev \
12 - && apt-get install -y ffmpeg libsm6 libxext6 libgl1 \
13 - && update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12
14 -
15 -# tcmalloc provides better memory allocation efficiency, e.g., holding memory in caches to speed up access of commonly-used objects.
16 -RUN --mount=type=cache,target=/root/.cache/pip \
17 - pip install py-cpuinfo # Use this to gather CPU info and optimize based on ARM Neoverse cores
18 -
19 -# Set LD_PRELOAD for tcmalloc on ARM
20 -ENV LD_PRELOAD="/usr/lib/aarch64-linux-gnu/libtcmalloc_minimal.so.4"
21 -
22 -RUN echo 'ulimit -c 0' >> ~/.bashrc
23 -
24 -WORKDIR /workspace
25 -
26 -ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/cpu"
27 -ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
28 -RUN --mount=type=cache,target=/root/.cache/pip \
29 - --mount=type=bind,src=requirements/build.txt,target=requirements/build.txt \
30 - pip install --upgrade pip && \
31 - pip install -r requirements/build.txt
32 -
33 -FROM cpu-test-arm AS build
34 -
35 -WORKDIR /workspace/vllm
36 -
37 -RUN --mount=type=cache,target=/root/.cache/pip \
38 - --mount=type=bind,src=requirements/common.txt,target=requirements/common.txt \
39 - --mount=type=bind,src=requirements/cpu.txt,target=requirements/cpu.txt \
40 - pip install -v -r requirements/cpu.txt
41 -
42 -COPY . .
43 -ARG GIT_REPO_CHECK=0
44 -RUN --mount=type=bind,source=.git,target=.git \
45 - if [ "$GIT_REPO_CHECK" != 0 ]; then bash tools/check_repo.sh ; fi
46 -
47 -# Disabling AVX512 specific optimizations for ARM
48 -ARG VLLM_CPU_DISABLE_AVX512="true"
49 -ENV VLLM_CPU_DISABLE_AVX512=${VLLM_CPU_DISABLE_AVX512}
50 -
51 -RUN --mount=type=cache,target=/root/.cache/pip \
52 - --mount=type=cache,target=/root/.cache/ccache \
53 - --mount=type=bind,source=.git,target=.git \
54 - VLLM_TARGET_DEVICE=cpu python3 setup.py bdist_wheel && \
55 - pip install dist/*.whl && \
56 - rm -rf dist
57 -
58 -WORKDIR /workspace/
59 -
60 -RUN ln -s /workspace/vllm/tests && ln -s /workspace/vllm/examples && ln -s /workspace/vllm/benchmarks
61 -
62 -ENTRYPOINT ["python3", "-m", "vllm.entrypoints.openai.api_server"]
\ No newline at end of file
docker/Dockerfile.cpu
+22 -2
@@ -1,4 +1,11 @@
1 -# This vLLM Dockerfile is used to construct image that can build and run vLLM on x86 CPU platform.
1 +# This vLLM Dockerfile is used to build images that can run vLLM on both x86_64 and arm64 CPU platforms.
2 +#
3 +# Supported platforms:
4 +# - linux/amd64 (x86_64)
5 +# - linux/arm64 (aarch64)
6 +#
7 +# Use the `--platform` option with `docker buildx build` to specify the target architecture, e.g.:
8 +# docker buildx build --platform=linux/arm64 -f docker/Dockerfile.cpu .
9 #
10 # Build targets:
11 # vllm-openai (default): used for serving deployment
@@ -53,7 +60,20 @@ RUN --mount=type=cache,target=/root/.cache/uv \
60 uv pip install --upgrade pip && \
61 uv pip install -r requirements/cpu.txt
62
56 -ENV LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4:/opt/venv/lib/libiomp5.so:$LD_PRELOAD"
63 +ARG TARGETARCH
64 +ENV TARGETARCH=${TARGETARCH}
65 +
66 +RUN if [ "$TARGETARCH" = "arm64" ]; then \
67 + PRELOAD_PATH="/usr/lib/aarch64-linux-gnu/libtcmalloc_minimal.so.4"; \
68 + else \
69 + PRELOAD_PATH="/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4:/opt/venv/lib/libiomp5.so"; \
70 + fi && \
71 + echo "export LD_PRELOAD=$PRELOAD_PATH" >> ~/.bashrc
72 +
73 +# Ensure that the LD_PRELOAD environment variable for export is in effect.
74 +SHELL ["/bin/bash", "-c"]
75 +
76 +ENV LD_PRELOAD=${LD_PRELOAD}
77
78 RUN echo 'ulimit -c 0' >> ~/.bashrc
79
docs/getting_started/installation/cpu/arm.inc.md
+1 -1
@@ -33,7 +33,7 @@ Testing has been conducted on AWS Graviton3 instances for compatibility.
33 # --8<-- [end:pre-built-images]
34 # --8<-- [start:build-image-from-source]
35 ```bash
36 -docker build -f docker/Dockerfile.arm \
36 +docker build -f docker/Dockerfile.cpu \
37 --tag vllm-cpu-env .
38
39 # Launching OpenAI server
requirements/cpu.txt
+5 -1
@@ -10,7 +10,8 @@ setuptools>=77.0.3,<80.0.0
10 --extra-index-url https://download.pytorch.org/whl/cpu
11 torch==2.6.0+cpu; platform_machine == "x86_64" # torch>2.6.0+cpu has performance regression on x86 platform, see https://github.com/pytorch/pytorch/pull/151218
12 torch==2.7.0; platform_system == "Darwin"
13 -torch==2.7.0; platform_machine == "ppc64le" or platform_machine == "aarch64"
13 +torch==2.7.0; platform_machine == "ppc64le"
14 +torch==2.6.0; platform_machine == "aarch64" # for arm64 CPUs, torch 2.7.0 has a issue: https://github.com/vllm-project/vllm/issues/17960
15
16 # required for the image processor of minicpm-o-2_6, this must be updated alongside torch
17 torchaudio; platform_machine != "ppc64le" and platform_machine != "s390x"
@@ -25,3 +26,6 @@ datasets # for benchmark scripts
26 intel-openmp==2024.2.1; platform_machine == "x86_64"
27 intel_extension_for_pytorch==2.6.0; platform_machine == "x86_64" # torch>2.6.0+cpu has performance regression on x86 platform, see https://github.com/pytorch/pytorch/pull/151218
28 triton==3.2.0; platform_machine == "x86_64" # Triton is required for torch 2.6+cpu, as it is imported in torch.compile.
29 +
30 +# Use this to gather CPU info and optimize based on ARM Neoverse cores
31 +py-cpuinfo; platform_machine == "aarch64"