将 TGI 与 AMD GPU 结合使用

TGI 在 AMD Instinct MI210、MI250 和 MI300 GPU 上受到支持和测试。未来可能会扩展支持范围。建议通过 Docker 使用。请务必查阅 AMD 文档，了解如何将 Docker 与 AMD GPU 结合使用。

在由 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.3.4-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 上略低，但需要预热，这可能会带来问题，因为它需要为每个新的提示长度再次进行。如果需要，可以在启动 TGI 的 Docker 容器时通过 --env ROCM_USE_FLASH_ATTN_V2_TRITON="0" 启用 FA Triton 实现。

自定义 PagedAttention

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

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

不支持的功能

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

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