diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 3f3862405..5ec1517d9 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/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() diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 45158e8f8..5c586c063 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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, diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index eb8f72bad..1474598d7 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/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):