超越移植：vLLM 如何在 AMD ROCm 上编排高性能推理

简介

长期以来，启用 AMD 支持意味着“移植”；即仅仅让代码能跑起来。那个时代已经结束了。

随着 AMD CDNA^TM 3 架构硬件（AMD Instinct^TM MI300X、Instinct MI325X、Instinct MI355X GPU）和诸如 DeepSeek 的 MLA 等复杂模型结构的出现，“能跑起来”已远远不够。这些工作负载需要架构协同设计，即软件编排与硬件原语的紧密配合。

vLLM 现在在 AMD ROCm^TM 软件上提供 7 种注意力后端。本文将详细介绍每一种后端：它们存在的原因、权衡及其适用场景。我们提供了对比所有后端的透明基准测试，并展示了用于 MHA（多头注意力）的 ROCM_AITER_FA 和 AITER MLA（多头潜在注意力）后端如何通过 AMD 的 AITER 原语和 vLLM 的内核编排实现 1.2 至 4.4 倍的吞吐量（TPS）提升。

挑战：批处理中的混合工作负载

在生产环境的 LLM 服务中，每个推理步骤都会处理包含不同请求类型的混合批次。业界已经意识到了这一挑战——目前存在各种解决方案，从具有复杂内部调度的统一内核到使用专用内核的多路径路由。AMD 的 ROCM_AITER_FA 采用了显式路由方案，使针对工作负载的优化成为首要设计原则，而非仅仅是内部内核的细节。

• Prefill（预填充）：到达服务器的新提示词。这些请求包含数千个需要同时计算注意力的输入词元。在此阶段，GPU 正在进行繁重的矩阵乘法运算，因此预填充是计算密集型的。

• Extend（扩展）：为 KV 缓存已部分构建的请求处理额外的输入词元（例如来自块预填充、前缀缓存重用或之前的轮次）。由于这些新词元必须同时关注缓存的上下文和新鲜输入，扩展阶段是一种混合工作负载。在在线服务中，调度器利用此阶段将较长的提示词任务拆分成片段，并将其与来自其他在途请求的解码任务交错处理，从而改善延迟与吞吐量之间的平衡。

• Decode（解码）：逐个生成输出词元。每个解码步骤都会从内存中加载整个 KV 缓存以产生单个词元。瓶颈在于内存带宽，因此解码是内存密集型的。

这些请求类型随机到达，并为了效率而被打包在一起。

优化挑战：预填充需要较大的分块大小和最大的 ALU 利用率，而解码则需要合并内存访问和最小的缓存提取。针对一种工作负载优化的内核，会降低另一种工作负载的性能。

这种混合工作负载场景正是 ROCM_AITER_FA 的三路径路由所解决的问题：它不是强制所有类型的请求都通过一个内核，而是将每种类型路由到针对该工作负载特性优化的专用内核。

其他 MHA 后端

在深入了解 ROCM_AITER_FA 之前，让我们先了解一下目前可用的其他 MHA 后端。

统一注意力后端

这些后端通过单个内核路径处理所有词元（预填充/扩展/解码）。

def forward():
    # Stage 1: Save Key/Value into KV-Cache
    reshape_and_cache_flush(new_key, new_value, ...)
    # Stage 2: Single kernel for all attention
    unified_attention_kernel(new_query, KV-Cache, ...)

ROCM_ATTN：传统双路径后端

ROCM_ATTN 使用双路径路由，每个阶段采用不同的内核。

• 预填充：Triton 内核
• 解码：HIP 分页注意力内核（受支持时）

该后端有两个重要特征：

1. 传统双路径架构：为预填充（Triton）和解码（HIP 分页注意力）使用独立的内核。请注意，HIP 分页注意力内核仅支持特定的 KV 头大小——对于不受支持的配置（如 Qwen3-235B），它会回退到 Triton 解码内核，导致性能显著降低。

2. Radeon GPU 支持：与 TRITON_ATTN 一起，该后端支持 Radeon GPU——对于 AITER 原语不可用的消费级硬件部署非常有用。

ROCM_AITER_FA 后端：针对 AMD 的内核编排

ROCM_AITER_FA 不仅仅是一个内核包装器，它是一个复杂的编排层，将请求路由到专用内核，将 vLLM 的高级管理与 AMD 的 AITER 原语相结合。

关键创新

1. 三路径路由：请求被动态分类为解码、预填充和扩展路径，每一条路径都具备优化的内核。

   • 预填充路径：新序列使用 flash_attn_varlen_func——利用 CDNA 矩阵核心进行计算密集型工作。
   • 扩展路径：连续序列使用带 LSE 合并的分块注意力——有效处理 100K+ 的上下文。
   • 解码路径：单词元生成使用 AITER 高度优化的内核，以充分利用内存带宽。

   点击展开

2. 批次重排序（模型运行器）：ROCM_AITER_FA 是少数在处理前对请求进行重排序的后端之一。vLLM 的模型运行器将请求重排序为 [decode:extend:prefill] 以实现连续内存访问。每个注意力后端都通过设置 reorder_batch_threshold 来选择加入此功能——ROCM_AITER_FA 将其设置为 1，确保每个混合批次在被三路径路由消耗之前都经过重排序。

3. 分块上下文处理：长序列以固定的每迭代词元预算（总计约 32K 词元）进行分块处理，跨越扩展请求；基于 LSE 的合并确保了数值稳定性。

4. 硬件优化的 KV 缓存布局：使用由 AMD AITER 内核团队设计的预洗牌（preshuffled）KV 缓存布局。

k_cache: [num_blocks, num_heads, head_dim // x, block_size, x]
v_cache: [num_blocks, num_heads, block_size // x, head_dim, x]

这种布局将内存访问模式与 AMD CDNA 架构对齐，使解码路径能够以零布局转换开销调用 AITER 的 pa_fwd_asm 内核——与标准 KV 缓存布局相比，解码吞吐量提升了 15-20%。

为什么需要显式的三路径路由？

ROCM_AITER_FA 做出了明确的架构选择：在软件层路由工作负载，而不是依赖单个内核处理所有情况。

这种显式方法提供了：

• 可调试性：每条路径都可以独立分析、调整和优化。
• 可移植性：相同的路由逻辑在 MI300X 到 MI325X 再到 MI355X 之间通用，无需针对硬件进行更改。
• 可扩展性：可以在不重新设计核心架构的情况下添加新的工作负载类型或内核变体。
• 可预测性：执行路径是确定性的，使得性能分析变得简单明了。

扩展路径尤为重要：前缀缓存和多轮对话现在是生产部署的标准。拥有一个具备分块上下文注意力的专用路径，确保了这些工作负载获得一流的优化。

三路径处理详解

预填充路径：查询/键/值采用标准的 [num_tokens, num_heads, head_dim] 布局，以便与高度优化的 AITER MHA 内核对齐，并避免任何额外的内存拷贝操作。

扩展路径：这是最具挑战性的路径。新词元必须与存储在洗牌 KV 缓存布局中的上下文词元一起计算注意力。由于洗牌布局与 AITER 的 MHA 长上下文计算内核不兼容，我们插入了一个额外的 KV 缓存提取算子（cp_mha_gather_cache）来提取并转换上下文键/值到标准布局。长上下文被分段为多个块以管理内存。

def extend_forward():
    # Stage 1: Attention for new tokens
    flash_attn_varlen_func()  # calling AITER MHA
 
    # Stage 2: Context Chunk Loop Processing
    for chunk in context_chunks:
        cp_mha_gather_cache()      # Triton gather kernel
        flash_attn_varlen_func()   # calling AITER MHA
        merge_attn_states()        # LSE-based merge
 
    # Stage 3: Get the final result
    merge_attn_states()

每个块产生一个输出和 LSE（log-sum-exp）。LSE 捕获了 softmax 分母，从而实现数值稳定的合并——注意力分数较高的块自然会在最终结果中占据主导地位。

解码路径：直接利用洗牌 KV 缓存布局。自定义的 reshape_and_cache_flush 算子确保缓存始终处于洗牌布局中，允许注意力后端以零布局转换开销调用 AITER 高性能的 pa_fwd_asm 内核。

交互式动画：ROCM_AITER_FA 请求流

以下动画演示了多个请求如何在 7 次迭代中流经 ROCM_AITER_FA 后端。使用控件来开始、暂停或跳至特定的迭代。

迭代指南

该动画显示了 ROCM_AITER_FA 如何根据请求的状态动态地将它们路由到预填充 -> 扩展 -> 解码路径，从而实现混合工作负载的高效批处理。

AITER MLA 后端：针对 DeepSeek 的优化

DeepSeek 和 Kimi 的 MLA 架构将 KV 缓存压缩到 576 维（标准 MHA 为约 8K）——实现了 14 倍的内存缩减。这种压缩改变了注意力的性能特征，需要与标准 MHA 不同的优化策略。

混合方法

vLLM 提供了两种基于 AITER 的 MLA 后端，具有不同的预填充实现。

基础 TRITON_MLA 后端在两个阶段都使用 vLLM 的默认 Triton 内核。AITER 后端将解码内核替换为手工调优的汇编（mla_decode_fwd），这是性能提升的主要来源。这两个 AITER 后端之间的唯一区别在于预填充路径：ROCM_AITER_MLA 调用 aiter.flash_attn_varlen_func，而 ROCM_AITER_TRITON_MLA 调用 aiter.ops.triton.mha.flash_attn_varlen_func。

吸收式与非吸收式方案

所有 MLA 后端都使用相同的基本处理策略。

• 预填充/扩展（非吸收式）：在未压缩的表示上使用标准 MHA 内核计算注意力。
• 解码（吸收式）：使用直接在 576 维压缩潜在空间上运行的专用 MLA 内核。

def _forward_prefill():
    # Stage 1: Attention for new tokens (non-absorbed)
    _run_prefill_new_tokens()
 
    # Stage 2: For extend path, context chunk loop
    for chunk in context_chunks:
        gather_and_maybe_dequant_cache()
        _run_prefill_context_chunk()
        merge_attn_states()
 
    # Stage 3: Final merge
    merge_attn_states()

在解码期间，模型一次生成一个词元。压缩的 KV 缓存意味着加载的数据量更少，但由于仍然受限于内存带宽，解码仍是内存密集型的。AITER 汇编内核（mla_decode_fwd）最大限度地利用了 HBM3 带宽的每一个字节，性能显著优于通用的 Triton 解码内核。

为何汇编解码内核至关重要

两个 AITER MLA 后端（ROCM_AITER_MLA 和 ROCM_AITER_TRITON_MLA）共享相同的汇编解码内核（mla_decode_fwd）。这就是性能提升的主要来源。

1.2 至 1.6 倍的加速主要来自于共享的汇编解码内核。由于 TPOT 是解码密集型的（OSL=1K 时为 1K 次迭代），优化解码带来的吞吐量增益最大。两个 AITER 后端在预填充内核上的差异对整体性能的影响微乎其微。

除了原始内核性能外，这些后端还继承了 FlashMLABackend 的全部功能集，包括 FULL_AND_PIECEWISE CUDA 图支持和 MTP 支持。另一个优势是几乎所有的 KV 缓存块大小都能实现近乎相同的性能——您可以将每个词元视为前缀缓存，而无需担心与细粒度缓存相关的典型性能损失。

性能基准测试

基准测试方法：所有基准测试均使用 rocm/vllm-dev:nightly_main_20260115 和 ROCm 7.0.0 运行。这是一个于 2026 年 1 月 15 日从 https://github.com/vllm-project/vllm 主分支构建的夜间版 Docker 镜像。我们首先用初始请求预热了内核；报告的结果排除了首次运行，以消除 JIT 编译开销。

export SAFETENSORS_FAST_GPU=1
export VLLM_ROCM_USE_AITER=1
export VLLM_RPC_TIMEOUT=1800000
export VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1
 
# Choose backend: TRITON_ATTN, ROCM_ATTN, ROCM_AITER_FA, ROCM_AITER_UNIFIED_ATTN
ATTN_BACKEND="ROCM_AITER_FA"
 
model_path=Qwen/Qwen3-235B-A22B-Instruct-2507-FP8
vllm serve $model_path \
    --tensor-parallel-size 8 \
    --max-num-batched-tokens 16384 \
    --trust-remote-code \
    --no-enable-prefix-caching \
    --enable-expert-parallel \
    --disable-log-requests \
    --gpu_memory_utilization 0.9 \
    --attention-backend ${ATTN_BACKEND} \
    --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
    --async-scheduling \
    --port 1234

export SAFETENSORS_FAST_GPU=1
export VLLM_ROCM_USE_AITER=1
export VLLM_RPC_TIMEOUT=1800000
 
# Choose backend: TRITON_MLA, ROCM_AITER_MLA, ROCM_AITER_TRITON_MLA
ATTN_BACKEND="ROCM_AITER_MLA"
 
model_path=deepseek-ai/DeepSeek-R1-0528
vllm serve $model_path \
    --tensor-parallel-size 8 \
    --max-num-batched-tokens 16384 \
    --trust-remote-code \
    --no-enable-prefix-caching \
    --disable-log-requests \
    --gpu_memory_utilization 0.9 \
    --attention-backend ${ATTN_BACKEND} \
    --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}' \
    --async-scheduling \
    --port 1234

MHA 基准测试结果

模型：Qwen3-235B-A22B-FP8，注意力 TP8 + MoE EP8 | 工作负载：ISL=10K，OSL=1K，64 和 128 并发请求

相比 ROCM_AITER_FA 的 TPS 慢多少倍（64 并发请求）

相比 ROCM_AITER_FA 的 TPS 慢多少倍（128 并发请求）

相对性能在各代 GPU 中保持一致。ROCM_AITER_UNIFIED_ATTN（单内核路径）在此统一工作负载场景下与 ROCM_AITER_FA（三路径路由）的差距在 5% 以内——在包含前缀缓存命中的混合工作负载中，三路径路由的优势会更加明显。

注意：ROCM_ATTN 的 TPS 慢 2.7-4.4 倍，因为 Qwen3-235B 的 KV 头大小不被 HIP 分页注意力支持，被迫回退到 Triton 解码内核。ROCM_ATTN 对于具有受支持头大小的模型来说，比 TRITON_ATTN 更快。

MLA 基准测试结果

模型：DeepSeek-R1-0528，TP8，block_size=16 | 工作负载：ISL=10K，OSL=1K，64 和 128 并发请求

相比 ROCM_AITER_MLA 的 TPS 慢多少倍（64 并发请求）

相比 ROCM_AITER_MLA 的 TPS 慢多少倍（128 并发请求）

两个 AITER MLA 后端均提供相似的整体性能。在 gfx942（MI300X/MI325X）上，ROCM_AITER_TRITON_MLA 显示出高出 2-3% 的 TPS。在 gfx950（MI355X）上，ROCM_AITER_MLA 因为使用 AITER 汇编 MHA 预填充，匹配或超过了 ROCM_AITER_TRITON_MLA。ROCM_AITER_MLA 还在 MI355X 上实现了最佳 TTFT。推荐所有工作负载自动使用 ROCM_AITER_MLA。

注意：这些基准测试使用统一的请求大小。带有前缀缓存、混合上下文长度和不同请求模式的生产工作负载将更充分地发挥三路径路由架构的优势。

合作：vLLM + AITER

性能增益并非来自单一优化，而是源于 vLLM 的编排层与 AMD AITER 原语的协同工作。理解这种协作，就能解释为什么“仅仅移植”是不够的。

创新归属

性能从何而来？两层协同工作

关键洞察：AITER 提供了专为 CDNA 构建的高度优化注意力原语。vLLM 的编排层增加了针对工作负载的路由和分块处理，释放了最终的性能阶梯。两者缺一不可。

开始使用

快速入门

# Recommended: Let vLLM auto-select optimized backends
export VLLM_ROCM_USE_AITER=1
vllm serve <your-model> --tensor-parallel-size <tp>

通过 VLLM_ROCM_USE_AITER=1，vLLM 会自动选择：

• 针对 MHA 模型（Llama, Qwen, Mistral）使用 ROCM_AITER_FA
• 针对 MLA 模型（DeepSeek, Kimi）使用 ROCM_AITER_MLA

显式选择后端

对于想要实验的高级用户，可以通过 --attention-backend 指定后端。

vllm serve deepseek-ai/DeepSeek-R1-0528 \
    --tensor-parallel-size 8 \
    --attention-backend ROCM_AITER_TRITON_MLA

我们的基准测试表明，两个 AITER MLA 后端提供相似的性能，因为它们共享相同的汇编解码内核。预填充内核因架构而略有不同，但由于解码是工作负载的重头，整体差异很小。对于大多数用户，自动选择的 ROCM_AITER_MLA 表现良好。

硬件支持

完整后端参考

vLLM 在 AMD ROCm 上提供 7 种注意力后端，每种都针对不同场景进行了优化。

总结

“仅仅移植”的时代已经结束。本文介绍了 vLLM 中 AMD ROCm 上可用的所有 7 种注意力后端，并提供了透明的基准测试来显示它们的权衡。

关键结果（ISL=10K, OSL=1K 基准）

• ROCM_AITER_FA：在 MHA 模型上比 ROCM_ATTN 高出 2.7-4.4 倍的 TPS
• ROCM_AITER_MLA：通过汇编解码内核在 DeepSeek MLA 上比 TRITON_MLA 高出 1.2-1.5 倍的 TPS
• 性能在 MI300X -> MI325X -> MI355X 之间随代际提升

我们的建议：只需使用 export VLLM_ROCM_USE_AITER=1，让 vLLM 自动选择最佳后端。默认设置（MHA 使用 ROCM_AITER_FA，MLA 使用 ROCM_AITER_MLA）在所有测试工作负载中均能提供出色的性能。

这就是 AMD 原生优化的样子：不是移植，而是专为 AMD 构建。三路径路由架构反映了深思熟虑的设计选择——在软件层显式分离工作负载，每一条路径都调用硬件优化的 AITER 原语。结果是一个可调试、跨代可移植且为生产 LLM 服务混合工作负载做好准备的系统。

致谢

我们感谢所有为此次合作做出贡献的杰出人才。

AMD：Hattie Wu, Yi Gan, Zejun Chen, Carlus Huang, Lingpeng Jin, Peng Sun 以及 AITER 团队。

Embedded LLM：Pin Siang Tan, Tun Jian Tan, Jun Kang Chow 以及 Embedded LLM 团队。

资源

• AITER 库 (AMD)
• vLLM 文档
• Qwen3-235B 模型
• DeepSeek-R1 模型

免责声明

截至 2026 年 1 月 29 日，由 AMD AI 框架团队进行测试，测量 AMD Instinct MI300X、MI325X、MI355X 平台上的 TPS 推理性能。

硬件配置

• MI300X：AMD EPYC 9654 96 核处理器服务器，配备 8x AMD Instinct MI300X (192GB, 750W) GPU，Supermicro AS-8125GS-TNMR2，NPS1 (每插槽 1 个 NUMA)，2.2TiB (24 DIMMs, 4800 mts 内存, 96 GiB/DIMM)，BIOS 版本：3.2

• MI325X：AMD EPYC 9575F 64 核处理器服务器，配备 8x AMD Instinct MI325X (256GB, 1000W) GPU，Supermicro AS-8125GS-TNMR2，NPS1 (每插槽 1 个 NUMA)，2.2TiB (24 DIMMs, 4800 mts 内存, 96 GiB/DIMM)，BIOS 版本：3.2

• MI355X：AMD EPYC 9575F 64 核处理器服务器，配备 8x AMD Instinct MI355X (288GB, 1400W) GPU，Supermicro AS-8125GS-TNMR2，NPS1 (每插槽 1 个 NUMA)，2.2TiB (24 DIMMs, 4800 mts 内存, 96 GiB/DIMM)，BIOS 版本：3.2

软件配置：Ubuntu 22.04 LTS，Linux 内核 5.15.0-116-generic，ROCm 7.0 版本软件，PyTorch 2.9.0a0，vLLM 0.14.0rc2（截至 2026 年 1 月 15 日）

服务器制造商可能更改配置，从而导致结果不同。性能可能会根据配置、软件、vLLM 版本以及是否使用最新驱动程序和优化而变化。