[Bugfix] Add device assertion to TorchSDPA (#5402)

This commit is contained in:
Li, Jiang authored and GitHub committed 2024-06-12 12:58:53 -07:00
1 parent 1a8bfd92d5
commit c3c2903e72
1 file changed
+3
+3
View File
@@ -58,6 +58,9 @@ def get_attn_backend(
ROCmFlashAttentionBackend)
return ROCmFlashAttentionBackend
elif backend == _Backend.TORCH_SDPA:
# TODO: make XPU backend available here.
assert is_cpu(), RuntimeError(
"Torch SDPA backend is only used for the CPU device.")
logger.info("Using Torch SDPA backend.")
from vllm.attention.backends.torch_sdpa import TorchSDPABackend
return TorchSDPABackend