执行摘要

• 一句话：修复 AITER 草稿扩展注意力 int32 溢出
• 推荐动作：建议精读此 PR，尤其是 PR 描述中关于问题定位和验证的部分。它展示了如何通过系统性的性能分析和精确的实验设计来定位一个看似是批处理大小问题、实则是数据类型溢出的隐藏 bug，对于调试类似的静默性数值问题很有参考价值。

功能与动机

PR 描述详细说明了回归的来源和影响。由于 #30105 将 EAGLE-v2 的 draft-extend 路由到 unified_attention 路径，而该路径使用 int32 seqused_k，当每层 KV 缓冲区超过 2 GiB 时会发生整型溢出，导致注意力输出产生 NaN 且无任何报错。这会使投机解码静默失败，TPOT 下降约 2.4 倍。修复该 bug 的核心动机是恢复投机解码的默认性能，并确保此类问题不会影响 AMD 用户的 AITER 后端。

实现拆解

1. 定位问题：在 aiter_backend.py 的 forward_extend 中，DRAFT_EXTEND_V2 分支下，seqused_k 从 kv_indptr 差值计算后被强制转换为 int32。
2. 修复方案：将 seqused_k 的类型改为 int64，使其与其他两处调用 unified_attention 的位置（verify 和 decode）保持一致。
3. 添加注释：在代码中添加详细注释，解释 seqused_k 必须为 int64 的原因，以及 kv_indptr 本身是 int32 因此需要显式扩宽。
4. 测试验证：PR 描述中提供了详尽的精度测试和 benchmark 数据，证明修复前 NaN、修复后正常，且性能与旧路径相当甚至略优。
5. 无其他改动：该 PR 仅修改了 aiter_backend.py 一个文件，未涉及测试或配置变更。

关键文件：

• python/sglang/srt/layers/attention/aiter_backend.py（模块 注意力后端；类别 source；类型 core-logic；符号 forward_extend）: 这是唯一的改动文件，直接修复了 draft-extend 注意力路径中的 int32 溢出问题，是本次 PR 的核心。

关键符号：forward_extend

关键源码片段

python/sglang/srt/layers/attention/aiter_backend.py

这是唯一的改动文件，直接修复了 draft-extend 注意力路径中的 int32 溢出问题，是本次 PR 的核心。

# aiter_backend.py 中 forward_extend 函数内，DRAFT_EXTEND_V2 分支
kv_indptr = self.forward_metadata.kv_indptr
# seqused_k 必须是 int64：kv_indptr 是 int32，所以差值也是 int32，必须显式扩宽。
# unified_attention 从该 dtype 推导每条 KV 的地址：如果 seqused_k 是 int32，
# 整个 K/V 偏移链会保持 32 位，一旦单层 KV 缓冲区达到 2 GiB 就会回绕，
# 静默返回 NaN。verify（seq_lens + max_q_len）和 decode（seq_lens）
# 调用点传的也都是 int64，原因相同。
seqused_k = (kv_indptr[1 : bs + 1] - kv_indptr[:bs]).to(torch.int64)

评论区精华

Review 评论较少，仅有两处 approve（来自 kkHuang-amd 和 bingxche），没有实质性的技术讨论。有一个 issue 评论是关于 CI 状态的，来自 amd-bot，指出该 PR 的 CI 并不完整，但失败均与本次改动无关。

• CI 覆盖不足 (testing): 所有失败的 CI 均与本 PR 无关，但缺少针对 2 GiB 以上 KV 缓存的测试。

风险与影响

• 风险：

  1. 回归风险：该改动是纯数据类型变更，在 2 GiB 阈值以下不会改变行为，因此回归风险极低。
  2. 性能影响：没有性能影响，因为 int64 索引计算开销可忽略。
  3. 兼容性：该改动只影响 AMD ROCm 上的 AITER 后端，且仅在启用 SGLANG_AITER_UNIFIED_DRAFT_EXTEND 时生效（默认开启）。
  4. 测试覆盖：PR 没有添加自动化测试，CI 也未覆盖 2 GiB 以上的场景，这是一个潜在的风险点。
     - 影响：对用户：AMD 平台使用 AITER 后端和 EAGLE 投机解码的用户，在 KV 缓存超过 2 GiB 时会遇到性能骤降，此修复可恢复预期性能。
     对系统：修复后，unified_attention 的 draft-extend 路径在大型 KV 缓存下表现稳定，有助于提升整体吞吐量和降低延迟。
     对团队：这是一个低风险、高价值的修复，有助于维护 AMD 后端的可靠性。

• 风险标记：核心路径变更, 缺少测试覆盖

关联脉络

• PR #30105 [AMD][Spec] Fix aiter GQA packing + split-KV routing in NEXTN spec attention (verify & draft_extend): 本 PR 修复的正是 #30105 引入的回归，需要了解其原始动机和实现细节。