diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index cf8d4224f..ab010f3ab 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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, ) diff --git a/vllm_omni/platforms/xpu/platform.py b/vllm_omni/platforms/xpu/platform.py index 84707df08..5500e013f 100644 --- a/vllm_omni/platforms/xpu/platform.py +++ b/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()