使用 TGI 与 AMD GPU

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:3.2.2-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 实现，第一种是基于 Composable Kernel (CK) 实现的 ROCm/flash-attention，第二种是 Triton 实现。

默认情况下，使用 Composable Kernel 实现。但是，Triton 实现在 MI250 和 MI300 上具有稍低的延迟，但需要预热，这可能会成为问题，因为它需要为每个新的提示长度再次执行。如果需要，可以使用 --env ROCM_USE_FLASH_ATTN_V2_TRITON="0" 在启动 TGI 的 Docker 容器时启用 FA Triton 实现。

自定义 PagedAttention

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

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

不支持的功能

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

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