We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Bug]: TorchSDPAMetadata is out of date
### Your current environment

alexr@roxy:~$ python collect_env.py
Collecting environment information...
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.26.3
Libc version: glibc-2.35

Python version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.5.119
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3070
Nvidia driver version: 535.161.08
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      39 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             16
On-line CPU(s) list:                0-15
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Core(TM) i7-10700K CPU @ 3.80GHz
CPU family:                         6
Model:                              165
Thread(s) per core:                 2
Core(s) per socket:                 8
Socket(s):                          1
Stepping:                           5
CPU max MHz:                        5100.0000
CPU min MHz:                        800.0000
BogoMIPS:                           7599.80
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts pku ospke md_clear flush_l1d arch_capabilities
L1d cache:                          256 KiB (8 instances)
L1i cache:                          256 KiB (8 instances)
L2 cache:                           2 MiB (8 instances)
L3 cache:                           16 MiB (1 instance)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-15
Vulnerability Gather data sampling: Mitigation; Microcode
Vulnerability Itlb multihit:        KVM: Mitigation: VMX unsupported
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence
Vulnerability Srbds:                Mitigation; Microcode
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.2
[pip3] nvidia-nccl-cu11==2.14.3
[pip3] sentence-transformers==2.2.2
[pip3] torch==2.0.0
[pip3] torchvision==0.15.1
[pip3] transformers==4.26.0
[pip3] triton==2.0.0
[conda] Could not collect
ROCM Version: Could not collect
Neuron SDK Version: N/A
vLLM Version: N/A
vLLM Build Flags:
CUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled
GPU Topology:
GPU0	CPU Affinity	NUMA Affinity	GPU NUMA ID
GPU0	 X 	0-15	0		N/A

Legend:

  X    = Self
  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)
  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node
  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)
  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)
  PIX  = Connection traversing at most a single PCIe bridge
  NV#  = Connection traversing a bonded set of # NVLinks


### 🐛 Describe the bug

I believe the TorchSDPA backend for attention is no longer compatible. I received a crash message that said as much and did some debugging. If you set `export VLLM_ATTENTION_BACKEND=TORCH_SDPA`, then the metadata is set with
```
                attn_metadata = self.attn_backend.make_metadata(
                    num_prefills=0,
                    num_prefill_tokens=0,
                    num_decode_tokens=batch_size,
                    slot_mapping=slot_mapping[:batch_size],
                    seq_lens=None,
                    seq_lens_tensor=seq_lens[:batch_size],
                    max_query_len=None,
                    max_prefill_seq_len=0,
                    max_decode_seq_len=self.max_seq_len_to_capture,
                    query_start_loc=None,
                    seq_start_loc=None,
                    context_lens_tensor=None,
                    block_tables=block_tables[:batch_size],
                    use_cuda_graph=True,
                )
```

which is sent to 

```
    @staticmethod
    def make_metadata(*args, **kwargs) -> "TorchSDPAMetadata":
        return TorchSDPAMetadata(*args, **kwargs)
```

and neither TorchSDPAMetadata (nor its superclasses) have most of those values
```
@dataclass
class TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):
    """Metadata for TorchSDPABackend.
    """
    # Currently, input sequences can only contain all prompts
    # or all decoding. True if all sequences are prompts.
    is_prompt: bool
    slot_mapping: torch.Tensor
    seq_lens: Optional[List[int]]
    # Maximum query length in the batch. None for decoding.
    max_query_len: Optional[int]
```

Am I correct in my diagnosis? I can send a crash log if necessary, but I believe this is pretty self evidently an error (and a CI test should be added or the backend should be removed probably.)

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `vllm/attention/selector.py`
Content:
```
1 import enum
2 from functools import lru_cache
3 from typing import Optional, Type
4 
5 import torch
6 
7 import vllm.envs as envs
8 from vllm.attention.backends.abstract import AttentionBackend
9 from vllm.logger import init_logger
10 from vllm.utils import is_cpu, is_hip
11 
12 logger = init_logger(__name__)
13 
14 
15 class _Backend(enum.Enum):
16     FLASH_ATTN = enum.auto()
17     XFORMERS = enum.auto()
18     ROCM_FLASH = enum.auto()
19     TORCH_SDPA = enum.auto()
20     FLASHINFER = enum.auto()
21 
22 
23 @lru_cache(maxsize=None)
24 def get_attn_backend(
25     num_heads: int,
26     head_size: int,
27     num_kv_heads: int,
28     sliding_window: Optional[int],
29     dtype: torch.dtype,
30     kv_cache_dtype: Optional[str],
31     block_size: int,
32     is_blocksparse: bool = False,
33 ) -> Type[AttentionBackend]:
34     """Selects which attention backend to use and lazily imports it."""
35 
36     if is_blocksparse:
37         logger.info("Using BlocksparseFlashAttention backend.")
38         from vllm.attention.backends.blocksparse_attn import (
39             BlocksparseFlashAttentionBackend)
40         return BlocksparseFlashAttentionBackend
41 
42     backend = which_attn_to_use(num_heads, head_size, num_kv_heads,
43                                 sliding_window, dtype, kv_cache_dtype,
44                                 block_size)
45     if backend == _Backend.FLASH_ATTN:
46         from vllm.attention.backends.flash_attn import (  # noqa: F401
47             FlashAttentionBackend)
48         return FlashAttentionBackend
49     if backend == _Backend.XFORMERS:
50         logger.info("Using XFormers backend.")
51         from vllm.attention.backends.xformers import (  # noqa: F401
52             XFormersBackend)
53         return XFormersBackend
54     elif backend == _Backend.ROCM_FLASH:
55         logger.info("Using ROCmFlashAttention backend.")
56         from vllm.attention.backends.rocm_flash_attn import (  # noqa: F401
57             ROCmFlashAttentionBackend)
58         return ROCmFlashAttentionBackend
59     elif backend == _Backend.TORCH_SDPA:
60         logger.info("Using Torch SDPA backend.")
61         from vllm.attention.backends.torch_sdpa import TorchSDPABackend
62         return TorchSDPABackend
63     elif backend == _Backend.FLASHINFER:
64         logger.info("Using Flashinfer backend.")
65         logger.warning("Eager mode is required for the Flashinfer backend. "
66                        "Please make sure --enforce-eager is set.")
67         from vllm.attention.backends.flashinfer import FlashInferBackend
68         return FlashInferBackend
69     else:
70         raise ValueError("Invalid attention backend.")
71 
72 
73 def which_attn_to_use(
74     num_heads: int,
75     head_size: int,
76     num_kv_heads: int,
77     sliding_window: Optional[int],
78     dtype: torch.dtype,
79     kv_cache_dtype: Optional[str],
80     block_size: int,
81 ) -> _Backend:
82     """Returns which flash attention backend to use."""
83 
84     # Default case.
85     selected_backend = _Backend.FLASH_ATTN
86 
87     # Check the environment variable and override if specified
88     backend_by_env_var: Optional[str] = envs.VLLM_ATTENTION_BACKEND
89     if backend_by_env_var is not None:
90         backend_members = _Backend.__members__
91         if backend_by_env_var not in backend_members:
92             raise ValueError(
93                 f"Invalid attention backend '{backend_by_env_var}'. "
94                 f"Available backends: {', '.join(backend_members)} "
95                 "(case-sensitive).")
96         selected_backend = _Backend[backend_by_env_var]
97 
98     if is_cpu():
99         if selected_backend != _Backend.TORCH_SDPA:
100             logger.info("Cannot use %s backend on CPU.", selected_backend)
101         return _Backend.TORCH_SDPA
102 
103     if is_hip():
104         # AMD GPUs.
105         selected_backend = (_Backend.ROCM_FLASH if selected_backend
106                             == _Backend.FLASH_ATTN else selected_backend)
107         if selected_backend == _Backend.ROCM_FLASH:
108             if torch.cuda.get_device_capability()[0] != 9:
109                 # not Instinct series GPUs.
110                 logger.info("flash_attn is not supported on NAVI GPUs.")
111         else:
112             logger.info("%s is not supported in AMD GPUs.", selected_backend)
113         return _Backend.ROCM_FLASH
114 
115     # FlashAttn in NVIDIA GPUs.
116     if selected_backend == _Backend.FLASH_ATTN:
117         if torch.cuda.get_device_capability()[0] < 8:
118             # Volta and Turing NVIDIA GPUs.
119             logger.info(
120                 "Cannot use FlashAttention-2 backend for Volta and Turing "
121                 "GPUs.")
122             selected_backend = _Backend.XFORMERS
123         elif dtype not in (torch.float16, torch.bfloat16):
124             logger.info(
125                 "Cannot use FlashAttention-2 backend for dtype other than "
126                 "torch.float16 or torch.bfloat16.")
127             selected_backend = _Backend.XFORMERS
128         elif kv_cache_dtype is not None and kv_cache_dtype.startswith("fp8"):
129             logger.info(
130                 "Cannot use FlashAttention-2 backend for FP8 KV cache.")
131             selected_backend = _Backend.XFORMERS
132         elif block_size % 16 != 0:
133             logger.info(
134                 "Cannot use FlashAttention-2 backend for block size not "
135                 "divisible by 16.")
136             selected_backend = _Backend.XFORMERS
137         elif sliding_window is not None:
138             logger.info(
139                 "Cannot use FlashAttention-2 backend due to sliding window.")
140             selected_backend = _Backend.XFORMERS
141 
142     # FlashAttn is valid for the model, checking if the package is installed.
143     if selected_backend == _Backend.FLASH_ATTN:
144         try:
145             import vllm_flash_attn  # noqa: F401
146 
147             from vllm.attention.backends.flash_attn import (  # noqa: F401
148                 FlashAttentionBackend)
149 
150             supported_sizes = FlashAttentionBackend.get_supported_head_sizes()
151             if head_size not in supported_sizes:
152                 logger.info(
153                     "Cannot use FlashAttention-2 backend for head size %d.",
154                     head_size)
155                 selected_backend = _Backend.XFORMERS
156         except ImportError:
157             logger.info(
158                 "Cannot use FlashAttention-2 backend because the "
159                 "vllm_flash_attn package is not found. "
160                 "`pip install vllm-flash-attn` for better performance.")
161             selected_backend = _Backend.XFORMERS
162 
163     return selected_backend
164
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/vllm/attention/selector.py b/vllm/attention/selector.py
--- a/vllm/attention/selector.py
+++ b/vllm/attention/selector.py
@@ -57,6 +57,9 @@
             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

{"golden_diff": "diff --git a/vllm/attention/selector.py b/vllm/attention/selector.py\n--- a/vllm/attention/selector.py\n+++ b/vllm/attention/selector.py\n@@ -57,6 +57,9 @@\n             ROCmFlashAttentionBackend)\n         return ROCmFlashAttentionBackend\n     elif backend == _Backend.TORCH_SDPA:\n+        # TODO: make XPU backend available here.\n+        assert is_cpu(), RuntimeError(\n+            \"Torch SDPA backend is only used for the CPU device.\")\n         logger.info(\"Using Torch SDPA backend.\")\n         from vllm.attention.backends.torch_sdpa import TorchSDPABackend\n         return TorchSDPABackend\n", "issue": "[Bug]: TorchSDPAMetadata is out of date\n### Your current environment\n\nalexr@roxy:~$ python collect_env.py\r\nCollecting environment information...\r\nPyTorch version: 2.0.0+cu117\r\nIs debug build: False\r\nCUDA used to build PyTorch: 11.7\r\nROCM used to build PyTorch: N/A\r\n\r\nOS: Ubuntu 22.04.1 LTS (x86_64)\r\nGCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0\r\nClang version: Could not collect\r\nCMake version: version 3.26.3\r\nLibc version: glibc-2.35\r\n\r\nPython version: 3.10.12 (main, Nov 20 2023, 15:14:05) [GCC 11.4.0] (64-bit runtime)\r\nPython platform: Linux-5.15.0-105-generic-x86_64-with-glibc2.35\r\nIs CUDA available: True\r\nCUDA runtime version: 11.5.119\r\nCUDA_MODULE_LOADING set to: LAZY\r\nGPU models and configuration: GPU 0: NVIDIA GeForce RTX 3070\r\nNvidia driver version: 535.161.08\r\ncuDNN version: Could not collect\r\nHIP runtime version: N/A\r\nMIOpen runtime version: N/A\r\nIs XNNPACK available: True\r\n\r\nCPU:\r\nArchitecture:                       x86_64\r\nCPU op-mode(s):                     32-bit, 64-bit\r\nAddress sizes:                      39 bits physical, 48 bits virtual\r\nByte Order:                         Little Endian\r\nCPU(s):                             16\r\nOn-line CPU(s) list:                0-15\r\nVendor ID:                          GenuineIntel\r\nModel name:                         Intel(R) Core(TM) i7-10700K CPU @ 3.80GHz\r\nCPU family:                         6\r\nModel:                              165\r\nThread(s) per core:                 2\r\nCore(s) per socket:                 8\r\nSocket(s):                          1\r\nStepping:                           5\r\nCPU max MHz:                        5100.0000\r\nCPU min MHz:                        800.0000\r\nBogoMIPS:                           7599.80\r\nFlags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts pku ospke md_clear flush_l1d arch_capabilities\r\nL1d cache:                          256 KiB (8 instances)\r\nL1i cache:                          256 KiB (8 instances)\r\nL2 cache:                           2 MiB (8 instances)\r\nL3 cache:                           16 MiB (1 instance)\r\nNUMA node(s):                       1\r\nNUMA node0 CPU(s):                  0-15\r\nVulnerability Gather data sampling: Mitigation; Microcode\r\nVulnerability Itlb multihit:        KVM: Mitigation: VMX unsupported\r\nVulnerability L1tf:                 Not affected\r\nVulnerability Mds:                  Not affected\r\nVulnerability Meltdown:             Not affected\r\nVulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT vulnerable\r\nVulnerability Retbleed:             Mitigation; Enhanced IBRS\r\nVulnerability Spec rstack overflow: Not affected\r\nVulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp\r\nVulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization\r\nVulnerability Spectre v2:           Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence\r\nVulnerability Srbds:                Mitigation; Microcode\r\nVulnerability Tsx async abort:      Not affected\r\n\r\nVersions of relevant libraries:\r\n[pip3] flake8==6.0.0\r\n[pip3] mypy-extensions==1.0.0\r\n[pip3] numpy==1.24.2\r\n[pip3] nvidia-nccl-cu11==2.14.3\r\n[pip3] sentence-transformers==2.2.2\r\n[pip3] torch==2.0.0\r\n[pip3] torchvision==0.15.1\r\n[pip3] transformers==4.26.0\r\n[pip3] triton==2.0.0\r\n[conda] Could not collect\r\nROCM Version: Could not collect\r\nNeuron SDK Version: N/A\r\nvLLM Version: N/A\r\nvLLM Build Flags:\r\nCUDA Archs: Not Set; ROCm: Disabled; Neuron: Disabled\r\nGPU Topology:\r\nGPU0\tCPU Affinity\tNUMA Affinity\tGPU NUMA ID\r\nGPU0\t X \t0-15\t0\t\tN/A\r\n\r\nLegend:\r\n\r\n  X    = Self\r\n  SYS  = Connection traversing PCIe as well as the SMP interconnect between NUMA nodes (e.g., QPI/UPI)\r\n  NODE = Connection traversing PCIe as well as the interconnect between PCIe Host Bridges within a NUMA node\r\n  PHB  = Connection traversing PCIe as well as a PCIe Host Bridge (typically the CPU)\r\n  PXB  = Connection traversing multiple PCIe bridges (without traversing the PCIe Host Bridge)\r\n  PIX  = Connection traversing at most a single PCIe bridge\r\n  NV#  = Connection traversing a bonded set of # NVLinks\r\n\n\n### \ud83d\udc1b Describe the bug\n\nI believe the TorchSDPA backend for attention is no longer compatible. I received a crash message that said as much and did some debugging. If you set `export VLLM_ATTENTION_BACKEND=TORCH_SDPA`, then the metadata is set with\r\n```\r\n                attn_metadata = self.attn_backend.make_metadata(\r\n                    num_prefills=0,\r\n                    num_prefill_tokens=0,\r\n                    num_decode_tokens=batch_size,\r\n                    slot_mapping=slot_mapping[:batch_size],\r\n                    seq_lens=None,\r\n                    seq_lens_tensor=seq_lens[:batch_size],\r\n                    max_query_len=None,\r\n                    max_prefill_seq_len=0,\r\n                    max_decode_seq_len=self.max_seq_len_to_capture,\r\n                    query_start_loc=None,\r\n                    seq_start_loc=None,\r\n                    context_lens_tensor=None,\r\n                    block_tables=block_tables[:batch_size],\r\n                    use_cuda_graph=True,\r\n                )\r\n```\r\n\r\nwhich is sent to \r\n\r\n```\r\n    @staticmethod\r\n    def make_metadata(*args, **kwargs) -> \"TorchSDPAMetadata\":\r\n        return TorchSDPAMetadata(*args, **kwargs)\r\n```\r\n\r\nand neither TorchSDPAMetadata (nor its superclasses) have most of those values\r\n```\r\n@dataclass\r\nclass TorchSDPAMetadata(AttentionMetadata, PagedAttentionMetadata):\r\n    \"\"\"Metadata for TorchSDPABackend.\r\n    \"\"\"\r\n    # Currently, input sequences can only contain all prompts\r\n    # or all decoding. True if all sequences are prompts.\r\n    is_prompt: bool\r\n    slot_mapping: torch.Tensor\r\n    seq_lens: Optional[List[int]]\r\n    # Maximum query length in the batch. None for decoding.\r\n    max_query_len: Optional[int]\r\n```\r\n\r\nAm I correct in my diagnosis? I can send a crash log if necessary, but I believe this is pretty self evidently an error (and a CI test should be added or the backend should be removed probably.)\n", "before_files": [{"content": "import enum\nfrom functools import lru_cache\nfrom typing import Optional, Type\n\nimport torch\n\nimport vllm.envs as envs\nfrom vllm.attention.backends.abstract import AttentionBackend\nfrom vllm.logger import init_logger\nfrom vllm.utils import is_cpu, is_hip\n\nlogger = init_logger(__name__)\n\n\nclass _Backend(enum.Enum):\n    FLASH_ATTN = enum.auto()\n    XFORMERS = enum.auto()\n    ROCM_FLASH = enum.auto()\n    TORCH_SDPA = enum.auto()\n    FLASHINFER = enum.auto()\n\n\n@lru_cache(maxsize=None)\ndef get_attn_backend(\n    num_heads: int,\n    head_size: int,\n    num_kv_heads: int,\n    sliding_window: Optional[int],\n    dtype: torch.dtype,\n    kv_cache_dtype: Optional[str],\n    block_size: int,\n    is_blocksparse: bool = False,\n) -> Type[AttentionBackend]:\n    \"\"\"Selects which attention backend to use and lazily imports it.\"\"\"\n\n    if is_blocksparse:\n        logger.info(\"Using BlocksparseFlashAttention backend.\")\n        from vllm.attention.backends.blocksparse_attn import (\n            BlocksparseFlashAttentionBackend)\n        return BlocksparseFlashAttentionBackend\n\n    backend = which_attn_to_use(num_heads, head_size, num_kv_heads,\n                                sliding_window, dtype, kv_cache_dtype,\n                                block_size)\n    if backend == _Backend.FLASH_ATTN:\n        from vllm.attention.backends.flash_attn import (  # noqa: F401\n            FlashAttentionBackend)\n        return FlashAttentionBackend\n    if backend == _Backend.XFORMERS:\n        logger.info(\"Using XFormers backend.\")\n        from vllm.attention.backends.xformers import (  # noqa: F401\n            XFormersBackend)\n        return XFormersBackend\n    elif backend == _Backend.ROCM_FLASH:\n        logger.info(\"Using ROCmFlashAttention backend.\")\n        from vllm.attention.backends.rocm_flash_attn import (  # noqa: F401\n            ROCmFlashAttentionBackend)\n        return ROCmFlashAttentionBackend\n    elif backend == _Backend.TORCH_SDPA:\n        logger.info(\"Using Torch SDPA backend.\")\n        from vllm.attention.backends.torch_sdpa import TorchSDPABackend\n        return TorchSDPABackend\n    elif backend == _Backend.FLASHINFER:\n        logger.info(\"Using Flashinfer backend.\")\n        logger.warning(\"Eager mode is required for the Flashinfer backend. \"\n                       \"Please make sure --enforce-eager is set.\")\n        from vllm.attention.backends.flashinfer import FlashInferBackend\n        return FlashInferBackend\n    else:\n        raise ValueError(\"Invalid attention backend.\")\n\n\ndef which_attn_to_use(\n    num_heads: int,\n    head_size: int,\n    num_kv_heads: int,\n    sliding_window: Optional[int],\n    dtype: torch.dtype,\n    kv_cache_dtype: Optional[str],\n    block_size: int,\n) -> _Backend:\n    \"\"\"Returns which flash attention backend to use.\"\"\"\n\n    # Default case.\n    selected_backend = _Backend.FLASH_ATTN\n\n    # Check the environment variable and override if specified\n    backend_by_env_var: Optional[str] = envs.VLLM_ATTENTION_BACKEND\n    if backend_by_env_var is not None:\n        backend_members = _Backend.__members__\n        if backend_by_env_var not in backend_members:\n            raise ValueError(\n                f\"Invalid attention backend '{backend_by_env_var}'. \"\n                f\"Available backends: {', '.join(backend_members)} \"\n                \"(case-sensitive).\")\n        selected_backend = _Backend[backend_by_env_var]\n\n    if is_cpu():\n        if selected_backend != _Backend.TORCH_SDPA:\n            logger.info(\"Cannot use %s backend on CPU.\", selected_backend)\n        return _Backend.TORCH_SDPA\n\n    if is_hip():\n        # AMD GPUs.\n        selected_backend = (_Backend.ROCM_FLASH if selected_backend\n                            == _Backend.FLASH_ATTN else selected_backend)\n        if selected_backend == _Backend.ROCM_FLASH:\n            if torch.cuda.get_device_capability()[0] != 9:\n                # not Instinct series GPUs.\n                logger.info(\"flash_attn is not supported on NAVI GPUs.\")\n        else:\n            logger.info(\"%s is not supported in AMD GPUs.\", selected_backend)\n        return _Backend.ROCM_FLASH\n\n    # FlashAttn in NVIDIA GPUs.\n    if selected_backend == _Backend.FLASH_ATTN:\n        if torch.cuda.get_device_capability()[0] < 8:\n            # Volta and Turing NVIDIA GPUs.\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for Volta and Turing \"\n                \"GPUs.\")\n            selected_backend = _Backend.XFORMERS\n        elif dtype not in (torch.float16, torch.bfloat16):\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for dtype other than \"\n                \"torch.float16 or torch.bfloat16.\")\n            selected_backend = _Backend.XFORMERS\n        elif kv_cache_dtype is not None and kv_cache_dtype.startswith(\"fp8\"):\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for FP8 KV cache.\")\n            selected_backend = _Backend.XFORMERS\n        elif block_size % 16 != 0:\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for block size not \"\n                \"divisible by 16.\")\n            selected_backend = _Backend.XFORMERS\n        elif sliding_window is not None:\n            logger.info(\n                \"Cannot use FlashAttention-2 backend due to sliding window.\")\n            selected_backend = _Backend.XFORMERS\n\n    # FlashAttn is valid for the model, checking if the package is installed.\n    if selected_backend == _Backend.FLASH_ATTN:\n        try:\n            import vllm_flash_attn  # noqa: F401\n\n            from vllm.attention.backends.flash_attn import (  # noqa: F401\n                FlashAttentionBackend)\n\n            supported_sizes = FlashAttentionBackend.get_supported_head_sizes()\n            if head_size not in supported_sizes:\n                logger.info(\n                    \"Cannot use FlashAttention-2 backend for head size %d.\",\n                    head_size)\n                selected_backend = _Backend.XFORMERS\n        except ImportError:\n            logger.info(\n                \"Cannot use FlashAttention-2 backend because the \"\n                \"vllm_flash_attn package is not found. \"\n                \"`pip install vllm-flash-attn` for better performance.\")\n            selected_backend = _Backend.XFORMERS\n\n    return selected_backend\n", "path": "vllm/attention/selector.py"}], "after_files": [{"content": "import enum\nfrom functools import lru_cache\nfrom typing import Optional, Type\n\nimport torch\n\nimport vllm.envs as envs\nfrom vllm.attention.backends.abstract import AttentionBackend\nfrom vllm.logger import init_logger\nfrom vllm.utils import is_cpu, is_hip\n\nlogger = init_logger(__name__)\n\n\nclass _Backend(enum.Enum):\n    FLASH_ATTN = enum.auto()\n    XFORMERS = enum.auto()\n    ROCM_FLASH = enum.auto()\n    TORCH_SDPA = enum.auto()\n    FLASHINFER = enum.auto()\n\n\n@lru_cache(maxsize=None)\ndef get_attn_backend(\n    num_heads: int,\n    head_size: int,\n    num_kv_heads: int,\n    sliding_window: Optional[int],\n    dtype: torch.dtype,\n    kv_cache_dtype: Optional[str],\n    block_size: int,\n    is_blocksparse: bool = False,\n) -> Type[AttentionBackend]:\n    \"\"\"Selects which attention backend to use and lazily imports it.\"\"\"\n\n    if is_blocksparse:\n        logger.info(\"Using BlocksparseFlashAttention backend.\")\n        from vllm.attention.backends.blocksparse_attn import (\n            BlocksparseFlashAttentionBackend)\n        return BlocksparseFlashAttentionBackend\n\n    backend = which_attn_to_use(num_heads, head_size, num_kv_heads,\n                                sliding_window, dtype, kv_cache_dtype,\n                                block_size)\n    if backend == _Backend.FLASH_ATTN:\n        from vllm.attention.backends.flash_attn import (  # noqa: F401\n            FlashAttentionBackend)\n        return FlashAttentionBackend\n    if backend == _Backend.XFORMERS:\n        logger.info(\"Using XFormers backend.\")\n        from vllm.attention.backends.xformers import (  # noqa: F401\n            XFormersBackend)\n        return XFormersBackend\n    elif backend == _Backend.ROCM_FLASH:\n        logger.info(\"Using ROCmFlashAttention backend.\")\n        from vllm.attention.backends.rocm_flash_attn import (  # noqa: F401\n            ROCmFlashAttentionBackend)\n        return ROCmFlashAttentionBackend\n    elif backend == _Backend.TORCH_SDPA:\n        # TODO: make XPU backend available here.\n        assert is_cpu(), RuntimeError(\n            \"Torch SDPA backend is only used for the CPU device.\")\n        logger.info(\"Using Torch SDPA backend.\")\n        from vllm.attention.backends.torch_sdpa import TorchSDPABackend\n        return TorchSDPABackend\n    elif backend == _Backend.FLASHINFER:\n        logger.info(\"Using Flashinfer backend.\")\n        logger.warning(\"Eager mode is required for the Flashinfer backend. \"\n                       \"Please make sure --enforce-eager is set.\")\n        from vllm.attention.backends.flashinfer import FlashInferBackend\n        return FlashInferBackend\n    else:\n        raise ValueError(\"Invalid attention backend.\")\n\n\ndef which_attn_to_use(\n    num_heads: int,\n    head_size: int,\n    num_kv_heads: int,\n    sliding_window: Optional[int],\n    dtype: torch.dtype,\n    kv_cache_dtype: Optional[str],\n    block_size: int,\n) -> _Backend:\n    \"\"\"Returns which flash attention backend to use.\"\"\"\n\n    # Default case.\n    selected_backend = _Backend.FLASH_ATTN\n\n    # Check the environment variable and override if specified\n    backend_by_env_var: Optional[str] = envs.VLLM_ATTENTION_BACKEND\n    if backend_by_env_var is not None:\n        backend_members = _Backend.__members__\n        if backend_by_env_var not in backend_members:\n            raise ValueError(\n                f\"Invalid attention backend '{backend_by_env_var}'. \"\n                f\"Available backends: {', '.join(backend_members)} \"\n                \"(case-sensitive).\")\n        selected_backend = _Backend[backend_by_env_var]\n\n    if is_cpu():\n        if selected_backend != _Backend.TORCH_SDPA:\n            logger.info(\"Cannot use %s backend on CPU.\", selected_backend)\n        return _Backend.TORCH_SDPA\n\n    if is_hip():\n        # AMD GPUs.\n        selected_backend = (_Backend.ROCM_FLASH if selected_backend\n                            == _Backend.FLASH_ATTN else selected_backend)\n        if selected_backend == _Backend.ROCM_FLASH:\n            if torch.cuda.get_device_capability()[0] != 9:\n                # not Instinct series GPUs.\n                logger.info(\"flash_attn is not supported on NAVI GPUs.\")\n        else:\n            logger.info(\"%s is not supported in AMD GPUs.\", selected_backend)\n        return _Backend.ROCM_FLASH\n\n    # FlashAttn in NVIDIA GPUs.\n    if selected_backend == _Backend.FLASH_ATTN:\n        if torch.cuda.get_device_capability()[0] < 8:\n            # Volta and Turing NVIDIA GPUs.\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for Volta and Turing \"\n                \"GPUs.\")\n            selected_backend = _Backend.XFORMERS\n        elif dtype not in (torch.float16, torch.bfloat16):\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for dtype other than \"\n                \"torch.float16 or torch.bfloat16.\")\n            selected_backend = _Backend.XFORMERS\n        elif kv_cache_dtype is not None and kv_cache_dtype.startswith(\"fp8\"):\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for FP8 KV cache.\")\n            selected_backend = _Backend.XFORMERS\n        elif block_size % 16 != 0:\n            logger.info(\n                \"Cannot use FlashAttention-2 backend for block size not \"\n                \"divisible by 16.\")\n            selected_backend = _Backend.XFORMERS\n        elif sliding_window is not None:\n            logger.info(\n                \"Cannot use FlashAttention-2 backend due to sliding window.\")\n            selected_backend = _Backend.XFORMERS\n\n    # FlashAttn is valid for the model, checking if the package is installed.\n    if selected_backend == _Backend.FLASH_ATTN:\n        try:\n            import vllm_flash_attn  # noqa: F401\n\n            from vllm.attention.backends.flash_attn import (  # noqa: F401\n                FlashAttentionBackend)\n\n            supported_sizes = FlashAttentionBackend.get_supported_head_sizes()\n            if head_size not in supported_sizes:\n                logger.info(\n                    \"Cannot use FlashAttention-2 backend for head size %d.\",\n                    head_size)\n                selected_backend = _Backend.XFORMERS\n        except ImportError:\n            logger.info(\n                \"Cannot use FlashAttention-2 backend because the \"\n                \"vllm_flash_attn package is not found. \"\n                \"`pip install vllm-flash-attn` for better performance.\")\n            selected_backend = _Backend.XFORMERS\n\n    return selected_backend\n", "path": "vllm/attention/selector.py"}]}