执行摘要

本PR通过引入aiter的paged_attention_common函数优化ROCm平台上的shuffle KV缓存性能，解决了Qwen模型在小并发下性能下降的问题。变更影响ROCm注意力后端，显著提升吞吐量，但需注意对特定模型（如Qwen3.5）的正确性风险和兼容性限制。建议团队关注后续修复和动态路由设计。

功能与动机

为什么做：根据PR body，对于Qwen/Qwen3-235B-A22B-Instruct-2507-FP8模型，当前启用VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1时，在小并发场景下性能比=0更差。PR引用aiter的paged_attention_common（见https://github.com/ROCm/aiter/pull/1821）来修复此问题，旨在提升ROCm平台的注意力计算效率。

实现拆解

关键改动点：

• vllm/_aiter_ops.py：新增paged_attention_common静态方法，封装aiter库的同名函数，跳过@is_aiter_supported装饰器以允许显式后端选择。
  

  @staticmethod
  def paged_attention_common(
      Q: torch.Tensor,
      K: torch.Tensor,
      ...
  ):
      from aiter import paged_attention_common
      return paged_attention_common(...)
  

• vllm/v1/attention/backends/rocm_aiter_fa.py：在forward函数中，当rocm_aiter_ops.is_shuffle_kv_cache_enabled()为真时，改用rocm_aiter_ops.paged_attention_common，并创建临时张量（tmp_out、exp_sums、max_logits）以适配新接口。关键参数调整包括KV缩放传递和缓存视图重构。

评论区精华

核心讨论线程：

1. 设计权衡：tjtanaa询问是否移除VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT标志，samutamm回应“shuffle-path在部分场景有收益”，最终团队决定保留标志，因为paged_attention_common不支持滑动窗口等特性。

   “So far VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT seems useful in cases we've seen.” – samutamm

2. 正确性问题：tjtanaa报告Qwen3.5-397B模型输出错误，tuukkjs指出pa_fwd_asm仅支持head_dim=128，需在aiter中添加路由逻辑，后续PR 39192跟进修复。

   “pa_fwd_asm supports only head_dim=128. This model has afaik head_dim=256.” – tuukkjs

3. 代码质量：gemini-code-assist[bot]指出变量重定义和死代码，提交中修复以提升可维护性。

风险与影响

具体风险：

• 回归风险：对于head_dim!=128的模型（如Qwen3.5），paged_attention_common可能路由到不支持的内核，导致输出错误或崩溃。
• 兼容性限制：与滑动窗口注意力不兼容，需保留环境变量标志，可能影响其他优化（如rope+kvcache融合）的集成。
• 性能依赖：动态路由阈值基于硬件配置，需在不同平台上验证以优化性能。

影响评估：

• 用户：ROCm用户启用shuffle kv cache后将看到性能提升（测试显示0.5%到13.4%改善），但需注意模型兼容性。
• 系统：改进注意力计算路径，减少延迟，但增加代码复杂性和外部依赖。
• 团队：需监控后续正确性修复，并可能调整其他ROCm相关优化策略。

关联脉络

与历史PR的关系：

• PR 39192：直接后续，修复本PR中paged_attention_common对Qwen3.5模型的正确性问题。
• 其他ROCm相关PR：如39088（XPU修复）、39087（AMD内核修复），显示团队在跨平台性能优化上的持续投入。

演进趋势：本PR是ROCm平台上注意力内核优化的一部分，通过引入统一接口paged_attention_common，为动态内核选择和性能调优奠定基础，未来可能扩展支持更多模型和特性。