Skip to content
This repository was archived by the owner on Oct 11, 2024. It is now read-only.

Commit e830048

Browse files
DamonFoolRobert Shaw
authored andcommitted
[CI/BUILD] Support non-AVX512 vLLM building and testing (vllm-project#5574)
1 parent dceff94 commit e830048

File tree

3 files changed

+20
-2
lines changed

3 files changed

+20
-2
lines changed

.buildkite/run-cpu-test.sh

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,20 @@ set -ex
44

55
# Try building the docker image
66
docker build -t cpu-test -f Dockerfile.cpu .
7+
docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" -t cpu-test-avx2 -f Dockerfile.cpu .
78

89
# Setup cleanup
9-
remove_docker_container() { docker rm -f cpu-test || true; }
10+
remove_docker_container() { docker rm -f cpu-test cpu-test-avx2 || true; }
1011
trap remove_docker_container EXIT
1112
remove_docker_container
1213

1314
# Run the image
1415
docker run -itd -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus=48-95 --cpuset-mems=1 --network host -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --name cpu-test cpu-test
16+
docker run -itd -v ~/.cache/huggingface:/root/.cache/huggingface --cpuset-cpus=48-95 --cpuset-mems=1 --network host -e HF_TOKEN --env VLLM_CPU_KVCACHE_SPACE=4 --name cpu-test-avx2 cpu-test-avx2
1517

1618
# offline inference
1719
docker exec cpu-test bash -c "python3 examples/offline_inference.py"
20+
docker exec cpu-test-avx2 bash -c "python3 examples/offline_inference.py"
1821

1922
# Run basic model test
2023
docker exec cpu-test bash -c "cd tests;

Dockerfile.cpu

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,10 @@ WORKDIR /workspace/vllm
2121

2222
RUN pip install -v -r requirements-cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu
2323

24+
# Support for building with non-AVX512 vLLM: docker build --build-arg VLLM_CPU_DISABLE_AVX512="true" ...
25+
ARG VLLM_CPU_DISABLE_AVX512
26+
ENV VLLM_CPU_DISABLE_AVX512=${VLLM_CPU_DISABLE_AVX512}
27+
2428
RUN VLLM_TARGET_DEVICE=cpu python3 setup.py install
2529

2630
WORKDIR /workspace/

cmake/cpu_extension.cmake

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,21 @@ function (find_isa CPUINFO TARGET OUT)
3333
endif()
3434
endfunction()
3535

36+
function (is_avx512_disabled OUT)
37+
set(DISABLE_AVX512 $ENV{VLLM_CPU_DISABLE_AVX512})
38+
if(DISABLE_AVX512 AND DISABLE_AVX512 STREQUAL "true")
39+
set(${OUT} ON PARENT_SCOPE)
40+
else()
41+
set(${OUT} OFF PARENT_SCOPE)
42+
endif()
43+
endfunction()
44+
45+
is_avx512_disabled(AVX512_DISABLED)
46+
3647
find_isa(${CPUINFO} "avx2" AVX2_FOUND)
3748
find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
3849

39-
if (AVX512_FOUND)
50+
if (AVX512_FOUND AND NOT AVX512_DISABLED)
4051
list(APPEND CXX_COMPILE_FLAGS
4152
"-mavx512f"
4253
"-mavx512vl"

0 commit comments

Comments
 (0)