This repository was archived by the owner on Oct 11, 2024. It is now read-only.
File tree Expand file tree Collapse file tree 3 files changed +20
-2
lines changed Expand file tree Collapse file tree 3 files changed +20
-2
lines changed Original file line number Diff line number Diff line change @@ -4,17 +4,20 @@ set -ex
4
4
5
5
# Try building the docker image
6
6
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 .
7
8
8
9
# 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 ; }
10
11
trap remove_docker_container EXIT
11
12
remove_docker_container
12
13
13
14
# Run the image
14
15
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
15
17
16
18
# offline inference
17
19
docker exec cpu-test bash -c " python3 examples/offline_inference.py"
20
+ docker exec cpu-test-avx2 bash -c " python3 examples/offline_inference.py"
18
21
19
22
# Run basic model test
20
23
docker exec cpu-test bash -c " cd tests;
Original file line number Diff line number Diff line change @@ -21,6 +21,10 @@ WORKDIR /workspace/vllm
21
21
22
22
RUN pip install -v -r requirements-cpu.txt --extra-index-url https://download.pytorch.org/whl/cpu
23
23
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
+
24
28
RUN VLLM_TARGET_DEVICE=cpu python3 setup.py install
25
29
26
30
WORKDIR /workspace/
Original file line number Diff line number Diff line change @@ -33,10 +33,21 @@ function (find_isa CPUINFO TARGET OUT)
33
33
endif ()
34
34
endfunction ()
35
35
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
+
36
47
find_isa(${CPUINFO} "avx2" AVX2_FOUND)
37
48
find_isa(${CPUINFO} "avx512f" AVX512_FOUND)
38
49
39
- if (AVX512_FOUND)
50
+ if (AVX512_FOUND AND NOT AVX512_DISABLED )
40
51
list (APPEND CXX_COMPILE_FLAGS
41
52
"-mavx512f"
42
53
"-mavx512vl"
You can’t perform that action at this time.
0 commit comments