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

c3c2903e · Li, Jiang · GitHub · 1a8bfd92 · c3c2903e
Unverified Commit c3c2903e authored Jun 13, 2024 by Li, Jiang Committed by GitHub Jun 12, 2024
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

vllm/attention/selector.py vllm/attention/selector.py +3 -0

No files found.
--- a/vllm/attention/selector.py
+++ b/vllm/attention/selector.py
@@ -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