| 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 | BUILD_WHEEL="${BUILD_WHEEL:-true}" |
| 9 | |
| 10 | if [[ -z "${FLASHINFER_GIT_REF}" ]]; then |
| 11 | echo "❌ FLASHINFER_GIT_REF must be specified" >&2 |
| 12 | exit 1 |
| 13 | fi |
| 14 | |
| 15 | if [[ -z "${CUDA_VERSION}" ]]; then |
| 16 | echo "❌ CUDA_VERSION must be specified" >&2 |
| 17 | exit 1 |
| 18 | fi |
| 19 | |
| 20 | echo "🏗️ Building FlashInfer ${FLASHINFER_GIT_REF} for CUDA ${CUDA_VERSION}" |
| 21 | |
| 22 | # Clone FlashInfer |
| 23 | git clone --depth 1 --recursive --shallow-submodules \ |
| 24 | --branch "${FLASHINFER_GIT_REF}" \ |
| 25 | ${FLASHINFER_GIT_REPO} flashinfer |
| 26 | |
| 27 | # Set CUDA arch list based on CUDA version |
| 28 | # Exclude CUDA arches for older versions (11.x and 12.0-12.7) |
| 29 | if [[ "${CUDA_VERSION}" == 11.* ]]; then |
| 30 | FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9" |
| 31 | elif [[ "${CUDA_VERSION}" == 12.[0-7]* ]]; then |
| 32 | FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a" |
| 33 | elif [[ "${CUDA_VERSION}" == 12.[8-9]* ]]; then |
| 34 | # CUDA 12.8–12.9 |
| 35 | FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0a 10.3a 12.0" |
| 36 | else |
| 37 | # CUDA 13.0+ |
| 38 | FI_TORCH_CUDA_ARCH_LIST="7.5 8.0 8.9 9.0a 10.0f 11.0 12.0f" |
| 39 | fi |
| 40 | |
| 41 | echo "🏗️ Building FlashInfer AOT for arches: ${FI_TORCH_CUDA_ARCH_LIST}" |
| 42 | |
| 43 | pushd flashinfer |
| 44 | # Make sure the wheel is built for the correct CUDA version |
| 45 | export UV_TORCH_BACKEND=cu$(echo "$CUDA_VERSION" | cut -d. -f1,2 | tr -d '.') |
| 46 | |
| 47 | # Build AOT kernels |
| 48 | export TORCH_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" |
| 49 | export FLASHINFER_CUDA_ARCH_LIST="${FI_TORCH_CUDA_ARCH_LIST}" |
| 50 | python3 -m flashinfer.aot |
| 51 | |
| 52 | if [[ "${BUILD_WHEEL}" == "true" ]]; then |
| 53 | # Build wheel for distribution |
| 54 | uv build --no-build-isolation --wheel --out-dir ../flashinfer-dist . |
| 55 | echo "✅ FlashInfer wheel built successfully in flashinfer-dist/" |
| 56 | else |
| 57 | # Install directly (for Dockerfile) |
| 58 | uv pip install --system --no-build-isolation --force-reinstall . |
| 59 | echo "✅ FlashInfer installed successfully" |
| 60 | fi |
| 61 | popd |
| 62 | |
| 63 | # Cleanup |
| 64 | rm -rf flashinfer |