Fix: selecting correct backend for MultiHeadAttention

vllm/attention/layer.py

@@ -191,11 +191,16 @@ def __init__(
                                         kv_cache_dtype=None,
                                         block_size=16,
                                         is_attention_free=False)
-        if attn_backend in {_Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1}:
-            attn_backend = _Backend.XFORMERS
 
-        self.attn_backend = attn_backend if attn_backend in {
-            _Backend.TORCH_SDPA, _Backend.XFORMERS
+        attn_backend_enum = backend_name_to_enum(attn_backend.get_name())
+
+        if attn_backend_enum in {
+                _Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1
+        }:
+            attn_backend_enum = _Backend.XFORMERS
+
+        self.attn_backend = attn_backend_enum if attn_backend_enum in {
+            _Backend.TORCH_SDPA, _Backend.XFORMERS, _Backend.HPU_ATTN
         } else _Backend.TORCH_SDPA
 
     def forward(
@@ -228,6 +233,28 @@ def forward(
                                                  value,
                                                  scale=self.scale)
             out = out.transpose(1, 2)
+        elif self.attn_backend == _Backend.HPU_ATTN:
+            from habana_frameworks.torch.hpex.kernels import FusedSDPA
+            from vllm_hpu_extension.utils import ModuleFusedSDPA
+
+            fsdpa_op = ModuleFusedSDPA(FusedSDPA)
+
+            query, key, value = (x.transpose(1, 2)
+                                 for x in (query, key, value))
+
+            out = fsdpa_op(query,
+                           key,
+                           value,
+                           None,
+                           dropout_p=0.0,
+                           is_causal=True,
+                           scale=self.scale,
+                           softmax_mode="fast",
+                           recompute_mode=True,
+                           valid_sequence_lengths=None)
+
+            out = out.transpose(1, 2).contiguous()
+
         return out.view(bsz, q_len, -1)