使用 ROCm 安装#

vLLM 支持使用 ROCm 6.1 的 AMD GPU。

要求#

操作系统：Linux
Python：3.8 - 3.11
GPU：MI200s (gfx90a)、MI300 (gfx942)、Radeon RX 7900 系列 (gfx1100)
ROCm 6.1

安装选项：

使用 Docker 从源代码构建
从源代码构建

选项 1：使用 Docker 从源代码构建（推荐）#

你可以从源代码构建并安装 vLLM。

首先，从 Dockerfile.rocm 构建一个 Docker 镜像，并从该镜像启动一个 Docker 容器。

Dockerfile.rocm 默认使用 ROCm 6.1，但也支持旧版 vLLM 分支中的 ROCm 5.7 和 6.0。它提供了灵活性，可以使用以下参数自定义 Docker 镜像的构建：

BASE_IMAGE：指定运行 docker build 时使用的基础镜像，具体来说是 ROCm 上的 PyTorch 基础镜像。
BUILD_FA：指定是否构建 CK flash-attention。默认值为 1。对于 Radeon RX 7900 系列 (gfx1100)，在 flash-attention 支持此目标之前，应将其设置为 0。
FX_GFX_ARCHS：指定用于构建 CK flash-attention 的 GFX 架构，例如，MI200 和 MI300 的 gfx90a;gfx942。默认值为 gfx90a;gfx942
FA_BRANCH：指定在 ROCm 的 flash-attention 仓库中构建 CK flash-attention 时使用的分支。默认值为 ae7928c
BUILD_TRITON：指定是否构建 triton flash-attention。默认值为 1。

它们的值可以在运行 docker build 时使用 --build-arg 选项传递。

要在 ROCm 6.1 上为 MI200 和 MI300 系列构建 vllm，可以使用默认值：

$ DOCKER_BUILDKIT=1 docker build -f Dockerfile.rocm -t vllm-rocm .

要在 ROCm 6.1 上为 Radeon RX7900 系列 (gfx1100) 构建 vllm，应将 BUILD_FA 指定为如下：

$ DOCKER_BUILDKIT=1 docker build --build-arg BUILD_FA="0" -f Dockerfile.rocm -t vllm-rocm .

要运行上述 Docker 镜像 vllm-rocm，请使用以下命令：

$ docker run -it \
   --network=host \
   --group-add=video \
   --ipc=host \
   --cap-add=SYS_PTRACE \
   --security-opt seccomp=unconfined \
   --device /dev/kfd \
   --device /dev/dri \
   -v <path/to/model>:/app/model \
   vllm-rocm \
   bash

<path/to/model> 是模型存储的位置，例如 llama2 或 llama3 模型的权重。

选项 2：从源代码构建#

安装先决条件（如果你已经在安装了以下内容的环境/docker 中，则跳过此步骤）：

对于安装 PyTorch，你可以从一个新的 docker 镜像开始，例如 rocm/pytorch:rocm6.1.2_ubuntu20.04_py3.9_pytorch_staging、rocm/pytorch-nightly。

或者，你可以使用 PyTorch 轮子安装 PyTorch。你可以在 PyTorch 的入门指南中查看 PyTorch 安装指南。

安装 Triton flash attention for ROCm

按照 ROCm/triton 中的说明安装 ROCm 的 Triton flash attention（默认的 triton-mlir 分支）。

可选地，如果你选择使用 CK flash attention，你可以安装 flash attention for ROCm

按照 ROCm/flash-attention 中的说明安装 ROCm 的 flash attention（v2.5.9.post1）。或者，可以在发布版本下访问为 vLLM 使用的轮子。

备注

你可能需要将 “ninja” 版本降级到 1.10，因为它在编译 flash-attention-2 时不会被使用（例如 pip install ninja==1.10.2.4）。

构建 vLLM。

$ cd vllm
$ pip install -U -r requirements-rocm.txt
$ python setup.py develop # This may take 5-10 minutes. Currently, `pip install .`` does not work for ROCm installation

小技巧

例如，vLLM v0.5.3 在 ROCM 6.1 上可以通过以下步骤构建：

$ pip install --upgrade pip

$ # Install PyTorch
$ pip uninstall torch -y
$ pip install --no-cache-dir --pre torch==2.5.0.dev20240726 --index-url https://download.pytorch.org/whl/nightly/rocm6.1

$ # Build & install AMD SMI
$ pip install /opt/rocm/share/amd_smi

$ # Install dependencies
$ pip install --upgrade numba scipy huggingface-hub[cli]
$ pip install "numpy<2"
$ pip install -r requirements-rocm.txt

$ # Apply the patch to ROCM 6.1 (requires root permission)
$ wget -N https://github.com/ROCm/vllm/raw/fa78403/rocm_patch/libamdhip64.so.6 -P /opt/rocm/lib
$ rm -f "$(python3 -c 'import torch; print(torch.__path__[0])')"/lib/libamdhip64.so*

$ # Build vLLM for MI210/MI250/MI300.
$ export PYTORCH_ROCM_ARCH="gfx90a;gfx942"
$ python3 setup.py develop

小技巧

默认情况下使用 Triton flash attention。为了进行基准测试，建议在收集性能数据之前运行一个预热步骤。
Triton flash attention 目前不支持滑动窗口注意力。如果使用半精度，请使用 CK flash-attention 来支持滑动窗口。
要使用 CK flash-attention 或 PyTorch naive attention，请使用此标志 export VLLM_USE_TRITON_FLASH_ATTN=0 来关闭 triton flash attention。
理想情况下，PyTorch 的 ROCm 版本应与 ROCm 驱动程序版本匹配。

小技巧

对于 MI300x (gfx942) 用户，为了获得最佳性能，请参考 MI300x 调整指南以获取有关系统和工作流程级别的性能优化和调整技巧。对于 vLLM，请参考 vLLM 性能优化。

使用 ROCm 安装

目录

使用 ROCm 安装#

要求#

选项 1：使用 Docker 从源代码构建（推荐）#

选项 2：从源代码构建#