分享一个 vllm v0.22.0版本的安装脚本

[Billy Shen]

#!/bin/bash
# ============================================================
#  setup_vllm.sh
#  vLLM 0.22.0 全新环境重建 & 部署脚本
#
#  包管理: uv
#  目标机器: Ubuntu 24.04 + NVIDIA RTX 5880 Ada (SM 8.9)
#            CUDA Driver 595.x  |  CUDA Toolkit 13.2
#            Python 3.12
#
#  用法:
#    chmod +x setup_vllm.sh
#    bash setup_vllm.sh
#
#  脚本执行流程:
#    1. 切换 nvcc → CUDA 13.2
#    2. 安装 / 升级 uv
#    3. 销毁旧虚拟环境
#    4. 用 uv 创建全新 Python 3.12 虚拟环境
#    5. 用 uv 安装 vLLM 0.22.0（--torch-backend=cu130）
#    6. Patch flashinfer sampling.cuh（FlagHeads → SubtractLeft）
#    7. 安装 flashinfer-jit-cache（消除运行时 JIT 编译）
#    8. 清除编译缓存
#    9. 验证环境
# ============================================================

set -euo pipefail

# ── 颜色输出 ──────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'
info()    { echo -e "${CYAN}[INFO]${NC}  $*"; }
ok()      { echo -e "${GREEN}[OK]${NC}    $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC}  $*"; }
error()   { echo -e "${RED}[ERROR]${NC} $*"; exit 1; }
section() {
    echo -e "\n${BOLD}══════════════════════════════════════════${NC}"
    echo -e "${BOLD}  $*${NC}"
    echo -e "${BOLD}══════════════════════════════════════════${NC}"
}

# ── 配置变量（按需修改）──────────────────────────────────
VLLM_VERSION="0.22.0"
CUDA_TOOLKIT_VER="13.2"                            # 系统已安装的 CUDA Toolkit
CUDA_PATH="/usr/local/cuda-${CUDA_TOOLKIT_VER}"   # nvcc 所在路径
TORCH_BACKEND="cu130"                              # uv --torch-backend 参数
VENV_DIR="$HOME/vllm_env"                         # 虚拟环境目录
PYTHON_VER="3.12"

# ── 开始 ──────────────────────────────────────────────────
echo ""
echo -e "${BOLD}╔══════════════════════════════════════════════╗${NC}"
echo -e "${BOLD}║   vLLM ${VLLM_VERSION} 全新环境重建脚本 (uv)       ║${NC}"
echo -e "${BOLD}╚══════════════════════════════════════════════╝${NC}"
echo ""

# ============================================================
# Step 1: 切换 nvcc 到 CUDA 13.2
# ============================================================
section "Step 1/9  切换 nvcc → CUDA ${CUDA_TOOLKIT_VER}"

if [ ! -d "$CUDA_PATH" ]; then
    warn "未找到 $CUDA_PATH，当前可用版本："
    ls /usr/local/ | grep cuda || true
    error "请修改脚本顶部 CUDA_TOOLKIT_VER 为实际已安装版本"
fi

# 幂等写入 bashrc
if ! grep -q "cuda-${CUDA_TOOLKIT_VER}/bin" "$HOME/.bashrc" 2>/dev/null; then
    {
        echo ""
        echo "# === CUDA ${CUDA_TOOLKIT_VER} (added by setup_vllm.sh) ==="
        echo "export PATH=${CUDA_PATH}/bin:\$PATH"
        echo "export LD_LIBRARY_PATH=${CUDA_PATH}/lib64:\$LD_LIBRARY_PATH"
        echo "export CUDA_HOME=${CUDA_PATH}"
    } >> "$HOME/.bashrc"
    ok "已写入 ~/.bashrc"
else
    ok "~/.bashrc 中已有 CUDA ${CUDA_TOOLKIT_VER} 配置，跳过"
fi

# 当前 shell 立即生效
export PATH="${CUDA_PATH}/bin:$PATH"
export LD_LIBRARY_PATH="${CUDA_PATH}/lib64:${LD_LIBRARY_PATH:-}"
export CUDA_HOME="${CUDA_PATH}"

NVCC_OUT=$(nvcc --version 2>/dev/null | grep "release" | awk '{print $5,$6}' | tr -d ',' || echo "未找到")
ok "nvcc: ${NVCC_OUT}"

# ============================================================
# Step 2: 安装 / 升级 uv
# ============================================================
section "Step 2/9  安装 / 升级 uv"

if command -v uv &>/dev/null; then
    UV_VER=$(uv --version 2>/dev/null | awk '{print $2}')
    info "已检测到 uv ${UV_VER}，执行升级..."
    uv self update 2>/dev/null && ok "uv 已升级" || warn "uv 升级失败，继续使用当前版本"
else
    info "未检测到 uv，正在安装..."
    curl -LsSf https://astral.sh/uv/install.sh | sh
    ok "uv 安装完成"
fi

# 确保 uv 在 PATH 中
export PATH="$HOME/.local/bin:$PATH"
if ! grep -q '\.local/bin' "$HOME/.bashrc" 2>/dev/null; then
    echo 'export PATH="$HOME/.local/bin:$PATH"' >> "$HOME/.bashrc"
fi

UV_VER=$(uv --version 2>/dev/null | awk '{print $2}' || echo "未知")
ok "uv ${UV_VER} 就绪"

# ============================================================
# Step 3: 销毁旧虚拟环境
# ============================================================
section "Step 3/9  销毁旧虚拟环境 (${VENV_DIR})"

if [ -d "$VENV_DIR" ]; then
    info "正在删除 ${VENV_DIR} ..."
    rm -rf "$VENV_DIR"
    ok "已删除旧虚拟环境"
else
    ok "旧虚拟环境不存在，跳过"
fi

# ============================================================
# Step 4: 用 uv 创建全新虚拟环境
# ============================================================
section "Step 4/9  用 uv 创建 Python ${PYTHON_VER} 虚拟环境"

uv venv "$VENV_DIR" --python "${PYTHON_VER}"
source "$VENV_DIR/bin/activate"
ok "虚拟环境已创建并激活: ${VENV_DIR}"

# ============================================================
# Step 5: 用 uv 安装 vLLM 0.22.0
# ============================================================
section "Step 5/9  用 uv 安装 vLLM ${VLLM_VERSION} (--torch-backend=${TORCH_BACKEND})"

# uv --torch-backend=cu130 自动从 PyTorch 官方 whl 索引拉取 torch 2.11.0+cu130
# vLLM wheel 会同时携带匹配版本的 flashinfer-python / flashinfer-cubin
info "正在安装 vllm==${VLLM_VERSION}（torch 2.11+cu130 约 600MB，请耐心等待）..."

uv pip install "vllm==${VLLM_VERSION}" \
    --torch-backend="${TORCH_BACKEND}" \
    2>&1 | tee /tmp/vllm_install.log \
         | grep -E "^(Resolved|Prepared|Installed|Uninstalled|error|Error|warning)" \
    || true

# 验证安装
INSTALLED_VLLM=$(uv pip show vllm 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
if [ "$INSTALLED_VLLM" != "$VLLM_VERSION" ]; then
    error "vLLM 安装失败（期望 ${VLLM_VERSION}，实际 ${INSTALLED_VLLM:-未安装}）\n       详情: /tmp/vllm_install.log"
fi
ok "vLLM ${INSTALLED_VLLM} 安装成功"

TORCH_VER=$(python -c "import torch; print(torch.__version__)" 2>/dev/null || echo "未知")
ok "torch: ${TORCH_VER}"

FI_PY=$(uv pip show flashinfer-python 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "未安装")
FI_CB=$(uv pip show flashinfer-cubin  2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "未安装")
ok "flashinfer-python: ${FI_PY}"
ok "flashinfer-cubin:  ${FI_CB}"

# ============================================================
# Step 6: Patch flashinfer sampling.cuh
#         FlagHeads → SubtractLeft
#         原因: flashinfer-cubin 0.6.x 内置 CUB v3.3.2 删除了
#               BlockAdjacentDifference::FlagHeads，改名为 SubtractLeft
# ============================================================
section "Step 6/9  Patch flashinfer sampling.cuh (FlagHeads → SubtractLeft)"

FLASHINFER_DIR=$(python -c \
    "import flashinfer, os; print(os.path.dirname(flashinfer.__file__))" \
    2>/dev/null || echo "")

if [ -z "$FLASHINFER_DIR" ]; then
    warn "无法找到 flashinfer 安装路径，跳过 patch"
else
    SAMPLING_CUH="${FLASHINFER_DIR}/data/include/flashinfer/sampling.cuh"

    if [ ! -f "$SAMPLING_CUH" ]; then
        warn "未找到 ${SAMPLING_CUH}，可能此版本已官方修复，跳过"
    else
        FLAG_COUNT=$(grep -c "\.FlagHeads(" "$SAMPLING_CUH" 2>/dev/null || echo 0)
        PATCHED_COUNT=$(grep -c "\.SubtractLeft(" "$SAMPLING_CUH" 2>/dev/null || echo 0)

        if [ "$FLAG_COUNT" -eq 0 ] && [ "$PATCHED_COUNT" -gt 0 ]; then
            ok "sampling.cuh 已是 patch 后版本，跳过"
        elif [ "$FLAG_COUNT" -eq 0 ]; then
            ok "未发现 FlagHeads，无需 patch"
        else
            cp "$SAMPLING_CUH" "${SAMPLING_CUH}.bak"
            sed -i 's/\.FlagHeads(/.SubtractLeft(/g' "$SAMPLING_CUH"
            ok "Patch 完成：${FLAG_COUNT} 处 FlagHeads → SubtractLeft（原文件已备份）"
        fi
    fi
fi

# ============================================================
# Step 7: 安装 flashinfer-jit-cache（消除运行时 JIT 编译）
# ============================================================
section "Step 7/9  安装 flashinfer-jit-cache（预编译 kernel）"

info "从 flashinfer.ai 安装 jit-cache (cu130)..."

uv pip install flashinfer-jit-cache \
    --index-url "https://flashinfer.ai/whl/cu130" \
    2>&1 | grep -E "^(Resolved|Prepared|Installed|error|Error|No solution|warning)" \
    || true

JIT_VER=$(uv pip show flashinfer-jit-cache 2>/dev/null | grep "^Version:" | awk '{print $2}' || echo "")
if [ -n "$JIT_VER" ]; then
    ok "flashinfer-jit-cache ${JIT_VER} 安装成功（首次启动无需 JIT 编译）"
else
    warn "flashinfer-jit-cache 未安装（cu130 wheel 可能尚未发布）"
    warn "首次启动会自动 JIT 编译，约需 2~5 分钟，属正常现象"
fi

# ============================================================
# Step 8: 清除编译缓存
# ============================================================
section "Step 8/9  清除所有编译缓存"

for CACHE_DIR in \
    "$HOME/.cache/flashinfer" \
    "$HOME/.cache/vllm/torch_compile_cache" \
    "$HOME/.cache/vllm/torch_aot_compile"
do
    if [ -d "$CACHE_DIR" ]; then
        rm -rf "$CACHE_DIR"
        ok "已清除: $CACHE_DIR"
    else
        info "不存在，跳过: $CACHE_DIR"
    fi
done

# ============================================================
# Step 9: 验证环境
# ============================================================
section "Step 9/9  验证环境"

echo ""
echo "  ┌───────────────────────────────────────────────────┐"
echo "  │  环境验证报告                                      │"
echo "  ├───────────────────────────────────────────────────┤"

_row() { printf "  │  %-22s %-24s │\n" "$1" "$2"; }

_row "uv:"                "$(uv --version 2>/dev/null | awk '{print $2}' || echo '未找到')"
_row "nvcc:"              "$(nvcc --version 2>/dev/null | grep 'release' | awk '{print $5,$6}' | tr -d ',' || echo '未找到')"
_row "Python:"            "$(python --version 2>&1)"
_row "torch:"             "$(python -c 'import torch; print(torch.__version__)' 2>/dev/null || echo '未安装')"
_row "torch CUDA:"        "$(python -c 'import torch; print(torch.version.cuda)' 2>/dev/null || echo 'N/A')"
_row "GPU:"               "$(python -c 'import torch; print(torch.cuda.get_device_name(0))' 2>/dev/null || echo '无法获取')"
_row "vllm:"              "$(uv pip show vllm 2>/dev/null | grep '^Version:' | awk '{print $2}' || echo '未安装')"
_row "flashinfer-python:" "$(uv pip show flashinfer-python 2>/dev/null | grep '^Version:' | awk '{print $2}' || echo '未安装')"
_row "flashinfer-cubin:"  "$(uv pip show flashinfer-cubin  2>/dev/null | grep '^Version:' | awk '{print $2}' || echo '未安装')"
_row "flashinfer-jit:"    "$(uv pip show flashinfer-jit-cache 2>/dev/null | grep '^Version:' | awk '{print $2}' || echo '未安装')"

echo "  ├───────────────────────────────────────────────────┤"

# fast_decode_plan 导入测试
if python -c "from flashinfer.decode import fast_decode_plan" 2>/dev/null; then
    _row "fast_decode_plan:" "✅ 可导入"
else
    _row "fast_decode_plan:" "❌ 导入失败"
fi

# sampling.cuh patch 检查
if [ -n "${FLASHINFER_DIR:-}" ]; then
    CUH="${FLASHINFER_DIR}/data/include/flashinfer/sampling.cuh"
    if [ -f "$CUH" ]; then
        REMAIN=$(grep -c "\.FlagHeads(" "$CUH" 2>/dev/null || echo 0)
        [ "$REMAIN" -eq 0 ] \
            && _row "sampling.cuh patch:" "✅ 已修复" \
            || _row "sampling.cuh patch:" "❌ 仍有 FlagHeads"
    fi
fi

echo "  └───────────────────────────────────────────────────┘"

# ============================================================
# 完成提示
# ============================================================
echo ""
echo -e "${BOLD}╔══════════════════════════════════════════════╗${NC}"
echo -e "${BOLD}║   ✅  环境重建完成                           ║${NC}"
echo -e "${BOLD}╚══════════════════════════════════════════════╝${NC}"
echo ""
echo "  激活环境并启动 vLLM："
echo ""
echo -e "  ${CYAN}source ${VENV_DIR}/bin/activate${NC}"
echo -e "  ${CYAN}bash ~/vllm/start_vllm.sh${NC}"
echo ""
echo "  启动成功标志（日志中出现以下两行）："
echo "    INFO [topk_topp_sampler.py] Using FlashInfer for top-p & top-k sampling."
echo "    INFO [api_server.py] Application startup complete."
echo ""

[terry]

帖子不错，下次发到对应版区，不要乱发。发帖前让AI整理成markdown格式，代码就用代码格式或者附件格式发布，这样便于阅读。