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
[UX] Add FlashInfer as default CUDA dependency (#26443)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Michael Goin committed
Oct 9, 2025 at 17:10 UTC
c9d33c60dcdc5fb18f4f464e01046da74e3a06e8
4 files changed
+20
-72
docker/Dockerfile
+8
-69
@@ -356,75 +356,14 @@ RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist
356
uv pip install --system dist/*.whl --verbose \
357
--extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
358
359
-# If we need to build FlashInfer wheel before its release:
360
-# $ # Note we remove 7.0 from the arch list compared to the list below, since FlashInfer only supports sm75+
361
-# $ export TORCH_CUDA_ARCH_LIST='7.5 8.0 8.9 9.0a 10.0a 12.0'
362
-# $ git clone https://github.com/flashinfer-ai/flashinfer.git --recursive
363
-# $ cd flashinfer
364
-# $ git checkout v0.2.6.post1
365
-# $ python -m flashinfer.aot
366
-# $ python -m build --no-isolation --wheel
367
-# $ ls -la dist
368
-# -rw-rw-r-- 1 mgoin mgoin 205M Jun 9 18:03 flashinfer_python-0.2.6.post1-cp39-abi3-linux_x86_64.whl
369
-# $ # upload the wheel to a public location, e.g. https://wheels.vllm.ai/flashinfer/v0.2.6.post1/flashinfer_python-0.2.6.post1-cp39-abi3-linux_x86_64.whl
370
-
371
-# Install FlashInfer from source
372
-ARG FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
373
-# Keep this in sync with "flashinfer" extra in setup.py
374
-ARG FLASHINFER_GIT_REF="v0.4.0"
375
-# Flag to control whether to compile FlashInfer AOT kernels
376
-# Set to "true" to enable AOT compilation:
377
-# docker build --build-arg FLASHINFER_AOT_COMPILE=true ...
378
-ARG FLASHINFER_AOT_COMPILE=false
379
-RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
380
- . /etc/environment
381
- git clone --depth 1 --recursive --shallow-submodules \
382
- --branch ${FLASHINFER_GIT_REF} \
383
- ${FLASHINFER_GIT_REPO} flashinfer
384
- # Exclude CUDA arches for older versions (11.x and 12.0-12.7)
385
- # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
386
- if [[ "${CUDA_VERSION}" == 11.* ]]; then
387
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
388
- elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then
389
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
390
- else
391
- # CUDA 12.8+ supports 10.0a and 12.0
392
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
393
- fi
394
- pushd flashinfer
395
- if [[ "${CUDA_VERSION}" == 12.8.* ]] && [ "$TARGETPLATFORM" = "linux/amd64" ] && [ "${FLASHINFER_GIT_REF}" = "v0.3.1" ]; then
396
- # NOTE: To make new precompiled wheels, see tools/flashinfer-build.sh
397
- echo "🏗️ Installing FlashInfer from pre-compiled wheel"
398
- uv pip install --system https://wheels.vllm.ai/flashinfer-python/flashinfer_python-0.3.1-cp39-abi3-manylinux1_x86_64.whl \
399
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
400
- if [ "${FLASHINFER_AOT_COMPILE}" = "true" ]; then
401
- # Download pre-compiled cubins
402
- TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
403
- python3 -m flashinfer --download-cubin || echo "WARNING: Failed to download flashinfer cubins."
404
- fi
405
- elif [ "${FLASHINFER_AOT_COMPILE}" = "true" ]; then
406
- echo "🏗️ Installing FlashInfer with AOT compilation for arches: ${FI_TORCH_CUDA_ARCH_LIST}"
407
- export FLASHINFER_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}"
408
- # HACK: We need these to run flashinfer.aot before installing flashinfer, get from the package in the future
409
- uv pip install --system cuda-python==$(echo $CUDA_VERSION | cut -d. -f1,2) pynvml==$(echo $CUDA_VERSION | cut -d. -f1) nvidia-nvshmem-cu$(echo $CUDA_VERSION | cut -d. -f1)
410
- # Build AOT kernels
411
- TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
412
- python3 -m flashinfer.aot
413
- # Install with no-build-isolation since we already built AOT kernels
414
- TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
415
- uv pip install --system --no-build-isolation . \
416
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
417
- # Download pre-compiled cubins
418
- TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
419
- python3 -m flashinfer --download-cubin || echo "WARNING: Failed to download flashinfer cubins."
420
- else
421
- echo "🏗️ Installing FlashInfer without AOT compilation in JIT mode"
422
- uv pip install --system . \
423
- --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
424
- fi
425
- popd
426
- rm -rf flashinfer
427
-BASH
359
+# Install FlashInfer pre-compiled kernel cache and binaries
360
+# https://docs.flashinfer.ai/installation.html
361
+RUN --mount=type=cache,target=/root/.cache/uv \
362
+ uv pip install --system flashinfer-cubin==0.4.0 \
363
+ && uv pip install --system flashinfer-jit-cache==0.4.0 \
364
+ --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \
365
+ && flashinfer show-config
366
+
367
COPY examples examples
368
COPY benchmarks benchmarks
369
COPY ./vllm/collect_env.py .
requirements/cuda.txt
+2
@@ -11,3 +11,5 @@ torchaudio==2.8.0
11
torchvision==0.23.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
12
# https://github.com/facebookresearch/xformers/releases/tag/v0.0.32.post1
13
xformers==0.0.32.post1; platform_system == 'Linux' and platform_machine == 'x86_64' # Requires PyTorch >= 2.8
14
+# FlashInfer should be updated together with the Dockerfile
15
+flashinfer-python==0.4.0
\ No newline at end of file
setup.py
+1
-2
@@ -714,8 +714,7 @@ setup(
714
"mistral_common[audio]",
715
], # Required for audio processing
716
"video": [], # Kept for backwards compatibility
717
- # FlashInfer should be updated together with the Dockerfile
718
- "flashinfer": ["flashinfer-python==0.4.0"],
717
+ "flashinfer": [], # Kept for backwards compatibility
718
# Optional deps for AMD FP4 quantization support
719
"petit-kernel": ["petit-kernel"],
720
},
vllm/utils/flashinfer.py
+9
-1
@@ -12,6 +12,7 @@ import functools
12
import importlib
13
import importlib.util
14
import os
15
+import shutil
16
from typing import Any, Callable, NoReturn
17
18
import requests
@@ -37,7 +38,14 @@ def has_flashinfer() -> bool:
38
"""Return ``True`` if FlashInfer is available."""
39
# Use find_spec to check if the module exists without importing it
40
# This avoids potential CUDA initialization side effects
40
- return importlib.util.find_spec("flashinfer") is not None
41
+ if importlib.util.find_spec("flashinfer") is None:
42
+ logger.debug_once("FlashInfer unavailable since package was not found")
43
+ return False
44
+ # Also check if nvcc is available since it's required to JIT compile flashinfer
45
+ if shutil.which("nvcc") is None:
46
+ logger.debug_once("FlashInfer unavailable since nvcc was not found")
47
+ return False
48
+ return True
49
50
51
def _missing(*_: Any, **__: Any) -> NoReturn: