[XPU] Enable FLASH_ATTN on XPU (#1332)

Signed-off-by: Yan Ma <yan.ma@intel.com>

tests/diffusion/attention/test_flash_attn.py

@@ -15,6 +15,9 @@ import torch
 from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
 from vllm_omni.diffusion.attention.backends.flash_attn import FlashAttentionImpl
 from vllm_omni.diffusion.attention.backends.sdpa import SDPAImpl
+from vllm_omni.platforms import current_omni_platform
+
+is_gpu = current_omni_platform.is_cuda_alike() or current_omni_platform.is_xpu()
 
 
 def create_attention_mask(batch_size: int, seq_len: int, valid_len: int, device: torch.device) -> torch.Tensor:
@@ -56,7 +59,7 @@ def pad_tensor(tensor: torch.Tensor, target_seq_len: int, pad_value: float = 0.0
     return torch.cat([tensor, padding], dim=1)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashAttention requires CUDA")
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
 def test_padding_equivalence():
     """
     Case 1: Test that padded and unpadded inputs produce similar outputs.
@@ -68,7 +71,7 @@ def test_padding_equivalence():
 
     Expected: Output A and Output B should be very close.
     """
-    device = torch.device("cuda")
+    device = torch.device(current_omni_platform.device_type)
     dtype = torch.bfloat16
 
     # Configuration
@@ -149,7 +152,7 @@ def test_padding_equivalence():
     print("✓ Case 1 PASSED: Padded and unpadded outputs are very close!")
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="FlashAttention requires CUDA")
+@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU")
 def test_fa_vs_sdpa():
     """
     Case 2: Compare FlashAttention and SDPA backends with padding.
@@ -162,7 +165,7 @@ def test_fa_vs_sdpa():
 
     Expected: FA and SDPA outputs should be very close.
     """
-    device = torch.device("cuda")
+    device = torch.device(current_omni_platform.device_type)
     dtype = torch.bfloat16
 
     # Configuration
@@ -265,8 +268,8 @@ if __name__ == "__main__":
     print("Running FlashAttention Padding Tests...")
     print("=" * 60)
 
-    # Try to run CUDA tests
-    if torch.cuda.is_available():
+    # Try to run tests
+    if is_gpu:
         try:
             print("\n[Running Case 1: Padding Equivalence for FA]")
             test_padding_equivalence()

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -48,6 +48,47 @@ class FlashAttentionImpl(AttentionImpl):
         self.causal = causal
         self.softmax_scale = softmax_scale
 
+    @staticmethod
+    def _unwrap_flash_output(out: torch.Tensor | tuple[torch.Tensor, ...]) -> torch.Tensor:
+        # FA3 may return (out, lse), FA2 returns out
+        return out[0] if isinstance(out, tuple) else out
+
+    def _forward_varlen_masked(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            _pad_input,
+            _unpad_input,
+            _upad_input,
+            flash_attn_varlen_func,
+        )
+
+        assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
+        query_length = query.size(1)
+        q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
+            query, key, value, attention_mask, query_length, _unpad_input
+        )
+
+        out_unpad = flash_attn_varlen_func(
+            q,
+            k,
+            v,
+            cu_seqlens_q=cu_seq_lens_q,
+            cu_seqlens_k=cu_seq_lens_k,
+            max_seqlen_q=max_length_q,
+            max_seqlen_k=max_length_k,
+            **{
+                "causal": self.causal,
+                "softmax_scale": self.softmax_scale,
+            },
+        )
+        out_unpad = self._unwrap_flash_output(out_unpad)
+        return _pad_input(out_unpad, indices_q, query.size(0), query_length)
+
     def forward_cuda(
         self,
         query: torch.Tensor,
@@ -56,15 +97,9 @@ class FlashAttentionImpl(AttentionImpl):
         attn_metadata: AttentionMetadata = None,
     ) -> torch.Tensor:
         """CUDA/ROCm flash attention implementation."""
-        # Import flash attention functions with fallback chain from utils/fa.py
-        # FA3 (fa3_fwd_interface) -> FA3 (flash_attn_interface) -> FA2 (flash_attn)
         from vllm_omni.diffusion.attention.backends.utils.fa import (
             HAS_FLASH_ATTN,
-            _pad_input,
-            _unpad_input,
-            _upad_input,
             flash_attn_func,
-            flash_attn_varlen_func,
         )
 
         if not HAS_FLASH_ATTN:
@@ -74,45 +109,76 @@ class FlashAttentionImpl(AttentionImpl):
                 "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
             )
 
-        query_length = query.size(1)
         attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
-        #  Contains at least one padding token in the sequence
+
         if attention_mask is not None and torch.any(~attention_mask):
-            assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)"
-            q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input(
-                query, key, value, attention_mask, query_length, _unpad_input
-            )
-
-            out_unpad = flash_attn_varlen_func(
-                q,
-                k,
-                v,
-                cu_seqlens_q=cu_seq_lens_q,
-                cu_seqlens_k=cu_seq_lens_k,
-                max_seqlen_q=max_length_q,
-                max_seqlen_k=max_length_k,
-                **{
-                    "causal": self.causal,
-                    "softmax_scale": self.softmax_scale,
-                },
-            )
-            if isinstance(out_unpad, tuple):
-                out_unpad = out_unpad[0]
-
-            out = _pad_input(out_unpad, indices_q, query.size(0), query_length)
-
-        else:
-            out = flash_attn_func(
+            return self._forward_varlen_masked(
                 query,
                 key,
                 value,
-                causal=self.causal,
-                softmax_scale=self.softmax_scale,
+                attention_mask,
             )
-            # FA3 may return (out, lse) tuple, FA2 returns just out
-            if isinstance(out, tuple):
-                out = out[0]
-        return out
+
+        out = flash_attn_func(
+            query,
+            key,
+            value,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+        return self._unwrap_flash_output(out)
+
+    def forward_xpu(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        """XPU flash attention implementation."""
+        from vllm_omni.diffusion.attention.backends.utils.fa import (
+            HAS_FLASH_ATTN,
+            flash_attn_varlen_func,
+        )
+
+        if not HAS_FLASH_ATTN:
+            raise ImportError(
+                "FlashAttentionBackend requires Flash Attention. "
+                "Please assure vllm-xpu-kernels properly installed. "
+                "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
+            )
+
+        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+
+        if attention_mask is not None and torch.any(~attention_mask):
+            return self._forward_varlen_masked(
+                query,
+                key,
+                value,
+                attention_mask,
+            )
+
+        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)
+        # b s ... -> (b s) ...
+        query = query.flatten(0, 1)
+        key = key.flatten(0, 1)
+        value = value.flatten(0, 1)
+
+        out = flash_attn_varlen_func(
+            query,
+            key,
+            value,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=q_len,
+            max_seqlen_k=q_len,
+            causal=self.causal,
+            softmax_scale=self.softmax_scale,
+        )
+        out = self._unwrap_flash_output(out)
+        # (b s) h d -> b s h d
+        return out.reshape(batch_size, q_len, *out.shape[1:])
 
     def forward_npu(
         self,

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -68,7 +68,7 @@ else:
 
 # If no FA backend available, SDPA backend will be selected at the platform level
 # flash_attn_func and flash_attn_varlen_func will be None
-HAS_FLASH_ATTN = flash_attn_func is not None
+HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
 
 
 def _index_first_axis(tensor, indices):