[Bug][ROCm]: ROCM_AITER_UNIFIED_ATTN gives wrong output on gfx950 when a layer’s KV cache is 2–4 GiB

快速结论：在 gfx950（MI355X）上启用 ROCM_AITER_UNIFIED_ATTN 且每层 KV cache 落在 2–4 GiB 区间时，aiter 的 unified attention kernel 会因 32 位偏移溢出而静默返回错误注意力结果，表现为输出乱码、精度指标骤降或 EAGLE3 接受长度变低。优先排查单层 KV cache 字节数与 attention backend 组合。

适用环境：gfx950（MI355X）；ROCm 镜像 vllm/vllm-openai-rocm:nightly-rocm100-0cbac6cd1（vLLM 0.30.1rc1.dev558，aiter 0.1.23）；vLLM 当前锁定的 aiter v0.1.25、v0.1.26 含相同代码；受影响模型包括 Qwen3-32B、Qwen3-8B 以及复现用的 Qwen3.8-27B-FP8；Qwen3 形状为 8 KV heads、head size 128、bf16，后端 block size 64。

最快修复方案：升级到包含 aiter v0.1.27 的 vLLM 版本（#63774 已将 AITER_BRANCH 提升到 v0.1.27，其中包含 ROCm/aiter#6661 的门控修复）。升级前可临时改用 --attention-backend ROCM_ATTN 或 TRITON_ATTN。

注意事项：用 ROCM_AITER_FA 替代并不能完全规避：它只在普通全注意力 decode 时绕开该 kernel，遇到 sliding-window 层、多 token verify（spec decode）、sinks 以及 head_dim < 64 时仍会调用同一个 unified_attention。临时自动选择 PR #62813 已关闭，不作为正式方案。

问题场景

在 gfx950 上运行 vLLM，使用 ROCM_AITER_UNIFIED_ATTN（main 上 gfx950 对 KV connector 运行会自动选中）推理 Qwen3-32B、Qwen3-8B 或 Qwen3.8-27B-FP8。当开启默认的 prefix caching、请求分配到的 block id ≥ 8192（即请求 KV 落到 2 GiB 之后），注意力输出与上下文无关。低并发短路请求或小 KV pool 不易触发，因此常规冒烟测试可能发现不了。

报错原文

[Bug][ROCm]: ROCM_AITER_UNIFIED_ATTN gives wrong output on gfx950 when a layer's KV cache is 2–4 GiB
No error is raised.
’    CORPORATIONtoa  Equation Impossible藁 Royals p犊/bower…

原因分析

已确认根因：gfx950 上 aiter 的 unified_attention 走 _unified_attention_gfx950，在 k.nelement() * k.element_size() <= 2**31 - 1 时使用 32 位 buffer 偏移。vLLM 把 K 和 V 作为同一块 K|V 打包 cache 的两个 strided view 传入：每个 view 元素数只有整块 cache 的一半，但地址范围覆盖整块 cache。于是 2–4 GiB 的 layer KV cache 能通过该 32 位检查，实际超过 2 GiB 的 block 偏移发生溢出，读写到错误位置。按 Qwen3 形状（8 KV heads、head size 128、bf16）和 block size 64 计算，所有 block id ≥ 8192 都受影响。超过 4 GiB 时检查不通过，kernel 反而“自愈”，这正是问题只在 2–4 GiB 窗口出现的原因。

环境排查

确认 GPU 为 gfx950 / MI355X，记录 rocminfo 中的 gfx 目标。
确认 vLLM 版本与 aiter 版本；aiter 需为 v0.1.27（含 ROCm/aiter#6661）才带修复，v0.1.25、v0.1.26 仍受影响。
确认当前 attention backend：ROCM_AITER_UNIFIED_ATTN 为受影响路径；ROCM_AITER_FA 部分场景受影响；ROCM_ATTN、TRITON_ATTN 正常。
估算每层 KV cache 字节数：Qwen3 每 token 每层 K 与 V 合计 4 KiB。默认配置（prefix caching 开、--gpu-memory-utilization 0.85、--max-model-len 131072 + YaRN）下 KV pool 约 700k token，约 2.7 GiB/层，落在窗口内。
确认请求是否实际用到 2 GiB 之后的 block：开启 prefix caching 时分配器会遍历整个 KV pool，总量超过约 2 GiB/层后后续请求即开始命中高危 block。

解决步骤

升级到已提升 AITER_BRANCH 到 v0.1.27 的 vLLM 版本，确认镜像内 aiter 版本为 v0.1.27 或更高，其中包含 ROCm/aiter#6661 的修复。这是 Issue 中确认关闭问题的方案。
若暂不能升级，启动服务时显式指定 --attention-backend ROCM_ATTN（或 TRITON_ATTN）作为临时规避。
不要以 ROCM_AITER_FA 作为规避手段：sliding-window 层、spec decode 的 verify 阶段、sinks 以及 head_dim < 64 仍会进入同一 kernel。
复现或验证窗口时，可按 Issue 方法用小 utilization 钉住形状，例如 --gpu-memory-utilization 0.35 --max-num-seqs 512，使每层 KV 约 3.9 GiB、共 16 层；utilization 过低或过高都会落到窗口之外。

验证方法

在同一 backend 下跑一组能让分配器越过 2 GiB 的评测，观察指标是否回到基线。Issue 中的基线：Qwen3-8B 最坏情况下 TRITON_ATTN 的 GSM8K 为 0.859，受影响 unified attention 仅 0.002，修复后 0.863；MRCR 分别为 0.538、0.178、0.523。Qwen3-32B 默认配置下，MRCR 修复前后为 0.327 → 0.502，第二轮 GSM8K 为 0.633 → 0.733；EAGLE3 接受长度（1k 输入、并发 64）为 1.64–1.72 → 2.30。也可直接对比 kernel 输出与 fp32 reference，修复前最大误差 0.49，TRITON_ATTN 为 0.009。