vLLM AMD ROCm 注意力后端优化

vLLM 已实现一系列针对 AMD ROCm 的优化注意力后端，超越了单纯移植，实现了架构协同设计。通过将 vLLM 的调度层与 AMD 的 AITER 原语相结合，系统在 AMD Instinct MI300X、MI325X 和 MI355X GPU 上的多头注意力（MHA）和多头潜在注意力（MLA）工作负载实现了 1.2‑4.4 倍更高的吞吐量（TPS）。

混合推理工作负载的挑战

生产环境的 LLM 服务涉及持续批处理，需要同时处理预填充（prefill）、扩展（extend）和解码（decode）令牌。每个阶段都有不同的性能瓶颈：

  • Prefill：计算受限；需要大块尺寸和最大 ALU 利用率来处理新提示。
  • Extend：混合工作负载；为拥有部分 KV 缓存的请求处理额外的提示侧令牌。
  • Decode：内存受限；一次生成一个令牌，需要合并的内存访问和最小的缓存读取。

由于针对某一阶段调优的 kernel 往往在其他阶段表现不佳，vLLM 使用显式路由将每种请求类型导向专用 kernel。

ROCM_AITER_FA：针对 MHA 的三路径调度

ROCM_AITER_FA 是一个复杂的调度层，能够将请求通过三条专用路径路由，以最大化 CDNA 架构上的硬件利用率。

技术创新

  • 三路径路由：请求被动态分类为 Prefill（使用 flash_attn_varlen_func 在矩阵核心上），Extend（使用块状注意力并通过 LSE 合并处理 100K+ 上下文），以及 Decode（使用 AITER 高度优化的内存带宽 kernel）。
  • 批次重排：vLLM Model Runner 将请求重新排序为 [decode:extend:prefill] 顺序，确保连续的内存访问并消除冗余的 KV 缓存读取。
  • 硬件优化的 KV 缓存布局：预洗牌的 KV 缓存布局使内存访问模式与 AMD 的 CDNA 架构对齐。这样解码路径可以直接调用 AITER 的 pa_fwd_asm kernel，零布局转换开销，使解码吞吐量提升 15‑20%。
  • 块状上下文处理：长序列以固定的每轮令牌预算（约 32K 令牌）处理，使用基于对数求和指数（Log‑Sum‑Exp，LSE）的合并保持数值稳定性。

MHA 后端对比

除了 ROCM_AITER_FA，vLLM 还提供其他 MHA 后端：

  • 统一后端（TRITON_ATTN、ROCM_AITER_UNIFIED_ATTN）：所有令牌通过单一 kernel 路径处理。
  • 传统 2 路径（ROCM_ATTN）：为预填充使用 Triton kernel，为解码使用 HIP 分页注意力 kernel。该后端支持 Radeon GPU，但在不支持的 KV 头尺寸下可能回退到较慢的 Triton 解码 kernel。

AITER MLA 后端：针对 DeepSeek 的优化

多头潜在注意力（MLA），在 DeepSeek 和 Kimi 中使用，将 KV 缓存压缩至 576 维，改变了性能瓶颈。vLLM 提供基于 AITER 的专用 MLA 后端来处理这种压缩。

混合处理策略

MLA 后端根据处理阶段采用分割策略：

  • Prefill/Extend（未吸收）：使用标准 MHA kernel 在未压缩的表示上计算注意力。
  • Decode（已吸收）：专用 MLA kernel 直接在压缩的 576 维潜在空间上运行。

汇编带来的性能提升

ROCM_AITER_MLA 和 ROCM_AITER_TRITON_MLA 的主要性能提升来源于 mla_decode_fwd 汇编 kernel。该手工调优的 kernel 最大化 HBM3 带宽，使得相较于 TRITON_MLA 基线，输出每个令牌的时间（TPOT）提升 1.2‑1.6 倍。

性能基准

基准测试在 ROCm 7.0.0 上使用 Qwen3-235B（MHA）和 DeepSeek‑R1（MLA）模型进行。

MHA 结果（Qwen3-235B）

ROCM_AITER_FA 在所有测试硬件上显著超越传统后端的输出吞吐量（TPS）：

硬件 ROCM_AITER_FA ROCM_AITER_UNIFIED_ATTN TRITON_ATTN ROCM_ATTN
MI300X (64 req) 1.00x 1.05x 1.30x 3.82x
MI325X (64 req) 1.00x 1.02x 1.19x 4.36x
MI355X (64 req) 1.00x 0.95x 1.08x 3.61x

ROCM_AITER_FA 相比 ROCM_ATTN 在该模型上实现了 2.7‑4.4 倍更高的吞吐量，部分原因是 ROCM_ATTN 在不支持的头尺寸下回退到 Triton 解码 kernel。

MLA 结果（DeepSeek‑R1）

AITER MLA 后端提供最高 1.5 倍的吞吐量提升相较于 TRITON_MLA：

硬件 ROCM_AITER_MLA ROCM_AITER_TRITON_MLA TRITON_MLA
MI300X (64 req) 1.00x 0.98x 1.33x
MI325X (64 req) 1.00x 0.98x 1.41x
MI355X (64 req) 1.00x 1.03x 1.52x

推荐在 MLA 工作负载中默认使用 ROCM_AITER_MLA，尤其在 MI355X 上可获得最佳的首令牌时间（TTFT）。

实现与部署

要启用优化的 AITER 后端，用户需设置以下环境变量：

export VLLM_ROCM_USE_AITER=1

设置后，vLLM 会自动为 MHA 模型（如 Llama、Qwen、Mistral）选择 ROCM_AITER_FA，为 MLA 模型（如 DeepSeek、Kimi）选择 ROCM_AITER_MLA。

硬件支持矩阵

GPU 内存 架构
MI300X 192GB HBM3 gfx942
MI325X 256GB HBM3e gfx942
MI355X 288GB HBM3e gfx950

Sources