Skip to content

GitLab

Unverified Commit c3c2903e authored by Li, Jiang's avatar Li, Jiang Committed by GitHub
Browse files

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

parent 1a8bfd92
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 0 deletions
vllm/attention/selector.py
View file @ c3c2903e
...@@ -58,6 +58,9 @@ def get_attn_backend( ...@@ -58,6 +58,9 @@ def get_attn_backend(
ROCmFlashAttentionBackend) ROCmFlashAttentionBackend)
return ROCmFlashAttentionBackend return ROCmFlashAttentionBackend
elif backend == _Backend.TORCH_SDPA: 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.") logger.info("Using Torch SDPA backend.")
from vllm.attention.backends.torch_sdpa import TorchSDPABackend from vllm.attention.backends.torch_sdpa import TorchSDPABackend
return TorchSDPABackend return TorchSDPABackend
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment