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 15, 02:15 UTC
Update Dockerfile FlashInfer to v0.2.8rc1 (#20718)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 11, 2025 at 00:09 UTC
4b9a9435bb6bb03c91ce097ae1e3f85b77e1c161
1 file changed
+5
-2
docker/Dockerfile
+5
-2
@@ -379,12 +379,15 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
379
ARG FLASHINFER_CUDA128_INDEX_URL="https://download.pytorch.org/whl/cu128/flashinfer"
380
ARG FLASHINFER_CUDA128_WHEEL="flashinfer_python-0.2.6.post1%2Bcu128torch2.7-cp39-abi3-linux_x86_64.whl"
381
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
382
-ARG FLASHINFER_GIT_REF="v0.2.6.post1"
382
+ARG FLASHINFER_GIT_REF="v0.2.8rc1"
383
+# Flag to control whether to use pre-built FlashInfer wheels (set to false to force build from source)
384
+# TODO: Currently disabled because the pre-built wheels are not available for FLASHINFER_GIT_REF
385
+ARG USE_FLASHINFER_PREBUILT_WHEEL=false
386
RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
387
. /etc/environment
388
if [ "$TARGETPLATFORM" != "linux/arm64" ]; then
389
# FlashInfer already has a wheel for PyTorch 2.7.0 and CUDA 12.8. This is enough for CI use
387
- if [[ "$CUDA_VERSION" == 12.8* ]]; then
390
+ if [[ "$CUDA_VERSION" == 12.8* ]] && [[ "$USE_FLASHINFER_PREBUILT_WHEEL" == "true" ]]; then
391
uv pip install --system ${FLASHINFER_CUDA128_INDEX_URL}/${FLASHINFER_CUDA128_WHEEL}
392
else
393
# Exclude CUDA arches for older versions (11.x and 12.0-12.7)