diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index a9f2d925c..dc54acf04 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -10,8 +10,26 @@ from vllm_omni.diffusion.attention.backends.abstract import ( AttentionMetadata, ) +# 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, +) + logger = init_logger(__name__) +if not HAS_FLASH_ATTN: + raise ImportError( + "FlashAttentionBackend requires Flash Attention. " + "Please install one of: fa3-fwd, flash-attention, or flash-attn. " + "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" + ) + class FlashAttentionBackend(AttentionBackend): accept_output_buffer: bool = True @@ -56,14 +74,6 @@ class FlashAttentionImpl(AttentionImpl): attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: """CUDA/ROCm flash attention implementation.""" - from vllm_omni.diffusion.attention.backends.utils.fa import ( - _pad_input, - _unpad_input, - _upad_input, - flash_attn_func, - flash_attn_varlen_func, - ) - 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 @@ -92,13 +102,16 @@ class FlashAttentionImpl(AttentionImpl): out = _pad_input(out_unpad, indices_q, query.size(0), query_length) else: - out: torch.Tensor = flash_attn_func( + out = flash_attn_func( query, key, value, causal=self.causal, softmax_scale=self.softmax_scale, ) + # FA3 may return (out, lse) tuple, FA2 returns just out + if isinstance(out, tuple): + out = out[0] return out def forward_npu( diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py index 7ec092785..d80d81fb4 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_globals.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_globals.py @@ -3,30 +3,53 @@ # Copyright (c) 2024, Jiarui Fang. # Adapted from https://github.com/feifeibear/long-context-attention - -# test if flash_attn is available +# test if flash_attn (FA2) is available try: import flash_attn # noqa: F401 - from flash_attn.flash_attn_interface import _flash_attn_backward, _flash_attn_forward # noqa: F401 + from flash_attn.flash_attn_interface import _flash_attn_forward # noqa: F401 HAS_FLASH_ATTN = True -except ImportError: +except (ImportError, ModuleNotFoundError): HAS_FLASH_ATTN = False -try: - from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward # noqa: F401 - from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper # noqa: F401 - from flash_attn_interface import flash_attn_func as flash3_attn_func # noqa: F401 +# FA3 detection: try multiple sources (forward only, no backward needed for inference) +# Source 1: flash_attn_interface (from flash-attention source build) +# Source 2: fa3_fwd_interface (from fa3-fwd PyPI package, supports Ampere/Ada/Hopper) +# Note: FA3 high-level API may or may not return softmax_lse depending on version. +# For Ring Attention which requires LSE, we fall back to low-level API if needed. +HAS_FA3 = False +fa3_fwd_func = None # Low-level forward function (_flash_attn_forward) +fa3_attn_func = None # High-level attention function (flash_attn_func) - HAS_FLASH_ATTN_HOPPER = True -except ImportError: - HAS_FLASH_ATTN_HOPPER = False +# Try flash_attn_interface first (from flash-attention source build) +try: + from flash_attn_interface import _flash_attn_forward as fa3_fwd_func # noqa: F401 + from flash_attn_interface import flash_attn_func as fa3_attn_func # noqa: F401 + + HAS_FA3 = True +except (ImportError, ModuleNotFoundError): + pass + +# Fallback: try fa3_fwd_interface (PyPI package, supports Ampere/Ada/Hopper) +if not HAS_FA3: + try: + from fa3_fwd_interface import _flash_attn_forward as fa3_fwd_func # noqa: F401 + from fa3_fwd_interface import flash_attn_func as fa3_attn_func # noqa: F401 + + HAS_FA3 = True + except (ImportError, ModuleNotFoundError): + pass + +# Legacy aliases for backward compatibility +HAS_FLASH_ATTN_HOPPER = HAS_FA3 +flash_attn_forward_hopper = fa3_fwd_func +flash3_attn_func = fa3_attn_func try: from flashinfer.prefill import single_prefill_with_kv_cache # noqa: F401 HAS_FLASHINFER = True -except ImportError: +except (ImportError, ModuleNotFoundError): HAS_FLASHINFER = False try: @@ -34,26 +57,26 @@ try: from aiter import flash_attn_func as flash_attn_func_aiter # noqa: F401 HAS_AITER = True -except ImportError: +except (ImportError, ModuleNotFoundError): HAS_AITER = False try: import sageattention # noqa: F401 HAS_SAGE_ATTENTION = True -except ImportError: +except (ImportError, ModuleNotFoundError): HAS_SAGE_ATTENTION = False try: import spas_sage_attn # noqa: F401 HAS_SPARSE_SAGE_ATTENTION = True -except ImportError: +except (ImportError, ModuleNotFoundError): HAS_SPARSE_SAGE_ATTENTION = False try: import torch_npu # noqa: F401 HAS_NPU = True -except ImportError: +except (ImportError, ModuleNotFoundError): HAS_NPU = False diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py index 2444dd894..ad61f9a20 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_kernels.py @@ -6,7 +6,13 @@ import math import torch -from .ring_globals import HAS_AITER, HAS_FLASH_ATTN, HAS_FLASH_ATTN_HOPPER, HAS_FLASHINFER, HAS_NPU +from .ring_globals import ( + HAS_AITER, + HAS_FA3, + HAS_FLASH_ATTN, + HAS_FLASHINFER, + fa3_fwd_func, +) _scaled_dot_product_flash_attention = torch.ops.aten._scaled_dot_product_flash_attention _scaled_dot_product_efficient_attention = torch.ops.aten._scaled_dot_product_efficient_attention @@ -26,23 +32,11 @@ if HAS_FLASH_ATTN: import flash_attn from flash_attn.flash_attn_interface import _flash_attn_forward -if HAS_FLASH_ATTN_HOPPER: - from flash_attn_interface import _flash_attn_backward as flash_attn_func_hopper_backward - from flash_attn_interface import _flash_attn_forward as flash_attn_forward_hopper - from flash_attn_interface import flash_attn_func as flash3_attn_func -else: - flash_attn_forward_hopper = None - flash_attn_func_hopper_backward = None - flash3_attn_func = None - if HAS_FLASHINFER: from flashinfer.prefill import single_prefill_with_kv_cache _LOG2_E = math.log2(math.e) -if HAS_NPU: - import torch_npu - def pytorch_attn_forward( q: torch.Tensor, @@ -146,52 +140,35 @@ def flash_attn_forward( return block_out, block_lse -def flash_attn3_func_forward( - q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax -): - assert HAS_FLASH_ATTN_HOPPER - # current signature of flash_attn_forward_hopper: - # (q, k, v, softmax_scale, causal, window_size, descale_q=None, descale_k=None, descale_v=None, gqa_parallel=False) +def fa3_forward(q, k, v, dropout_p, softmax_scale, causal, window_size, softcap, alibi_slopes, return_softmax): + """FA3 forward pass for inference. - out, softmax_lse, *unused = flash_attn_forward_hopper( - q=q, - k=k, - v=v, - k_new=None, - v_new=None, - qv=None, - out=None, - cu_seqlens_q=None, - cu_seqlens_k=None, - cu_seqlens_k_new=None, - seqused_q=None, - seqused_k=None, - max_seqlen_q=None, - max_seqlen_k=None, - page_table=None, - kv_batch_idx=None, - leftpad_k=None, - rotary_cos=None, - rotary_sin=None, - seqlens_rotary=None, - q_descale=None, - k_descale=None, - v_descale=None, + FA3 supports Ampere, Ada, and Hopper GPUs. Dropout is ignored since FA3 is inference-only. + Uses low-level API (_flash_attn_forward) which always returns softmax_lse, + required for Ring Attention's correct accumulation. + """ + assert HAS_FA3, "FA3 is not available" + assert fa3_fwd_func is not None, "FA3 low-level API (fa3_fwd_func) not available" + + # Low-level API always returns (out, softmax_lse, S_dmask, rng_state) + out, softmax_lse, *_ = fa3_fwd_func( + q, + k, + v, softmax_scale=softmax_scale, - causal=False, - window_size=(-1, -1), - attention_chunk=0, - softcap=0.0, - rotary_interleaved=True, - scheduler_metadata=None, - num_splits=0, - pack_gqa=None, - sm_margin=0, + causal=causal, + window_size_left=window_size[0] if window_size else -1, + window_size_right=window_size[1] if window_size else -1, + softcap=softcap if softcap else 0.0, ) return out, softmax_lse +# Legacy alias for backward compatibility +flash_attn3_func_forward = fa3_forward + + def flash_attn_forward_aiter( q, k, @@ -264,20 +241,3 @@ def flashinfer_attn_forward( raise ValueError(f"Invalid input shape: {q.shape}") lse = lse / _LOG2_E return out, lse - - -def npu_attn_forward(q, k, v, softmax_scale=None, layout="BSND"): - assert HAS_NPU, "torch_npu is not available" - softmax_scale = q.shape[-1] ** (-0.5) - block_out, block_lse = torch_npu.npu_fused_infer_attention_score( - q, - k, - v, - num_heads=q.shape[-2], - input_layout=layout, - scale=softmax_scale, - softmax_lse_flag=True, - pre_tokens=65535, - next_tokens=65535, - ) - return block_out, block_lse.squeeze(dim=-1) diff --git a/vllm_omni/diffusion/attention/backends/ring/ring_selector.py b/vllm_omni/diffusion/attention/backends/ring/ring_selector.py index 56e4a2e6b..77189c61b 100644 --- a/vllm_omni/diffusion/attention/backends/ring/ring_selector.py +++ b/vllm_omni/diffusion/attention/backends/ring/ring_selector.py @@ -10,7 +10,6 @@ from functools import partial import torch from .ring_globals import ( - HAS_NPU, HAS_SAGE_ATTENTION, HAS_SPARSE_SAGE_ATTENTION, ) @@ -28,9 +27,6 @@ if HAS_SAGE_ATTENTION: if HAS_SPARSE_SAGE_ATTENTION: from spas_sage_attn.autotune import SparseAttentionMeansim -if HAS_NPU: - from torch_npu import npu_fused_infer_attention_score - class AttnType(Enum): AITER = "aiter" @@ -44,7 +40,6 @@ class AttnType(Enum): SAGE_FP8 = "sage_fp8" SAGE_FP8_SM90 = "sage_fp8_sm90" SPARSE_SAGE = "sparse_sage" - NPU = "npu" @classmethod def from_string(cls, s: str): @@ -157,11 +152,6 @@ def select_flash_attn_impl( return fn - elif impl_type == AttnType.NPU: - if not HAS_NPU: - raise ImportError("torch_npu is not available!") - return npu_fused_infer_attention_score - elif attn_processor is not None: return attn_processor diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index ea12bc0f1..3344c9963 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -17,17 +17,53 @@ import torch.nn.functional as F from vllm_omni.platforms import current_omni_platform -if current_omni_platform.is_rocm(): - from vllm._aiter_ops import is_aiter_found_and_supported +# Flash Attention function detection with fallback chain +flash_attn_func = None +flash_attn_varlen_func = None - # Choose to enable this by default on ROCm - # Whenever possible as it is the fastest backend - if is_aiter_found_and_supported(): - from aiter import flash_attn_func, flash_attn_varlen_func # noqa: F401 - else: - raise ImportError("Aiter is not found and supported on currentROCm device.") +if current_omni_platform.is_rocm(): + # ROCm: try Aiter first + try: + from vllm._aiter_ops import is_aiter_found_and_supported + + if is_aiter_found_and_supported(): + from aiter import flash_attn_func, flash_attn_varlen_func # noqa: F401 + except (ImportError, ModuleNotFoundError): + pass else: - from fa3_fwd_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 + # 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 + 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 + 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 + 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 + +# 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 def _index_first_axis(tensor, indices): diff --git a/vllm_omni/diffusion/attention/parallel/ring.py b/vllm_omni/diffusion/attention/parallel/ring.py index 2d63f40c0..090d2ef61 100644 --- a/vllm_omni/diffusion/attention/parallel/ring.py +++ b/vllm_omni/diffusion/attention/parallel/ring.py @@ -10,7 +10,7 @@ import torch from vllm.logger import init_logger # import torch.distributed as dist # Not used directly here, but good practice if needed -from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FLASH_ATTN +from vllm_omni.diffusion.attention.backends.ring.ring_globals import HAS_FA3, HAS_FLASH_ATTN from vllm_omni.diffusion.attention.backends.ring.ring_selector import AttnType from vllm_omni.diffusion.attention.parallel.base import ( ParallelAttentionContext, @@ -116,11 +116,14 @@ class RingParallelAttention: except Exception: backend_pref = None - # Fallback for FP32 or if Flash Attention is not available - if query.dtype == torch.float32 or not HAS_FLASH_ATTN: - if not HAS_FLASH_ATTN and backend_pref != "sdpa": + # Determine attention type with fallback chain: FA3 -> FA2 -> SDPA + # FP32 is not supported by Flash Attention, force SDPA + if query.dtype == torch.float32: + backend_pref = "sdpa" + elif not HAS_FA3 and not HAS_FLASH_ATTN: + if backend_pref != "sdpa": logger = init_logger(__name__) - logger.warning_once("Flash Attention is not available! Force enabling SDPA.") + logger.warning_once("Flash Attention (FA2/FA3) is not available! Force enabling SDPA.") backend_pref = "sdpa" # Extract joint tensors @@ -150,6 +153,9 @@ class RingParallelAttention: from vllm_omni.diffusion.attention.backends.ring_flash_attn import ring_flash_attn_func + # Prefer FA3 over FA2 for better performance (FA3 supports Ampere/Ada/Hopper) + attn_type = AttnType.FA3 if HAS_FA3 else AttnType.FA + return ring_flash_attn_func( query, key, @@ -161,9 +167,8 @@ class RingParallelAttention: softcap=0.0, alibi_slopes=None, deterministic=False, - # return_attn_probs=False, # Removed as it might not be supported in signature group=self._sp_group.ring_group, - attn_type=AttnType.FA, + attn_type=attn_type, joint_tensor_key=joint_key, joint_tensor_value=joint_value, joint_strategy=joint_strategy, diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index be2b58ae8..dd566b059 100644 --- a/vllm_omni/diffusion/envs.py +++ b/vllm_omni/diffusion/envs.py @@ -66,14 +66,32 @@ class PackagesEnvChecker: if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name: return False + # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn) + # Try FA3 from fa3-fwd PyPI package + try: + import fa3_fwd_interface # noqa: F401 + + return True + except (ImportError, ModuleNotFoundError): + pass + + # Try FA3 from flash-attention source build + try: + import flash_attn_interface # noqa: F401 + + return True + except (ImportError, ModuleNotFoundError): + pass + + # Try FA2 from flash-attn package from flash_attn import __version__ if __version__ < "2.6.0": raise ImportError("install flash_attn >= 2.6.0") return True - except ImportError: + except (ImportError, ModuleNotFoundError): if not packages_info.get("has_aiter", False): - logger.warning('Flash Attention library "flash_attn" not found, using pytorch attention implementation') + logger.warning("No Flash Attention backend found, using pytorch SDPA implementation") return False def get_packages_info(self) -> dict: diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 9c65877be..89e403f99 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -39,22 +39,34 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): selected_backend: str | None, head_size: int, ) -> str: + from vllm_omni.diffusion.envs import PACKAGES_CHECKER + # Check compute capability for Flash Attention support # Flash Attention requires compute capability >= 8.0 and < 10.0 compute_capability = cls.get_device_capability() - flash_attn_supported = False + compute_supported = False if compute_capability is not None: major, minor = compute_capability capability = major * 10 + minor - flash_attn_supported = 80 <= capability < 100 + compute_supported = 80 <= capability < 100 + + # Check if FA packages are available + packages_info = PACKAGES_CHECKER.get_packages_info() + packages_available = packages_info.get("has_flash_attn", False) + + # Both compute capability and packages must be available for FA + flash_attn_supported = compute_supported and packages_available if selected_backend is not None: backend_upper = selected_backend.upper() if backend_upper == "FLASH_ATTN" and not flash_attn_supported: - logger.warning( - "Flash Attention requires GPU with compute capability >= 8.0 " - "and < 10.0. Falling back to TORCH_SDPA backend." - ) + if not compute_supported: + logger.warning( + "Flash Attention requires GPU with compute capability >= 8.0 " + "and < 10.0. Falling back to TORCH_SDPA backend." + ) + elif not packages_available: + logger.warning("Flash Attention packages not available. Falling back to TORCH_SDPA backend.") logger.info("Defaulting to diffusion attention backend SDPA") return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path() backend = DiffusionAttentionBackendEnum[backend_upper]