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 14, 05:30 UTC
Update to flashinfer-python==0.2.12 and disable AOT compile for non-release image (#23129)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Aug 20, 2025 at 08:05 UTC
50df09fe13c93b520c64c581de4f0b469995f7b9
3 files changed
+35
-21
.buildkite/release-pipeline.yaml
+1
-1
@@ -68,7 +68,7 @@ steps:
68
queue: cpu_queue_postmerge
69
commands:
70
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
71
- - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT --target vllm-openai --progress plain -f docker/Dockerfile ."
71
+ - "DOCKER_BUILDKIT=1 docker build --build-arg max_jobs=16 --build-arg USE_SCCACHE=1 --build-arg GIT_REPO_CHECK=1 --build-arg CUDA_VERSION=12.8.1 --build-arg FLASHINFER_AOT_COMPILE=true --build-arg INSTALL_KV_CONNECTORS=true --tag public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT --target vllm-openai --progress plain -f docker/Dockerfile ."
72
- "docker push public.ecr.aws/q9t5s3a7/vllm-release-repo:$BUILDKITE_COMMIT"
73
74
- label: "Annotate release workflow"
docker/Dockerfile
+33
-19
@@ -372,31 +372,45 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
372
373
# Install FlashInfer from source
374
ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
375
-# Keep this in sync with https://github.com/vllm-project/vllm/blob/main/requirements/cuda.txt
376
-# We use `--force-reinstall --no-deps` to avoid issues with the existing FlashInfer wheel.
377
-ARG FLASHINFER_GIT_REF="v0.2.11"
375
+# Keep this in sync with "flashinfer" extra in setup.py
376
+ARG FLASHINFER_GIT_REF="v0.2.12"
377
+# Flag to control whether to compile FlashInfer AOT kernels
378
+# Set to "true" to enable AOT compilation:
379
+# docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
380
+ARG FLASHINFER_AOT_COMPILE=false
381
RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
382
. /etc/environment
383
git clone --depth 1 --recursive --shallow-submodules \
384
--branch ${FLASHINFER_GIT_REF} \
385
${FLASHINFER_GIT_REPO} flashinfer
383
- # Exclude CUDA arches for older versions (11.x and 12.0-12.7)
384
- # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
385
- if [[ "${CUDA_VERSION}" == 11.* ]]; then
386
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
387
- elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then
388
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
389
- else
390
- # CUDA 12.8+ supports 10.0a and 12.0
391
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
392
- fi
393
- echo "🏗️ Building FlashInfer for arches: ${FI_TORCH_CUDA_ARCH_LIST}"
394
- # Needed to build AOT kernels
386
pushd flashinfer
396
- TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
397
- python3 -m flashinfer.aot
398
- TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
399
- uv pip install --system --no-build-isolation --force-reinstall --no-deps .
387
+ if [ "${FLASHINFER_AOT_COMPILE}" = "true" ]; then
388
+ # Exclude CUDA arches for older versions (11.x and 12.0-12.7)
389
+ # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
390
+ if [[ "${CUDA_VERSION}" == 11.* ]]; then
391
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
392
+ elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then
393
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
394
+ else
395
+ # CUDA 12.8+ supports 10.0a and 12.0
396
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
397
+ fi
398
+ echo "🏗️ Installing FlashInfer with AOT compilation for arches: ${FI_TORCH_CUDA_ARCH_LIST}"
399
+ # Build AOT kernels
400
+ TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
401
+ python3 -m flashinfer.aot
402
+ # Install with no-build-isolation since we already built AOT kernels
403
+ TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
404
+ uv pip install --system --no-build-isolation . \
405
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
406
+ # Download pre-compiled cubins
407
+ TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
408
+ python3 -m flashinfer --download-cubin || echo "WARNING: Failed to download flashinfer cubins."
409
+ else
410
+ echo "🏗️ Installing FlashInfer without AOT compilation in JIT mode"
411
+ uv pip install --system . \
412
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
413
+ fi
414
popd
415
rm -rf flashinfer
416
BASH
setup.py
+1
-1
@@ -685,7 +685,7 @@ setup(
685
"mistral_common[audio]"], # Required for audio processing
686
"video": [], # Kept for backwards compatibility
687
# FlashInfer should be updated together with the Dockerfile
688
- "flashinfer": ["flashinfer-python==0.2.11"],
688
+ "flashinfer": ["flashinfer-python==0.2.12"],
689
},
690
cmdclass=cmdclass,
691
package_data=package_data,