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
Add flashinfer-build.sh and register precompiled cu128 wheel in Dockerfile (#25782)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Sep 26, 2025 at 21:54 UTC
92da847cf5f4eedf0bc9fed45d7c076be78b8c1f
2 files changed
+83
-10
docker/Dockerfile
+20
-10
@@ -391,18 +391,28 @@ RUN --mount=type=cache,target=/root/.cache/uv bash - <<'BASH'
391
git clone --depth 1 --recursive --shallow-submodules \
392
--branch ${FLASHINFER_GIT_REF} \
393
${FLASHINFER_GIT_REPO} flashinfer
394
+ # Exclude CUDA arches for older versions (11.x and 12.0-12.7)
395
+ # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
396
+ if [[ "${CUDA_VERSION}" == 11.* ]]; then
397
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
398
+ elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then
399
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
400
+ else
401
+ # CUDA 12.8+ supports 10.0a and 12.0
402
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
403
+ fi
404
pushd flashinfer
395
- if [ "${FLASHINFER_AOT_COMPILE}" = "true" ]; then
396
- # Exclude CUDA arches for older versions (11.x and 12.0-12.7)
397
- # TODO: Update this to allow setting TORCH_CUDA_ARCH_LIST as a build arg.
398
- if [[ "${CUDA_VERSION}" == 11.* ]]; then
399
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
400
- elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then
401
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
402
- else
403
- # CUDA 12.8+ supports 10.0a and 12.0
404
- FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
405
+ if [[ "${CUDA_VERSION}" == 12.8.* ]] && [ "$TARGETPLATFORM" = "linux/amd64" ]; then
406
+ # NOTE: To make new precompiled wheels, see tools/flashinfer-build.sh
407
+ echo "🏗️ Installing FlashInfer from pre-compiled wheel"
408
+ uv pip install --system https://wheels.vllm.ai/flashinfer-python/flashinfer_python-0.3.1-cp39-abi3-manylinux1_x86_64.whl \
409
+ --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
410
+ if [ "${FLASHINFER_AOT_COMPILE}" = "true" ]; then
411
+ # Download pre-compiled cubins
412
+ TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" \
413
+ python3 -m flashinfer --download-cubin || echo "WARNING: Failed to download flashinfer cubins."
414
fi
415
+ elif [ "${FLASHINFER_AOT_COMPILE}" = "true" ]; then
416
echo "🏗️ Installing FlashInfer with AOT compilation for arches: ${FI_TORCH_CUDA_ARCH_LIST}"
417
export FLASHINFER_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}"
418
# HACK: We need these to run flashinfer.aot before installing flashinfer, get from the package in the future
tools/flashinfer-build.sh
new
+63
@@ -0,0 +1,63 @@
1
+#!/usr/bin/env bash
2
+# This script is used to build FlashInfer wheels with AOT kernels
3
+
4
+set -ex
5
+
6
+# FlashInfer configuration
7
+FLASHINFER_GIT_REPO="https://github.com/flashinfer-ai/flashinfer.git"
8
+FLASHINFER_GIT_REF="${FLASHINFER_GIT_REF}"
9
+CUDA_VERSION="${CUDA_VERSION}"
10
+BUILD_WHEEL="${BUILD_WHEEL:-true}"
11
+
12
+if [[ -z "${FLASHINFER_GIT_REF}" ]]; then
13
+ echo "❌ FLASHINFER_GIT_REF must be specified" >&2
14
+ exit 1
15
+fi
16
+
17
+if [[ -z "${CUDA_VERSION}" ]]; then
18
+ echo "❌ CUDA_VERSION must be specified" >&2
19
+ exit 1
20
+fi
21
+
22
+echo "🏗️ Building FlashInfer ${FLASHINFER_GIT_REF} for CUDA ${CUDA_VERSION}"
23
+
24
+# Clone FlashInfer
25
+git clone --depth 1 --recursive --shallow-submodules \
26
+ --branch ${FLASHINFER_GIT_REF} \
27
+ ${FLASHINFER_GIT_REPO} flashinfer
28
+
29
+# Set CUDA arch list based on CUDA version
30
+# Exclude CUDA arches for older versions (11.x and 12.0-12.7)
31
+if [[ "${CUDA_VERSION}" == 11.* ]]; then
32
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9"
33
+elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then
34
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a"
35
+else
36
+ # CUDA 12.8+ supports 10.0a and 12.0
37
+ FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 12.0"
38
+fi
39
+
40
+echo "🏗️ Building FlashInfer AOT for arches: ${FI_TORCH_CUDA_ARCH_LIST}"
41
+
42
+pushd flashinfer
43
+ # Make sure the wheel is built for the correct CUDA version
44
+ export UV_TORCH_BACKEND=cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.')
45
+
46
+ # Build AOT kernels
47
+ export TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}"
48
+ export FLASHINFER_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}"
49
+ python3 -m flashinfer.aot
50
+
51
+ if [[ "${BUILD_WHEEL}" == "true" ]]; then
52
+ # Build wheel for distribution
53
+ uv build --no-build-isolation --wheel --out-dir ../flashinfer-dist .
54
+ echo "✅ FlashInfer wheel built successfully in flashinfer-dist/"
55
+ else
56
+ # Install directly (for Dockerfile)
57
+ uv pip install --system --no-build-isolation --force-reinstall .
58
+ echo "✅ FlashInfer installed successfully"
59
+ fi
60
+popd
61
+
62
+# Cleanup
63
+rm -rf flashinfer
\ No newline at end of file