在 AMD GPU 上使用 TGI

TGI 在 AMD Instinct MI210、MI250 和 MI300 GPU 上受支持并经过测试。支持可能会在将来扩展。建议通过 Docker 使用。请务必查看 AMD 文档，了解如何在 AMD GPU 上使用 Docker。

在由 AMD GPU 支持的服务器上，可以使用以下命令启动 TGI：

model=teknium/OpenHermes-2.5-Mistral-7B
volume=$PWD/data # share a volume with the Docker container to avoid downloading weights every run

docker run --rm -it --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
    --device=/dev/kfd --device=/dev/dri --group-add video \
    --ipc=host --shm-size 256g --net host -v $volume:/data \
    ghcr.io/huggingface/text-generation-inference:2.3.1-rocm \
    --model-id $model

然后可以从客户端查询已启动的 TGI 服务器，请务必查看 使用 TGI 指南。

TunableOp

TGI 的 AMD GPU Docker 镜像集成了 PyTorch 的 TunableOp，它允许进行额外的预热以从 rocBLAS 或 hipBLASLt 中选择性能最佳的矩阵乘法（GEMM）内核。

经过实验，我们发现，在 MI300X 上，在 ROCm 6.1 和 PyTorch 2.3 之上使用 TunableOp 时，延迟降低了 6-8%。

TunableOp 默认启用，预热可能需要 1-2 分钟。如果您想禁用 TunableOp，请在启动 TGI 的 Docker 容器时传递 --env PYTORCH_TUNABLEOP_ENABLED="0"。

Flash Attention 实现

ROCm 提供了两种 Flash Attention 实现，第一种是基于 可组合内核 (CK) 实现的 ROCm/flash-attention，第二种是 Triton 实现。

默认情况下，使用可组合内核实现。但是，Triton 实现的 MI250 和 MI300 延迟略低，但需要预热，这可能是很繁重的，因为它需要针对每个新的提示长度重新执行。如果需要，可以使用 --env ROCM_USE_FLASH_ATTN_V2_TRITON="0" 启动 TGI 的 Docker 容器，从而启用 FA Triton 实现。

自定义分页注意力

为了在 ROCm 上获得更好的性能，提供了一个自定义的分页注意力内核，默认情况下已启用。要禁用它并回退到 PagedAttention v2 内核，请设置环境变量 ROCM_USE_CUSTOM_PAGED_ATTN=0。

自定义内核支持 bf16 和 fp16 数据类型，块大小为 16，头部大小为 128，最大上下文长度为 16k，GQA 比率在 1 到 16 之间。对于其他配置，我们使用 PagedAttention v2 内核。

不支持的功能

以下功能目前在 ROCm 版本的 TGI 中不支持，将来可能会扩展支持

• 加载 AWQ 检查点。
• 用于滑动窗口注意力的内核（Mistral）