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 2, 19:30 UTC
[ci] Add arm64 ci image (#41303)
Signed-off-by: khluu <khluu000@gmail.com> Signed-off-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Kevin H. Luu committed
May 26, 2026 at 14:38 UTC
e19b9b1045f6b8c2b5a647cc3a819bb919a718ac
4 files changed
+64
-3
.buildkite/image_build/image_build.yaml
+18
@@ -98,3 +98,21 @@ steps:
98
limit: 2
99
- exit_status: -10 # Agent was lost
100
limit: 2
101
+
102
+ - label: ":docker: Build arm64 image"
103
+ key: arm64-image-build
104
+ depends_on: []
105
+ source_file_dependencies:
106
+ - ".buildkite/image_build/image_build.yaml"
107
+ - ".buildkite/image_build/image_build_arm64.sh"
108
+ - "docker/Dockerfile"
109
+ commands:
110
+ - .buildkite/image_build/image_build_arm64.sh $REGISTRY $REPO $BUILDKITE_COMMIT
111
+ env:
112
+ DOCKER_BUILDKIT: "1"
113
+ retry:
114
+ automatic:
115
+ - exit_status: -1 # Agent was lost
116
+ limit: 2
117
+ - exit_status: -10 # Agent was lost
118
+ limit: 2
.buildkite/image_build/image_build_arm64.sh
new
+37
@@ -0,0 +1,37 @@
1
+#!/bin/bash
2
+set -e
3
+
4
+if [[ $# -lt 3 ]]; then
5
+ echo "Usage: $0 <registry> <repo> <commit>"
6
+ exit 1
7
+fi
8
+
9
+REGISTRY=$1
10
+REPO=$2
11
+BUILDKITE_COMMIT=$3
12
+
13
+# authenticate with AWS ECR
14
+aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$REGISTRY" || true
15
+
16
+# skip build if image already exists
17
+if [[ -z $(docker manifest inspect "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64) ]]; then
18
+ echo "Image not found, proceeding with build..."
19
+else
20
+ echo "Image found"
21
+ exit 0
22
+fi
23
+
24
+# build (Grace/GH200 is the arm64 GPU target; sm_90)
25
+docker build --file docker/Dockerfile \
26
+ --platform linux/arm64 \
27
+ --build-arg max_jobs=16 \
28
+ --build-arg nvcc_threads=4 \
29
+ --build-arg torch_cuda_arch_list="9.0" \
30
+ --build-arg USE_SCCACHE=1 \
31
+ --build-arg buildkite_commit="$BUILDKITE_COMMIT" \
32
+ --tag "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64 \
33
+ --target test \
34
+ --progress plain .
35
+
36
+# push
37
+docker push "$REGISTRY"/"$REPO":"$BUILDKITE_COMMIT"-arm64
docker/Dockerfile
+6
@@ -868,6 +868,7 @@ FROM vllm-base AS test
868
ADD . /vllm-workspace/
869
870
ARG PYTHON_VERSION
871
+ARG TARGETPLATFORM
872
873
ARG PIP_INDEX_URL UV_INDEX_URL
874
ARG PIP_EXTRA_INDEX_URL UV_EXTRA_INDEX_URL
@@ -907,6 +908,11 @@ RUN --mount=type=cache,target=/opt/uv/cache \
908
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/nightly/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
909
else \
910
echo "Installing dev requirements..." \
911
+ && if [ "$TARGETPLATFORM" = "linux/arm64" ]; then \
912
+ echo "Recompiling test requirements for arm64..." \
913
+ && uv pip compile requirements/test/cuda.in -o requirements/test/cuda.txt --index-strategy unsafe-best-match \
914
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
915
+ fi \
916
&& uv pip install --system -r requirements/dev.txt \
917
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.'); \
918
fi \
requirements/test/cuda.in
+3
-3
@@ -53,12 +53,12 @@ tritonclient>=2.51.0
53
grpcio==1.78.0
54
grpcio-reflection==1.78.0
55
56
-arctic-inference == 0.1.1 # Required for suffix decoding test
56
+arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test
57
numba == 0.65.0 # Required for N-gram speculative decoding
58
numpy
59
runai-model-streamer[s3,gcs,azure]==0.15.7
60
-fastsafetensors>=0.2.2 # 0.2.2 contains important fixes for multi-GPU mem usage
61
-instanttensor>=0.1.5
60
+fastsafetensors>=0.2.2; platform_machine == "x86_64" # 0.2.2 contains important fixes for multi-GPU mem usage
61
+instanttensor>=0.1.5; platform_machine == "x86_64"
62
pydantic>=2.12 # 2.11 leads to error on python 3.13
63
decord==0.6.0; platform_machine == "x86_64"
64
# terratorch is temporarily disabled while PyPI has the `lightning` package