Skip to content

Commit 808a876

Browse files
huydhnminpeter
authored andcommitted
Use prebuilt FlashInfer x86_64 PyTorch 2.7 CUDA 12.8 wheel for CI (vllm-project#18537)
Signed-off-by: Huy Do <huydhn@gmail.com> Signed-off-by: minpeter <kali2005611@gmail.com>
1 parent 1130584 commit 808a876

File tree

1 file changed

+8
-9
lines changed

1 file changed

+8
-9
lines changed

docker/Dockerfile

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -257,18 +257,17 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
257257
RUN --mount=type=cache,target=/root/.cache/uv \
258258
. /etc/environment && \
259259
if [ "$TARGETPLATFORM" != "linux/arm64" ]; then \
260-
# uv pip install --system https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.4/flashinfer_python-0.2.4+cu124torch2.6-cp38-abi3-linux_x86_64.whl ; \
261-
# TESTING: install FlashInfer from source to test 2.7.0 final RC
260+
# FlashInfer alreary has a wheel for PyTorch 2.7.0 and CUDA 12.8. This is enough for CI use
262261
if [[ "$CUDA_VERSION" == 12.8* ]]; then \
263-
export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0 10.0+PTX'; \
262+
uv pip install --system https://download.pytorch.org/whl/cu128/flashinfer/flashinfer_python-0.2.5%2Bcu128torch2.7-cp38-abi3-linux_x86_64.whl; \
264263
else \
265264
export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0+PTX'; \
266-
fi; \
267-
CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
268-
if [ "$CUDA_MAJOR" -lt 12 ]; then \
269-
export FLASHINFER_ENABLE_SM90=0; \
270-
fi; \
271-
uv pip install --system --no-build-isolation "git+https://github.com/flashinfer-ai/flashinfer@21ea1d2545f74782b91eb8c08fd503ac4c0743fc" ; \
265+
CUDA_MAJOR="${CUDA_VERSION%%.*}"; \
266+
if [ "$CUDA_MAJOR" -lt 12 ]; then \
267+
export FLASHINFER_ENABLE_SM90=0; \
268+
fi; \
269+
uv pip install --system --no-build-isolation "git+https://github.com/flashinfer-ai/flashinfer@21ea1d2545f74782b91eb8c08fd503ac4c0743fc" ; \
270+
fi \
272271
fi
273272
COPY examples examples
274273
COPY benchmarks benchmarks

0 commit comments

Comments
 (0)