执行摘要

• 一句话：修复 ROCm CUDA graph capture 时 CPU query offset 被清除
• 推荐动作：值得精读，尤其是关注 ROCm CUDA graph capture 细节的开发者。该 PR 揭示了一个隐藏在共享 metadata 状态下的时序问题，并提供了最小化的修复方案，展示在性能和安全捕获之间的权衡。可作为处理类似共享状态清理问题的参考案例。

功能与动机

在 ROCm 上，完整 CUDA graph capture 在混合注意力模型的服务器启动阶段可能失败，因为后续 metadata builder 读取到的 query 长度为 0。根本原因是 RocmAttentionMetadataBuilder 同时清除了设备端与 CPU 端的 query offsets，以防止 prefix prefill 内核的非法内存访问。设备端必须置零以确保安全捕获，但 CPU 端 offset 不被该内核使用，且与后续 builder 共享，清除后使后续 builder 无法计算有效的 query 长度。

实现拆解

1. 定位问题：在 vllm/v1/attention/backends/rocm_attn.py 的 build_for_cudagraph_capture 方法中，原代码同时调用 query_start_loc.zero_() 和 query_start_loc_cpu.zero_() 将设备端与 CPU 端 query start locations 清零。
2. 修改逻辑：仅保留设备端 query_start_loc.zero_() 的调用，移除 CPU 端的清零操作。同时更新注释说明，明确设备端置零是为了避免 prefix prefill 内核在 graph capture 期间的非法内存访问（见 #25985），并保留 CPU 端 offset 供后续 metadata builder 使用。
3. 影响范围：此改动仅涉及 RocmAttentionMetadataBuilder，该 builder 被 ROCM_ATTN 和 ROCM_AITER_UNIFIED_ATTN 两个 attention 后端共享，因此对 ROCm 上的 CUDA graph capture 行为统一生效。
4. 测试配套：PR 未附带自动化测试，但作者在 4 张 MI325 GPU 上使用 EmbeddedLLM/MiniMax-M3-FP8-dynamic 模型进行了手动验证，修复前 CUDA graph profiling 失败（断言 decode_query_len > 0），修复后捕获成功并完成请求。

关键文件：

• vllm/v1/attention/backends/rocm_attn.py（模块 注意力后端；类别 source；类型 core-logic；符号 build_for_cudagraph_capture）: 该文件包含 RocmAttentionMetadataBuilder，是本次修复的核心位置。修复使其在 CUDA graph capture 时仅清零设备端 query start locations，同时保留 CPU 端 offset，避免后续 metadata builder 因读取到零 query 长度而崩溃。

关键符号：build_for_cudagraph_capture

关键源码片段

vllm/v1/attention/backends/rocm_attn.py

该文件包含 RocmAttentionMetadataBuilder，是本次修复的核心位置。修复使其在 CUDA graph capture 时仅清零设备端 query start locations，同时保留 CPU 端 offset，避免后续 metadata builder 因读取到零 query 长度而崩溃。

# vllm/v1/attention/backends/rocm_attn.py
class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadata]):
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.ALWAYS
​
    def build_for_cudagraph_capture(
        self, common_attn_metadata: CommonAttentionMetadata
    ) -> RocmAttentionMetadata:
        attn_metadata = self.build(0, common_attn_metadata)
        # 将 seq_lens 填充为 1，避免完整图捕获时因 max_model_len 过慢。
        attn_metadata.seq_lens.fill_(1)
​
        # 仅将设备端 query start locations 清零，避免图捕获时
        # prefix prefill 内核触发非法内存访问（见 #25985）。
        # CPU 端的 offset 保留，供后续 metadata builder 计算 query 长度。
        common_attn_metadata.query_start_loc.zero_()
​
        return attn_metadata

评论区精华

1. pre-commit 争议：mergify bot 尝试 rebase 分支后导致 pre-commit 失败，作者 akii96 指出其初始提交已通过，并对流程表示困惑，最终在其他维护者协助下手动 rebase 并触发 CI。
2. 经验验证：作者在 issue 评论中确认，InfX 的 amd/MiniMax-M3-MXFP4 recipe 使用 Triton attention 后端也存在相同问题，验证了修复的通用性。
3. 维护者核准：tjtanaa 审查后批准该 PR，未提出代码层面的异议。

• pre-commit 失败与 rebase 流程 (other): 在 AndreasKaratzas 协助下手动 rebase 并重新触发 CI，最终通过。

风险与影响

• 风险：
  1. 回归风险：修改仅移除了 CPU 端 query_start_loc 的置零操作，设备端置零保持不变，因此对 prefix prefill 内核的捕获行为无影响。但需确认后续 metadata builder 是否依赖 CPU 端 offset 被清零的预期，若存在此类依赖可能引入行为变化。
  2. 性能影响：CPU 端 offset 保留后，后续 builder 能正确计算 query 长度，可能增加少量计算量，但影响极小。
  3. 兼容性：该改动适用于所有使用 RocmAttentionMetadataBuilder 的 ROCm 后端，但未涵盖其他平台（如 NVIDIA）的类似问题，若其他平台存在同样问题需另行处理。
• 影响：
  1. 用户层面：修复了 ROCm 平台上使用完整 CUDA graph capture 的混合注意力模型（如 MiniMax-M3）的启动崩溃问题，这些模型现在可以正常启动并完成推理请求。
  2. 系统层面：改动仅限于 ROCm attention 后端的 metadata 构建逻辑，不涉及设备端内核或数据流，对系统整体稳定性和性能无显著影响。
  3. 团队层面：该修复统一解决了 ROCM_ATTN 与 ROCM_AITER_UNIFIED_ATTN 后端的共享问题，减少后续相关 bug 排查成本，但缺少自动化测试可能增加回归风险。
     - 风险标记：缺少自动化测试, 平台限定（仅 ROCm）

关联脉络

• PR #53004 [ROCm][CI] Speed up test_rocm_aiter_qk_norm_rope_kvcache_fusion: 同为 ROCm 平台相关的改动，涉及 CUDA graph 捕获后端的性能优化，与本 PR 关注的 ROCm 捕获行为有潜在关联。
• PR #52819 [ROCm]: Bump triton 3.7 commit: 涉及 ROCm 基础镜像和 Triton 升级，可能影响 attention 后端的捕获行为，与本 PR 的 ROCm 捕获修复有间接关联。