From 68a484589d2c5f50107ce4466a792fc205e28e70 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 1 May 2026 17:40:05 -0600 Subject: [PATCH] [CI] Add vLLM FA to Flash Attention Selection (Bagel H100 Failures) (#3302) Signed-off-by: Alex Brooks --- .../attention/backends/flash_attn.py | 84 ++++++++++--------- .../diffusion/attention/backends/utils/fa.py | 27 +++++- 2 files changed, 69 insertions(+), 42 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index d38ea4f6e..53ddb1753 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project + import torch from vllm.logger import init_logger @@ -9,6 +10,13 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionImpl, AttentionMetadata, ) +from vllm_omni.diffusion.attention.backends.utils.fa import ( + HAS_FLASH_ATTN, + _pad_input, + _unpad_input, + _upad_input, + flash_attn_varlen_func, +) logger = init_logger(__name__) @@ -60,13 +68,6 @@ class FlashAttentionImpl(AttentionImpl): 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( @@ -89,6 +90,42 @@ class FlashAttentionImpl(AttentionImpl): out_unpad = self._unwrap_flash_output(out_unpad) return _pad_input(out_unpad, indices_q, query.size(0), query_length) + def _forward_varlen_dense( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + ) -> torch.Tensor: + """Common wrapper for calling flash_attn_varlen_func for XPU and CUDA in vLLM. + + NOTE: careful to keep the kwargs on these aligned and to pass everything as a keyword + argument, because some of the args differ positionally at the moment. + + https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/vllm_flash_attn/flash_attn_interface.py#L176 + https://github.com/vllm-project/vllm/blob/v0.20.0/vllm/_xpu_ops.py#L310 + """ + 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( + q=query, + k=key, + v=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_cuda( self, query: torch.Tensor, @@ -97,11 +134,6 @@ class FlashAttentionImpl(AttentionImpl): attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: """CUDA/ROCm/MUSA flash attention implementation.""" - from vllm_omni.diffusion.attention.backends.utils.fa import ( - HAS_FLASH_ATTN, - flash_attn_func, - ) - if not HAS_FLASH_ATTN: raise ImportError( "FlashAttentionBackend requires Flash Attention. " @@ -119,14 +151,11 @@ class FlashAttentionImpl(AttentionImpl): attention_mask, ) - out = flash_attn_func( + return self._forward_varlen_dense( query, key, value, - causal=self.causal, - softmax_scale=self.softmax_scale, ) - return self._unwrap_flash_output(out) def forward_xpu( self, @@ -136,11 +165,6 @@ class FlashAttentionImpl(AttentionImpl): 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. " @@ -158,27 +182,11 @@ class FlashAttentionImpl(AttentionImpl): 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( + return self._forward_varlen_dense( 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 776d3723a..4360fcaac 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -51,28 +51,37 @@ else: # CUDA: try FA3 -> FA2 fallback chain # Try FA3 from fa3-fwd PyPI package try: - from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 + from fa3_fwd_interface import flash_attn_varlen_func # noqa: F401 except (ImportError, ModuleNotFoundError): pass # Fallback: Try FA3 from flash-attention source build if flash_attn_func is None: try: - from flash_attn_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 + from flash_attn_interface import flash_attn_varlen_func # noqa: F401 except (ImportError, ModuleNotFoundError): pass # Fallback: Try FA2 from flash-attn package (try multiple import paths) if flash_attn_func is None: try: - from flash_attn import flash_attn_func, flash_attn_varlen_func # noqa: F401 + from flash_attn import flash_attn_varlen_func # noqa: F401 except (ImportError, ModuleNotFoundError): pass if flash_attn_func is None: try: from flash_attn.flash_attn_interface import ( # noqa: F401 - flash_attn_func, + flash_attn_varlen_func, + ) + except (ImportError, ModuleNotFoundError): + pass + + # Fallback: Try vLLM's encapsulated flash attention dispatcher + # TODO discuss priority and remove potentially redundant fallbacks + if flash_attn_func is None: + try: + from vllm.vllm_flash_attn import ( # noqa: F401 flash_attn_varlen_func, ) except (ImportError, ModuleNotFoundError): @@ -113,6 +122,16 @@ def is_flash_attn_installed() -> bool: if __version__ < "2.6.0": raise ImportError("install flash_attn >= 2.6.0") return True + except (ImportError, ModuleNotFoundError): + pass + + # Try vLLM's flash attention wrapper + try: + from vllm.vllm_flash_attn import ( # noqa: F401 + flash_attn_varlen_func, + ) + + return True except (ImportError, ModuleNotFoundError): logger.warning("No Flash Attention backend found, using pytorch SDPA implementation") return False