[XPU] set flash_attn as default diffusion attn backend and fix k_len for cross_attn (#3525)

Signed-off-by: Chendi Xue <chendi.xue@intel.com>

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -136,7 +136,9 @@ class FlashAttentionImpl(AttentionImpl):
         )
 
         batch_size, q_len = query.size()[:2]
-        cu_seqlens = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
+        k_len = key.size(1)
+        cu_seqlens_q = torch.arange(0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device)
+        cu_seqlens_k = torch.arange(0, (batch_size + 1) * k_len, step=k_len, dtype=torch.int32, device=query.device)
         # b s ... -> (b s) ...
         query = query.flatten(0, 1)
         key = key.flatten(0, 1)
@@ -146,10 +148,10 @@ class FlashAttentionImpl(AttentionImpl):
             q=query,
             k=key,
             v=value,
-            cu_seqlens_q=cu_seqlens,
-            cu_seqlens_k=cu_seqlens,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
             max_seqlen_q=q_len,
-            max_seqlen_k=q_len,
+            max_seqlen_k=k_len,
             causal=self.causal,
             softmax_scale=self.softmax_scale,
         )

vllm_omni/platforms/xpu/platform.py

@@ -36,12 +36,20 @@ class XPUOmniPlatform(OmniPlatform, XPUPlatform):
         selected_backend: str | None,
         head_size: int,
     ) -> str:
+        compute_capability = torch.xpu.get_device_capability()
+        # Intel Max 1100 and 1550 will not support flash_attn currently
+        flash_attn_supported = compute_capability["architecture"] not in [13136561920]
+
         if selected_backend is not None:
             backend_upper = selected_backend.upper()
             backend = DiffusionAttentionBackendEnum[backend_upper]
             logger.debug("Using diffusion attention backend '%s'", backend_upper)
             return backend.get_path()
 
+        if flash_attn_supported:
+            logger.debug("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
         logger.debug("Defaulting to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()