[Bug]: [ROCm] Qwen4Exp (AMD path): PLE n-gram IDs read a stale query_start_loc tail during PIECEWISE CUDA-graph replay

快速结论：该现象通常出现在 AMD ROCm 后端运行 Qwen4Exp（Qwen3.8-Flash-Next）并使用 PIECEWISE CUDA graph 的场景；当短批次重放时，PLE 的 n-gram ID 会读取到 query_start_loc 中上一次留下的旧尾部数据。应优先确认是否为该竞态，并检查是否有包含对应修复的 PR。

适用环境：Issue 中确认的环境为：vLLM ROCm build（基于 0.29），2× AMD Radeon AI PRO R9700（gfx1201），ROCm 10，Qwen3.8-Flash-Next（量化），MTP speculative decoding，CUDA graphs 使用 FULL_AND_PIECEWISE。Python、PyTorch 版本未在 Issue 中给出。

最快修复方案：暂无确认的一步修复方案。Issue 中报告者提到维护者已提交 PR #62675，可优先尝试该 PR；但该 PR 的具体改动内容与验证结果未在本 Issue 讨论中说明。

注意事项：该问题只在特定批次形态下出现，长 prompt 和单请求 decode 的输出可能完全正常，因此不能靠“能跑通”判断已修复。Issue 中的修复建议属于代码层面改动，不要直接照抄到生产环境；替代方案（把 AMD n-gram ID 算子注册为 splitting op）同样未在本 Issue 中验证。

问题场景

用户在 AMD ROCm 后端用 vLLM 部署 Qwen3.8-Flash-Next（Qwen4Exp 的 AMD 路径），开启 mamba-aligned prefix caching、MTP speculative decoding，并启用 FULL_AND_PIECEWISE CUDA graphs。当出现 token 数较少且请求数不均的批次时触发：例如短 prompt、prefill 的短尾块、prefix-caching 切分后的小尾步、小的混合 prefill/decode 批次。症状是首个 token 或 logprobs 依赖服务历史——在新启动的服务上跑一组探测，再发生一次 8 路并发突发后重跑同一组探测，10 个探测全部不一致。

报错原文

[Bug]: [ROCm] Qwen4Exp (AMD path): PLE n-gram IDs read a stale query_start_loc tail during PIECEWISE CUDA-graph replay

原因分析

问题出在 PIECEWISE CUDA graph 的重放机制与 PLE n-gram ID 生成之间的配合：

PIECEWISE graph 只按补齐后的 token 总数分发，捕获时使用 dummy 的 num_reqs = min(num_tokens, max_num_seqs)。因此当重放的请求数少于捕获时的请求数时，kernel 仍会读取 ple_query_start_loc[0:R_cap + 1] 和 ngram_context[0:R_cap]。
Qwen4ExpModelState.prepare_inputs 只刷新 ple_query_start_loc[:num_reqs_padded + 1]，而 PIECEWISE 下 num_reqs_padded == num_reqs，剩余条目保留之前步骤的旧值。
当这段旧尾部不单调，或小于当前步的 token 数时，真实 token 会被错误映射到旧的请求行，从而得到错误的 n-gram ID。例如 V2 warmup_kernels 之后缓冲区为 [0, 1, 2, 12, 16, ..., 32]，一个 8 token 的单请求步会看到 [0, 8, 2, 12, ...]，于是 token 2 到 7 被映射到第 2 行。
FULL decode graph 不受影响，因为捕获时 num_reqs == num_reqs_padded，正是被刷新的范围。
对照之下，NVIDIA 路径是故意把依赖请求数的 ID 生成放在 PIECEWISE graph 之外（作为 splitting op vllm::qwen4_exp_compute_ple_ngram_ids）。

报告者同时给出最小修复建议：在 prepare_inputs 的 [:num_reqs_padded + 1] 拷贝之后，将 ple_query_start_loc[num_reqs_padded + 1:] 填充为 input_batch.num_tokens（与 runner 自身对 query_start_loc 的非递减填充一致）；在 _prepare_ngram_context 中把 ngram_context 的每一行填成 EOS。另一种方案是把 AMD 的 n-gram ID 算子注册为 splitting op。

环境排查

确认 vLLM 是否为 ROCm build，以及是否基于 0.29；Issue 中核对过 vllm/models/qwen4_exp/amd/model_state.py（sha1 81e31496）与 ple_layer.py 对应版本。
确认 GPU 为 AMD Radeon AI PRO R9700（gfx1201），可用 rocminfo | grep gfx 核对架构。
确认 ROCm 版本（Issue 为 ROCm 10）。
确认 CUDA graphs 模式是否为 FULL_AND_PIECEWISE；仅 FULL 模式不会触发此问题。
确认是否启用量化模型、MTP speculative decoding、prefix caching、tensor-parallel（Issue 测试覆盖 TP=1 与 TP=2）。
确认是否包含 PR #62675 或对应的 model_state.py 修复；Issue 正文未提供 Python、PyTorch 及安装方式，排查时不要假设这些版本。

解决步骤

先在测试服务上复现：短 prompt（3–32 token）与短尾 prefill 块各跑一次，记录首 token 与 logprobs；然后发起一次 8 路并发请求突发，再重跑同一组探测。
若探测结果在突发前后不一致，即可确认命中该竞态，而不是输出随机性。
可优先尝试 Issue 中维护者提供的 PR #62675（合并到对应 vLLM 版本或自行构建）。
若需自行改代码（仅限验证环境），按 Issue 建议在 Qwen4ExpModelState.prepare_inputs 中把 ple_query_start_loc[num_reqs_padded + 1:] 填充为 input_batch.num_tokens，并在 _prepare_ngram_context 中把 ngram_context 所有行填为 EOS；或改用 NVIDIA 路径的 splitting op 方案。
修改后重跑步骤 1 的探测，并复测长 prompt 与单请求 decode，确认这些场景输出与修改前逐位一致。

验证方法

在新启动的服务上，先跑一遍短 prompt / 短尾块探测并记录首 token 与 logprobs；发起一次 8 路并发突发后，再跑同一组探测。修复后这组探测应在突发前后结果一致（报告者在 CPU 模型中 46 个探测步的 94 处真实 token ID 不匹配降为 0；GPU 上 10 个探测从 0/10 一致变为 10/10 一致）。同时确认长 prompt 和单请求 decode 的对比输出与修复前逐位相同，以排除修复引入的回归。