[Bug][ROCm]: ROCM_AITER_MLA can select an aiter MLA decode kernel that silently returns uncorrelated output (gqa_ratio=16, fp8 Q+KV, decode_

快速结论：该问题出现在使用 ROCm 版 vLLM（ROCM_AITER_MLA 后端）+ aiter MLA decode 的场景：当 aiter 以非持久（non-persistent）路径、fp8 Q + fp8 KV、gqa_ratio=16、decode_qlen 为 3–4 调用 mla_decode_fwd 时，会选中 mla_a8w8_qh64_qseqlen4_gqaratio16 内核并返回与参考结果不相关的输出（无崩溃、无警告）。优先排查你的调用是否缺少 persistent metadata（work_meta_data），以及是否绕过了 vLLM 正常后端直接调用 aiter。

适用环境：ROCm，GPU 架构 gfx950（Issue 评论中使用 MI355X 验证），aiter 0.1.28.post1，vLLM main 分支，MLA 模型 + fp8 KV、投机解码（MTP/spec，depth 3）。Issue 未确认 Python / PyTorch / CUDA 具体版本。

最快修复方案：暂无确认的一步修复方案。经 vLLM 维护者验证，正常走 vLLM 时该 shape 上的 decode 调用都会带上 persistent metadata，不会选到出错的非持久内核；因此“优先尝试”的是确认调用入口是否绕过了 vLLM、直接使用非持久方式的 mla_decode_fwd。

注意事项：Issue 后续把结论修正为：vLLM 侧未复现，原始复现是直接用 aiter 的 op_tests/test_mla.py 调用 mla_decode_fwd 且不传 work_meta_data，属于 aiter 层问题。此外维护者在持久内核上也观测到 qlen=2 的静默错误行（10/10 复现），以及 qlen 4 + LSE（DCP）偶发失败（1/10），这些与本题非持久 qlen 3–4 问题不是同一路径，不能由同一修复覆盖。

问题场景

在 gfx950 上使用 ROCm 版 vLLM 的 ROCM_AITER_MLA 后端跑 MLA 模型的 decode/MTP 验证阶段时，aiter 的 mla_decode_fwd 会在特定 shape 上分派到非持久内核并输出错误结果。原始复现不是通过 vLLM 端到端完成，而是直接运行 aiter 自带的 op_tests/test_mla.py（仅改成打印 cos_diff 而不做断言），以 fp8 dtype、fp8 KV、-c 8192 -b 1、-n 16,3 / 16,4 驱动 mla_decode_fwd，并且不提供 work_meta_data。vLLM 维护者用离线引擎（DeepSeek-V2-Lite、TP1 = 16 heads/rank、fp8 KV、spec depth 3）统计到 894 次多 token decode 调用，全部带 persistent metadata，未触发该内核。

报错原文

[Bug][ROCm]: ROCM_AITER_MLA can select an aiter MLA decode kernel that silently returns uncorrelated output (gqa_ratio=16, fp8 Q+KV, decode_qlen 3-4, non-persistent)
>>>COS 0.840287
>>>COS 0.897949
mla_a8w8_qh64_qseqlen4_gqaratio16
aiter's own cos_diff is 0.84-0.90 against aiter's own pass bar of 0.03

原因分析

Issue 后续把根因定位为 aiter 中 mla_asm_decode_remap_gfx950 的两条独立分支链相互作用，而不是单一条件缺失：

链 1：gqa_ratio 分派。当 gqa_ratio == 16 且 q_type == "fp8" 时，qlen 3 和 4 落入 else if(max_seqlen_q <= 4) 分支，该分支既没有 persistent 条件也没有 arch 条件，直接设置 sub_Q = 64、config_max_seqlen_q = 4。
链 2：gfx950 覆盖。在第一条链结束后另有独立 if 链，处理 gfx950 + fp8 Q/KV + gqa_ratio == 16 + qlen 3/4 的 shape，但要求 persistent；持久路径会被它救回到受支持内核。

结果是：持久调用走链 2，落到受支持内核；非持久调用保留链 1 的 sub_Q = 64，分派到 mla_a8w8_qh64_qseqlen4_gqaratio16（名字中缺 _ps，即非持久变体），实测 cos_diff 0.84–0.90。同样的组合在 gqa_ratio=32 上会被已有检查显式 reject，所以表现为“32 报错、16 静默算错”。

Issue 评论中的待查疑点：维护者指出在 kv_split == 1 时同一内核测得结果正常（5.4e-4），因此作者提出的最小修复（给 gqa_ratio=16 的 qlen 3–4 分支补上 && persistent && arch_id == "gfx950"）会把该单 split 情形也一并 reject；建议改为限制 kv_split、保留单 split 路径，但这一点尚未有定论。

环境排查

确认 GPU 架构是否为 gfx950（Issue 中为 MI355X）。
确认 aiter 版本（Issue 验证环境为 0.1.28.post1）。
确认 vLLM 版本/分支（Issue 基于 main）。
确认运行的是 MLA 模型，且 _decode_num_heads 为 16（例如 128 头 MLA 模型在 TP8 下为 16 头/rank）。
确认 KV cache dtype 为 fp8、Q 为 fp8。
确认解码长度：decode_qlen 是否为 3 或 4（MTP/spec-decode verify 阶段）。
确认调用路径是否带 persistent metadata（has_persistent_metadata / work_meta_data）：这是本题非持久路径的关键区分点。
若走 vLLM，确认是否在 ROCm/aiter 层直接绕过 vLLM 调用 mla_decode_fwd。

解决步骤

先确认问题入口：检查你的复现是经由 vLLM 正常后端，还是直接调用 aiter 的 mla_decode_fwd。Issue 中原始复现属于后者（未传 work_meta_data），且在 vLLM main 上未能端到端复现。
若经由 vLLM：在 aiter 的 mla_decode_fwd 上加日志，确认 decode_qlen 为 3–4 时是否始终传入 persistent metadata。维护者的离线引擎统计显示该 shape 下 894 次调用全部为 persistent，此时不会分派到 qh64 非持久内核。
若确认为非持久入口触发：关注 aiter 侧的分派修复，而不是在 vLLM 侧改配置。作者提出的候选改动（可优先尝试、尚未合并验证）是给链 1 的 else if(max_seqlen_q <= 4) 补上 && persistent && arch_id == "gfx950"，使非持久 qlen 3–4 走已有 reject(...) 而不是静默算错；同时该分支的 reject 文案（qo_len <= 4）需相应改写。注意此改动对 kv_split == 1 会有副作用，维护者倾向于改为限制 kv_split。
单独核查其他已知问题，不要与本题混为一谈：持久内核上 qlen=2（qh16_qseqlen2_gqaratio16_ps）10/10 出现错误行，以及 qlen 4 + LSE（..._lse_v3_ps，仅 DCP）1/10 失败，相关项见 ROCm/aiter#7007、ROCm/aiter#7621。

验证方法

用 aiter 的 op_tests/test_mla.py 复测并观察 cos_diff：对 gqa_ratio=16、fp8 Q/KV，qlen 1/2 应约为 5e-4 与 1e-3 量级，qlen 3/4 在非持久路径下不应再出现 0.84–0.90 的跳变。同时确认实际分派的内核名：正常持久路径应为带 _ps 的变体（如 mla_a8w8_qh16_qseqlen4_gqaratio16_v3_ps），若仍出现 mla_a8w8_qh64_qseqlen4_gqaratio16 则说明非持久路径未被拦住。通过 vLLM 验证时，应确认 decode_qlen 为 3–4 的调用都带 persistent metadata。