[ROCm] rocm_unquantized_gemm crashes on CPU tensors (dispatch ignores tensor device)

快速结论：该报错通常出现在 ROCm 构建的 vLLM 上，代码路径却把 CPU 张量派发给了只注册在 CUDA/HIP 后端的算子。[ROCm] rocm_unquantized_gemm crashes on CPU tensors (dispatch ignores tensor device) 的核心问题是 dispatch_unquantized_gemm() 只看 current_platform.is_rocm()，没有看输入张量实际在哪个设备上。

适用环境：Issue 中确认的环境是 vLLM 的 ROCm 构建，涉及 ROCm CI job amd-mi355-dpx-language-models-extended-generation（标记为 optional: true），触发用例为 tests/models/language/generation/test_gemma.py::test_gemma4_dspark_loads_confidence_head。Issue 未提供具体的 Python、PyTorch、ROCm 版本号，这里不做补充。

最快修复方案：Issue 中唯一指向的修复是 PR #58923：在 ROCm 分支里增加设备判断，当输入张量不在 CUDA 上时回退到 default_unquantized_gemm(...)，与 CUDA/CPU 分支已有的回退行为保持一致。该 PR 是否已合入你的版本，需要按实际发布情况确认。

注意事项：不要只靠改测试或跳过用例来绕过，问题根因在算子派发逻辑；在未包含该修复的版本上，任何在 ROCm 构建中把 CPU 张量喂给 rocm_unquantized_gemm 的路径都可能复现。

问题场景

在 ROCm 构建的 vLLM 上运行 Gemma 相关测试时触发。具体是执行 tests/models/language/generation/test_gemma.py::test_gemma4_dspark_loads_confidence_head，该用例构造 Gemma4DSparkForCausalLM，并直接以普通 CPU 张量调用 compute_confidence(hidden, markov)，没有先 .cuda()。这条路径会经过 DSparkConfidenceHead.proj（一个 ReplicatedLinear），在 ROCm 构建上被派发到 rocm_unquantized_gemm，从而崩溃。该用例由 #57263 引入。

报错原文

NotImplementedError: Could not run 'vllm::rocm_unquantized_gemm' with arguments
from the 'CPU' backend. This could be because the operator doesn't exist for
this backend, or was omitted during the selective/custom build process...
'vllm::rocm_unquantized_gemm' is only available for these backends: [CUDA, ...]

原因分析

根因在 vllm/model_executor/layers/utils.py 的 dispatch_unquantized_gemm()：它仅依据 current_platform.is_rocm() 就返回 rocm_unquantized_gemm，完全不参考输入张量的实际设备。而 rocm_unquantized_gemm 这个自定义 op 只注册在 CUDA/HIP dispatch key 上，因此任何 CPU 张量进入都会抛 NotImplementedError。相比之下，CUDA 和 CPU 分支已有回退能力（default_unquantized_gemm / cpu_unquantized_gemm 最终都走 torch.nn.functional.linear），只有 ROCm 分支缺少这种处理。

该问题此前未被发现，是因为这个 ROCm CI job 本身 optional: true，在 #57263 提交时并没有运行。

环境排查

确认当前 vLLM 是否为 ROCm 构建，以及所用版本是否已包含 PR #58923 的修复。
确认触发用例或脚本传入的 hidden、markov 等张量是否位于 CPU；compute_confidence 调用前是否需要将其移到 GPU。
确认对应 ROCm CI job（amd-mi355-dpx-language-models-extended-generation）是否为可选任务、是否被跳过，避免因此漏测。
确认 dispatch_unquantized_gemm() 所在文件是否仍按平台判断、未加入设备检查。
Issue 未给出 Python、PyTorch、ROCm、显卡型号等版本信息，这些项目需按你自己环境另行确认。

解决步骤

先定位 vllm/model_executor/layers/utils.py 中的 dispatch_unquantized_gemm()，确认 ROCm 分支是否只判断 current_platform.is_rocm()。
参考 PR #58923 的修法：在返回 rocm_unquantized_gemm 前检查输入张量设备，if not x.is_cuda: return default_unquantized_gemm(...)，使其与 CUDA/CPU 分支的回退行为一致。
若你的版本尚未包含该修复，可优先尝试升级到已合入 #58923 的 vLLM 版本；无法升级时再考虑按上述逻辑本地打补丁（属推测性操作，需自行验证）。
修复后重新运行 tests/models/language/generation/test_gemma.py::test_gemma4_dspark_loads_confidence_head，确认不再触发 NotImplementedError。
让原本跳过的 ROCm CI job 实际跑一次，避免同类“按平台判断、不看设备”的派发问题再次漏检。

验证方法

在 ROCm 构建上重新执行 test_gemma4_dspark_loads_confidence_head，确认用例通过且不再出现 Could not run 'vllm::rocm_unquantized_gemm' with arguments from the 'CPU' backend。同时确认 CPU 张量路径确实回退到了 default_unquantized_gemm/torch.nn.functional.linear，而非继续命中 ROCm 专用算子。