执行摘要

• 一句话：修复 AITER MLA 小头 verify 展平缺失因果掩码
• 推荐动作：值得精读。一是 PR body 对静默缺陷的论证方式（为什么输出正常、为什么表现为吞吐提升、blast radius 的 divisor/non-divisor 划分）是排查同类隐性正确性 bug 的优秀范本；二是测试用 patch + spy 截获传给 kernel 的元数据，既不依赖 AITER 构建产物又能锁定行为，设计精巧；三是关注 min_kv_seq_len 语义从请求级到行级的迁移，以及后续是否有必要把 num_heads >= 16 与 divisor-count 场景补上实机验证。

功能与动机

PR body 明确指出：paged_kv_indptr 是 seq_lens 的 cumsum，而 seq_lens 已包含本次调度的 verify 块，因此每行获得的整个 KV 范围已经跨越待验证的块，verify 位置 t 可以 attend 到它本应检查的 draft token 之后的 token；同时 'Nothing raises and nothing warns: the output is plausible, just conditioned on tokens that have not been verified yet, and a target that has already seen the draft tokens accepts more of them, so the defect presents as a throughput win'。该缺陷来自同日合并的 #50000，无关联 issue。PR body 还分析了 blast radius：head 数整除 16 的配置（如 DeepSeek 128 头在 TP=16 时每 rank 8 头）属于回归，而 Kimi-K3 每 rank 12 头这类非整除配置此前根本无法启动。

实现拆解

1. 变更入口：vllm/v1/attention/backends/mla/rocm_aiter_mla.py 的 AiterMLAImpl.forward_mqa 中 num_heads < _AITER_MIN_MLA_HEADS 且 max_qo_len > 1 的 small-head verify flatten 分支（约 1006-1079 行）。
2. 核心逻辑：原 row_len = per_req_len[row_req] 将整条请求的 KV 长度广播给展平后的每一行；现改为 per_req_len.unsqueeze(1) - (qlen - 1) + torch.arange(qlen) 再 clamp_(min=0) 的向量化计算，使行 r*qlen + t 的窗口长度恰为 seq_len_r - (qlen - 1) + t，即已提交上下文加 t+1 个已确认 token，严格因果。clamp 保证 cudagraph padding 请求（seq_len = 0）的所有行为空窗口。由于 paged_kv_indices 每请求页按位置升序排列，因果窗口是该请求切片的前缀，无需新索引。
3. 配套调整：min_kv_seq_len 由 int(per_req_len.min()) 改为 int(row_len.min())，从“最短请求长度”变为“最短实际提交行长度”；并修正了分支顶部和后段共 3 处注释，移除原注释中错误的“committed prefix / non-causal”表述。
4. 测试配套：新增 tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py（257 行），驱动真实 AiterMLAMetadataBuilder 与 AiterMLAImpl.forward_mqa，通过 unittest.mock.patch 把 _get_mla_gluon 替换为 spy，只捕获下发给 Gluon kernel 的 page_table、seq_info、min_kv_seq_len，在 4 条真实上下文请求（含 seq_len 0 的新请求）加 1 条 cudagraph padding 请求组成的 5 请求 batch 上逐行断言因果窗口、padding clamp 与 min_kv_seq_len。测试被 pytestmark 门控在 ROCm + AITER 环境。
5. 无配置/部署配套改动：不涉及 API、schema 或构建配置变更，num_heads >= 16 的路径（走 mla_decode_fwd）完全不受影响。

关键文件：

• vllm/v1/attention/backends/mla/rocm_aiter_mla.py（模块 注意力后端；类别 source；类型 core-logic；符号 forward_mqa）: 核心修复所在：small-head（<16）多 token verify 展平分支的行长度计算由整请求广播改为严格因果窗口，并同步修正 min_kv_seq_len 与注释。
• tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py（模块 回归测试；类别 test；类型 test-coverage；符号 _on_rocm_with_aiter, _seq_lens, _expected_row_lens, _run_verify_block）: 新增回归测试，使用真实 builder + forward_mqa 并以 spy 替换 Gluon kernel，逐行断言因果窗口、padding clamp 与 min_kv_seq_len，在未修复的 main 上失败。

关键符号：forward_mqa, _expected_row_lens, test_verify_flatten_rows_are_causal, spy

关键源码片段

tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py

新增回归测试，使用真实 builder + forward_mqa 并以 spy 替换 Gluon kernel，逐行断言因果窗口、padding clamp 与 min_kv_seq_len，在未修复的 main 上失败。

# tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
# 4 条真实请求的已提交上下文 + 1 条 cudagraph padding（seq_len 0）；
# QLEN = 8 即 1 + num_speculative_tokens，也就是 verify 块长度。
def _seq_lens() -> list[int]:
    return [c + QLEN for c in CONTEXT_LENS] + [0] * PADDING_ROWS
​
# 因果行长度：展平行 r*qlen + t 只能看到 seq_len_r - (qlen - 1) + t 个 KV 条目
# （已提交上下文 + t + 1 个已确认 token），padding 请求全部 clamp 为 0。
def _expected_row_lens() -> list[int]:
    return [max(0, s - (QLEN - 1) + t) for s in _seq_lens() for t in range(QLEN)]
​
# 用 spy 替换 Gluon kernel：只截获实际上交给 kernel 的元数据，
# 既不运行真实 kernel，也就不依赖 AITER 构建产物。
def spy(**kwargs):
    captured["page_table"] = kwargs["page_table"].detach().clone()
    captured["seq_info"] = kwargs["seq_info"].detach().clone()
    captured["min_kv_seq_len"] = kwargs["min_kv_seq_len"]
​
with patch(
    "vllm.v1.attention.backends.mla.rocm_aiter_mla._get_mla_gluon",
    lambda: spy,
):
    impl.forward_mqa((q_nope, q_pe), kv_cache, metadata, layer=None)
​
# 测试主体：由 seq_info 差分还原每行 KV 长度，再与严格因果期望比对。
indptr = captured["seq_info"].tolist()
got_row_lens = [indptr[i + 1] - indptr[i] for i in range(len(indptr) - 1)]
assert got_row_lens == want_row_lens, (
    "AITER MLA verify flatten is not causal: verify row r*qlen+t must get "
    f"seq_len_r - {QLEN - 1} + t KV entries."
    "Rows longer than expected let a verify position attend to the draft "
    "tokens it is supposed to be checking."
)
# padding 请求（seq_len 0）的每一行都必须 clamp 到空窗口
padding_rows = got_row_lens[len(CONTEXT_LENS) * QLEN:]
assert padding_rows == [0] * (PADDING_ROWS * QLEN)
# min_kv_seq_len 必须等于最短实际提交行的长度
assert captured["min_kv_seq_len"] == min(want_row_lens)

评论区精华

本 PR 无行内 review 评论。作者在 issue 评论区主动向 #50000 相关维护者 @ZJY0516 补充缺陷分析："It's silent, and because the target then agrees with the drafter more often, it presents as higher acceptance and reads like a throughput win rather than a bug." 并说明 CI 尚未运行、需要 ready 标签。维护者 dllehr-amd 直接批准："Thanks @yudigege86 This one looks good to me!"。claude[bot] 提示 fork 来源 PR 自动 review 被禁用。PR body 自述验证局限：divisor-count 回归的论断与 num_heads >= 16 不受影响的断言均来自读代码而非实机执行，两个分支只在 12 头下运行过。

• 作者向 #50000 维护者同步静默缺陷分析 (other): 维护者 dllehr-amd 审阅后直接批准（"This one looks good to me!"）；无进一步追问。
• fork PR 自动 review 禁用与快速批准 (other): 无需人工介入即完成合入流程。

风险与影响

• 风险：
  1. 行为回调风险：修复后 verify 位置不再能看到 draft token，目标模型接受率会下降，端到端吞吐可能相比"带缺陷版本"看起来变差；这是正确性修复的预期代价，但需要向用户澄清，避免被误读为性能回退。
  2. kernel 容忍性风险：padding 请求行窗口 clamp 为 0，page_table 中这些行为空切片、min_kv_seq_len 可能为 0；PR 中 Gluon kernel 被 spy 替代未实际执行，kernel 对空窗口的实际行为只经推理未验证。
  3. 前提假设风险：因果前缀方案依赖 paged_kv_indices 按位置升序排列页；若未来索引器改变页序（如乱序重排），"窗口即前缀"的假设会失效。
  4. 覆盖范围局限：测试门控在 ROCm + AITER 环境，非 ROCm CI 不覆盖；且 num_heads >= 16 与 divisor-count 回归的验证仅为代码推演，缺少执行证据。
     - 影响：影响用户：ROCm 平台上启用 AITER MLA 后端并运行推测解码、每 rank 查询头数小于 16 的部署，主要是 DeepSeek-V3/R1 在 TP=16（8 头/rank）以及 Kimi-K3 在 TP=8（12 头/rank）的场景；修复后验证阶段的 attention 掩码正确，接受率将由"虚高"回归到真实水平，输出质量不再受未验证 token 污染。对系统而言，每行 KV 窗口缩短，实际减少 Gluon kernel 的 KV 读取工作量，方向上是性能优化。对团队而言，该 PR 为 small-head MLA 展平路径补上了正确性基线，其"spy 替换 kernel、驱动真实 builder"的测试范式可复用到其他后端路径。
     - 风险标记：静默正确性缺陷修复, kernel 空窗口容忍性仅代码推演, 因果方案依赖页序升序假设, num_heads>=16 路径未实机验证, 测试依赖 ROCm + AITER 环境

关联脉络

• PR #50000 [ROCm][MLA] 引入 Gluon decode 路径与 small-head verify flatten（原标题未在材料中提供）: PR body 明确指出本缺陷由该 PR 引入：Gluon decode 路径、verify flatten 与放宽的 head 数检查同日合并；#50000 之前 head 数不整除 16 的配置根本无法启动，本 PR 修复了它造成的回归。
• PR #50302 [Bugfix] Universally align block table width to 128 tokens: 同为 v1 引擎中 DeepSeek-family MLA 注意力正确性修复，涉及 paged-KV 索引与 MLA 索引器（mla/indexer.py）的边界对齐；与本 PR 同属 v1 MLA 索引/掩码正确性演进脉络，可对照阅读。