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 12, 22:15 UTC
[cpu][ci] Add initial set of tests for Arm CPUs (#28657)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Nov 20, 2025 at 02:37 UTC
3168285fcaaee09bc93dce7bc9ae6ee823c71652
2 files changed
+74
.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh
new
+64
@@ -0,0 +1,64 @@
1
+#!/bin/bash
2
+
3
+# This script build the CPU docker image and run the offline inference inside the container.
4
+# It serves a sanity check for compilation and basic model usage.
5
+set -ex
6
+
7
+# allow to bind to different cores
8
+CORE_RANGE=${CORE_RANGE:-0-16}
9
+OMP_CORE_RANGE=${OMP_CORE_RANGE:-0-16}
10
+NUMA_NODE=${NUMA_NODE:-0}
11
+
12
+export CMAKE_BUILD_PARALLEL_LEVEL=32
13
+
14
+# Setup cleanup
15
+remove_docker_container() {
16
+ set -e;
17
+ docker rm -f cpu-test-"$NUMA_NODE" || true;
18
+}
19
+trap remove_docker_container EXIT
20
+remove_docker_container
21
+
22
+# Try building the docker image
23
+numactl -C "$CORE_RANGE" -N "$NUMA_NODE" docker build --tag cpu-test-"$NUMA_NODE" --target vllm-test -f docker/Dockerfile.cpu .
24
+
25
+# Run the image, setting --shm-size=4g for tensor parallel.
26
+docker run -itd --cpuset-cpus="$CORE_RANGE" --cpuset-mems="$NUMA_NODE" --entrypoint /bin/bash -v ~/.cache/huggingface:/root/.cache/huggingface --privileged=true -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=16 --env VLLM_CPU_CI_ENV=1 -e E2E_OMP_THREADS="$OMP_CORE_RANGE" --shm-size=4g --name cpu-test-"$NUMA_NODE" cpu-test-"$NUMA_NODE"
27
+
28
+function cpu_tests() {
29
+ set -e
30
+ export NUMA_NODE=$2
31
+
32
+ docker exec cpu-test-"$NUMA_NODE" bash -c "
33
+ set -e
34
+ pip list"
35
+
36
+ # offline inference
37
+ docker exec cpu-test-"$NUMA_NODE" bash -c "
38
+ set -e
39
+ python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m"
40
+
41
+ # Run kernel tests
42
+ docker exec cpu-test-"$NUMA_NODE" bash -c "
43
+ set -e
44
+ pytest -x -v -s tests/kernels/test_onednn.py
45
+ pytest -x -v -s tests/kernels/attention/test_cpu_attn.py"
46
+
47
+ # basic online serving
48
+ docker exec cpu-test-"$NUMA_NODE" bash -c '
49
+ set -e
50
+ VLLM_CPU_OMP_THREADS_BIND=$E2E_OMP_THREADS vllm serve meta-llama/Llama-3.2-3B-Instruct --max-model-len 2048 &
51
+ server_pid=$!
52
+ timeout 600 bash -c "until curl localhost:8000/v1/models; do sleep 1; done" || exit 1
53
+ vllm bench serve \
54
+ --backend vllm \
55
+ --dataset-name random \
56
+ --model meta-llama/Llama-3.2-3B-Instruct \
57
+ --num-prompts 20 \
58
+ --endpoint /v1/completions
59
+ kill -s SIGTERM $server_pid &'
60
+}
61
+
62
+# All of CPU tests are expected to be finished less than 40 mins.
63
+export -f cpu_tests
64
+timeout 2h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
docker/Dockerfile.cpu
+10
@@ -37,6 +37,7 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked \
37
&& update-alternatives --install /usr/bin/gcc gcc /usr/bin/gcc-12 10 --slave /usr/bin/g++ g++ /usr/bin/g++-12 \
38
&& curl -LsSf https://astral.sh/uv/install.sh | sh
39
40
+ENV CC=/usr/bin/gcc-12 CXX=/usr/bin/g++-12
41
ENV CCACHE_DIR=/root/.cache/ccache
42
ENV CMAKE_CXX_COMPILER_LAUNCHER=ccache
43
@@ -122,6 +123,15 @@ WORKDIR /workspace/vllm
123
RUN --mount=type=bind,src=requirements/test.in,target=requirements/test.in \
124
cp requirements/test.in requirements/cpu-test.in && \
125
sed -i '/mamba_ssm/d' requirements/cpu-test.in && \
126
+ remove_packages_not_supported_on_aarch64() { \
127
+ case "$(uname -m)" in \
128
+ aarch64|arm64) \
129
+ sed -i '/decord/d' requirements/cpu-test.in; \
130
+ sed -i '/terratorch/d' requirements/cpu-test.in; \
131
+ ;; \
132
+ esac; \
133
+ }; \
134
+ remove_packages_not_supported_on_aarch64 && \
135
sed -i 's/^torch==.*/torch==2.8.0/g' requirements/cpu-test.in && \
136
sed -i 's/torchaudio.*/torchaudio/g' requirements/cpu-test.in && \
137
sed -i 's/torchvision.*/torchvision/g' requirements/cpu-test.in && \