超越移植：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 的 3 路径路由所解决的问题：它不是强迫所有请求类型通过一个内核，而是将每种类型路由到针对该工作负载特性优化的专用内核。

其他 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：遗留的 2 路径后端

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

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

此后端具有两个重要特性

1. 遗留的 2 路径架构：预填充（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 内核团队设计的预洗牌 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%。

为什么选择显式 3 路径路由？

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

这种显式方法提供了

• 可调试性：每个路径都可以独立进行分析、调整和优化
• 可移植性：相同的路由逻辑在 MI300X → MI325X → MI355X 上均可工作，无需特定于硬件的更改
• 可扩展性：无需重新设计核心架构即可添加新的工作负载类型或内核变体
• 可预测性：执行路径是确定性的，使性能分析变得简单直接

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

三路径处理详解

预填充路径：Query/Key/Value 处于标准 [num_tokens, num_heads, head_dim] 布局中，以与高度优化的 AITER MHA 内核对齐，并避免任何额外的内存拷贝操作。

扩展路径：这是最具挑战性的路径。新令牌必须使用存储在洗牌 KV 缓存布局中的上下文令牌来计算注意力。由于洗牌布局与 AITER 的 MHA 内核不兼容长上下文计算，我们插入了一个额外的 KV 缓存获取运算符（cp_mha_gather_cache）来获取并将上下文 Key/Value 转换为标准布局。长上下文被分块成段以管理内存

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（AITER MHA 自动分发到 CK 或汇编内核），而 ROCM_AITER_TRITON_MLA 调用 aiter.ops.triton.mha.flash_attn_varlen_func（AITER Triton MHA）。

吸收式与非吸收式配方

所有 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 7.0.0 的 rocm/vllm-dev:nightly_main_20260115 运行。这是一个于 2026 年 1 月 15 日从 https://github.com/vllm-project/vllm 的 main 分支构建的夜间 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 个并发请求

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

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

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

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

MLA 基准测试结果

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

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

TPS 相比 ROCM_AITER_MLA 慢多少倍（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。

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

合作：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 基准测试）

• MHA 模型上 ROCM_AITER_FA 的 TPS 比 ROCM_ATTN 高 2.7-4.4 倍
• 通过汇编解码内核，DeepSeek MLA 上 ROCM_AITER_MLA 的 TPS 比 TRITON_MLA 高 1.2-1.5 倍
• 性能随 MI300X → MI325X → MI355X 扩展

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

这就是原生 AMD 优化的样子：不是移植的，而是专门构建的。3 路径路由架构反映了一个慎重的设计选择——在软件层进行显式的工作负载分离，每个路径调用硬件优化的 AITER 原语。结果是一个可调试、跨 GPU 代际可移植，并为生产 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 个 DIMM，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 个 DIMM，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 个 DIMM，4800 mts 内存，96 GiB/DIMM），BIOS 版本：3.2

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

服务器制造商可能会改变配置，从而产生不同的结果。性能可能因配置、软件、vLLM 版本以及使用最新的驱动程序和优化而异。