fix vit attn for models like THUDM/GLM-4v-9B on xpu (#339)

Signed-off-by: Yan Ma <[email protected]>

Author: yma11

vllm/attention/layer.py

@@ -522,8 +522,7 @@ def __init__(
         use_upstream_fa = False
 
         if current_platform.is_xpu():
-            # currently, only torch_sdpa is supported on xpu
-            self.attn_backend = _Backend.TORCH_SDPA
+            self.attn_backend = _Backend.IPEX
         else:
             self.attn_backend = (
                 backend
@@ -611,7 +610,10 @@ def forward(
             out = xops.memory_efficient_attention_forward(
                 query, key, value, scale=self.scale
             )
-        elif self.attn_backend == _Backend.TORCH_SDPA:
+        elif (
+            self.attn_backend == _Backend.TORCH_SDPA
+            or self.attn_backend == _Backend.IPEX
+        ):
             query, key, value = (x.transpose(1, 2) for x in (query, key, value))
             out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)
             out = out.transpose(1, 2)

vllm/model_executor/models/qwen2_5_vl.py

@@ -718,6 +718,7 @@ def __init__(
             _Backend.TORCH_SDPA,
             _Backend.XFORMERS,
             _Backend.ROCM_AITER_FA,
+            _Backend.IPEX,
         }:
             raise RuntimeError(
                 f"Qwen2.5-VL does not support {self.attn_backend} backend now."
@@ -859,6 +860,7 @@ def compute_attn_mask_seqlen(
         if (
             self.attn_backend == _Backend.FLASH_ATTN
             or self.attn_backend == _Backend.ROCM_AITER_FA
+            or self.attn_backend == _Backend.IPEX
         ):
             max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
         elif self.attn_backend == _Backend.XFORMERS:

vllm/model_executor/models/qwen2_vl.py

@@ -822,12 +822,11 @@ def compute_attn_mask_seqlen(
         if (
             self.attn_backend == _Backend.FLASH_ATTN
             or self.attn_backend == _Backend.ROCM_AITER_FA
+            or self.attn_backend == _Backend.IPEX
         ):
             max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
         elif self.attn_backend == _Backend.XFORMERS:
             seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
-        elif self.attn_backend == _Backend.IPEX:
-            max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
         return max_seqlen, seqlens
 
     def forward(

vllm/platforms/xpu.py

@@ -116,7 +116,9 @@ def get_device_total_memory(cls, device_id: int = 0) -> int:
         return device_props.total_memory
 
     @classmethod
-    def get_vit_attn_backend(cls, support_fa: bool = False) -> _Backend:
+    def get_vit_attn_backend(cls, head_size: int, dtype: torch.dtype) -> _Backend:
+        from vllm.attention.backends.registry import _Backend
+
         return _Backend.IPEX
 
     @classmethod