[Bug]: AMD MI350X/MI355X Kimi K2.5 MXFP4 decode ITL regression v0.18 → v0.19+ (distinct from #43153)

问题场景

用户在 AMD MI350X / MI355X（gfx950）上，以 TP=4、concurrency=16 运行 amd/Kimi-K2.5-MXFP4，负载为长前缀 decode 为主（约 40k 共享前缀 + 可变输入，50k prefill / 1k decode 形态，开启 prefix caching）。从 vLLM ROCm v0.18.0 升级到 v0.19.1 后，decode TPOT 从 19.9 ms 恶化到 29.2 ms，output tok/s 从 714 降到 496；到 v0.22.0 仍为 31.6 ms / 466 tok/s，decode 性能没有恢复。该问题被明确标记为与 #43153（INT4 / Triton MoE kernel 路径）不同。

报错原文

[Bug]:  AMD MI350X/MI355X Kimi K2.5 MXFP4 decode ITL regression v0.18 → v0.19+ (distinct from #43153)

v0.18.0  decode TPOT p50: 19.9 ms   output tok/s: 714   baseline
v0.19.1  decode TPOT p50: 29.2 ms   output tok/s: 496   +47% latency, −30% throughput
v0.22.0  decode TPOT p50: 31.6 ms   output tok/s: 466   +58% latency, −35% throughput

v0.22 does not recover decode.

原因分析

根因是 vLLM 的 ROCM_AITER_MLA 后端在 v0.18 之后切换了 AITER MLA decode kernel：v0.18 中 _rocm_aiter_mla_decode_fwd_impl 调用 aiter.mla.mla_decode_fwd(...) 时不带 persistent metadata，走非持久化 kernel mla_dec_stage1；v0.21+（含 v0.24）改为构建 persistent metadata（get_mla_metadata_v1）并传入 work_meta_data，走持久化 kernel mla_a16w16_qh16。两个镜像自带的 AITER 都支持持久化 API，因此切换由 vLLM 后端决定，而非 AITER 可用性问题。

在单一 v0.24.0 镜像内（ROCm/torch/AITER 保持一致）微基准显示，持久化 kernel 在大总上下文下最多比非持久化 kernel 慢约 1.48×；总上下文达到约 800k token 后差距明显，而低总上下文时持久化 kernel 更快，交叉点约在 200k–400k token。因此该回退只在长前缀、decode 密集、高并发场景暴露，此前不易被发现。在 v0.24 镜像中强制切回 mla_dec_stage1 后，output tok/s 有约 20% 提升，进一步支持该结论。

环境排查

• 确认 GPU 型号与架构：应为 MI350X / MI355X，即 gfx950。
• 确认 vLLM 版本：v0.18.0 正常，v0.19.1 与 v0.22.0 出现回退；Issue 中 v0.21+ 同样走持久化 kernel。
• 确认 PyTorch / ROCm / triton：Issue 环境为 PyTorch 2.10.0、ROCm 7.2.1、triton 3.6.0、vLLM 0.19.1；对照版本 v0.18.0 为 ROCm 7.0.0、torch 2.9.1、triton 3.4.0。
• 确认 AITER 相关环境变量：VLLM_ROCM_USE_AITER=1、VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4、VLLM_ROCM_USE_AITER_RMSNORM=0。
• 确认使用的 MLA 后端是否为 ROCM_AITER_MLA，以及实际调用的 decode kernel 名称。
• 确认测试负载的总上下文与并发度：Issue 中回退在 batch × ctx 达到较大总上下文（如 8×50k 以上）时明显，低负载可能不暴露。
• 确认测量条件：warm median、至少 3 次重复、profiler 关闭。

解决步骤

1. 先排除准确率问题：用 Issue 中的 lm_eval 命令检查 v0.18.0 / v0.19.1 / v0.22.0 的 gsm8k 准确率。Issue 结果为 0.9204 / 0.9280 / 0.9371，没有明显差异，说明不是“计算垃圾结果反而快”之类的准确率回归。
2. 在复现镜像中确认 MLA decode 实际走的 kernel：检查 ROCM_AITER_MLA 后端是否调用 get_mla_metadata_v1 并传入 work_meta_data，若是则走持久化 kernel mla_a16w16_qh16。
3. 用 Issue 中的微基准方法，在同一镜像内分别调用 mla_dec_stage1（不带持久化 metadata）和 mla_a16w16_qh16（带持久化 metadata），参数为 bf16 q/kv、16 query heads、kv_lora_rank=512、qk_rope=64、v=512、page_size=1、decode_qlen=1，取 50 次中位数。
4. 按 batch × ctx 覆盖测试总上下文：1×50k、4×50k、8×50k、16×50k、64×50k、128×50k、256×50k、256×1k。确认回退是否在总上下文约 800k token 以上、且随并发增加而加剧。
5. 可优先尝试：在 v0.24 镜像中强制使用 mla_dec_stage1 非持久化 kernel 做端到端验证。Issue 中这样做的运行结果为 output tok/s 从 306.4 提升到 358.8，ITL p50 从 31.3 ms 降到 24.1 ms。注意这是绕过修改，不是官方修复。
6. 跟踪 AITER kernel 选择启发式的更新：作者已在 ROCm/aiter#4144 提交 PR，目标是低负载用 kernel B、高总上下文负载用 kernel A。合入后再按官方版本验证，不要长期依赖手工替换 kernel。

验证方法

在相同硬件、TP=4、concurrency=16、相同长前缀 decode 负载和相同 warm-up 条件下，比较切换前后的 decode TPOT p50、ITL p50、output tok/s。若强制使用 mla_dec_stage1 后 output tok/s 明显回升（Issue 中约为 +20%）、TPOT / ITL 下降，说明回退与持久化 MLA decode kernel 选择一致。同时保持 gsm8k 等准确率检查在噪声范围内，确认性能变化不是牺牲精度换来的。

参考来源

vllm-project/vllm #45943

ROCm/aiter PR #4144

AI 工具推荐

想把多个 AI 模型放在一个入口？

GamsGo AI 集成 ChatGPT、DeepSeek、Gemini、Claude、Midjourney、Veo 等常用模型，适合写作、绘图、视频和日常 AI 工作流。

了解 GamsGo AI

推广链接：通过此链接购买，我可能获得佣金，不影响你的价格。