[ROCm][Bugfix] Fall back to default GEMM for CPU tensors on ROCm builds (#58923)
Signed-off-by: fai <fangzhouai@gmail.com>
Fangzhou Ai committed
Sep 27, 2026 at 14:47 UTC
187c81b32d2233b8caccc0eabc2519301e99e955
1 file changed
+7
vllm/model_executor/layers/utils.py
+7
@@ -377,6 +377,13 @@ def rocm_unquantized_gemm(
377
weight: torch.Tensor,
378
bias: torch.Tensor | None = None,
379
) -> torch.Tensor:
380
+ # dispatch_unquantized_gemm() picks this backend for any ROCm build,
381
+ # regardless of the tensor's actual device; the custom op below is only
382
+ # registered for the CUDA/HIP dispatch key, so route CPU tensors (e.g.
383
+ # from unit tests that build layers without moving them to the GPU)
384
+ # through the device-agnostic fallback instead of hard-crashing.
385
+ if not x.is_cuda:
386
+ return default_unquantized_gemm(layer, x, weight, bias)
387
return torch.ops.vllm.rocm_unquantized_gemm(x, weight, bias)
388
389