From 52f17bf29263ebba3c985a918bbfaa0eabf2ff29 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Mon, 27 Apr 2026 12:07:34 +0800 Subject: [PATCH] [Refactor] Extract has_flash_attn_pkg in OmniPlatform (#3068) Signed-off-by: Xiaodong Ye Co-authored-by: Canlin Guo --- .../diffusion/attention/backends/utils/fa.py | 38 ++++++++++++-- vllm_omni/diffusion/envs.py | 52 ++----------------- vllm_omni/platforms/cuda/platform.py | 14 +++++ vllm_omni/platforms/interface.py | 5 ++ vllm_omni/platforms/musa/platform.py | 37 ++++++++++--- vllm_omni/platforms/rocm/platform.py | 6 +++ 6 files changed, 94 insertions(+), 58 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index 18886871a..776d3723a 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -16,9 +16,12 @@ from functools import lru_cache import torch import torch.nn.functional as F +from vllm.logger import init_logger from vllm_omni.platforms import current_omni_platform +logger = init_logger(__name__) + # Flash Attention function detection with fallback chain flash_attn_func = None flash_attn_varlen_func = None @@ -81,9 +84,38 @@ HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not No @lru_cache(maxsize=1) -def is_mate_available() -> bool: - """Check if MATE (MUSA AI Tensor Engine) is available.""" - return current_omni_platform.is_musa() and flash_attn_func is not None or flash_attn_varlen_func is not None +def is_flash_attn_installed() -> bool: + """Return whether a Flash Attention backend package is importable. + + Shared by CUDA/ROCm/MUSA platforms. + """ + try: + # 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, ModuleNotFoundError): + logger.warning("No Flash Attention backend found, using pytorch SDPA implementation") + return False def _index_first_axis(tensor, indices): diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index ea7b2c24c..8ed8fbefd 100644 --- a/vllm_omni/diffusion/envs.py +++ b/vllm_omni/diffusion/envs.py @@ -5,9 +5,6 @@ import os from collections.abc import Callable from typing import TYPE_CHECKING, Any -from vllm.logger import init_logger - -from vllm_omni.diffusion.attention.backends.utils.fa import is_mate_available from vllm_omni.platforms import current_omni_platform if TYPE_CHECKING: @@ -30,8 +27,6 @@ environment_variables: dict[str, Callable[[], Any]] = { "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")), } -logger = init_logger(__name__) - class PackagesEnvChecker: """Singleton class for checking package availability.""" @@ -46,58 +41,17 @@ class PackagesEnvChecker: def initialize(self): packages_info = {} - packages_info["has_flash_attn"] = self._check_flash_attn(packages_info) + packages_info["has_flash_attn"] = self._check_flash_attn() self.packages_info = packages_info - def _check_flash_attn(self, packages_info) -> bool: + def _check_flash_attn(self) -> bool: """Check if flash attention is available and compatible.""" platform = current_omni_platform - # MUSA uses MATE for flash attention - if platform.is_musa(): - return is_mate_available() - - # Flash attention requires CUDA-like platforms (CUDA or ROCm) - if not platform.is_cuda_alike(): - return False - - # Check if devices are available if platform.get_device_count() == 0: return False - try: - gpu_name = platform.get_device_name() - # Turing/Tesla/T4 GPUs don't support flash attention well - 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, ModuleNotFoundError): - if not packages_info.get("has_aiter", False): - logger.warning("No Flash Attention backend found, using pytorch SDPA implementation") - return False + return platform.has_flash_attn_package() def get_packages_info(self) -> dict: """Get the packages info dictionary.""" diff --git a/vllm_omni/platforms/cuda/platform.py b/vllm_omni/platforms/cuda/platform.py index 6bf740a01..6d955dcfb 100644 --- a/vllm_omni/platforms/cuda/platform.py +++ b/vllm_omni/platforms/cuda/platform.py @@ -33,6 +33,20 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase): def get_default_stage_config_path(cls) -> str: return "vllm_omni/model_executor/stage_configs" + @classmethod + def has_flash_attn_package(cls) -> bool: + from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed + + # Turing/Tesla/T4 GPUs don't support flash attention well + gpu_name = cls.get_device_name() + if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name: + return False + + if not is_flash_attn_installed(): + return False + + return True + @classmethod def get_diffusion_attn_backend_cls( cls, diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 11eec76ac..16a33f9be 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -105,6 +105,11 @@ class OmniPlatform(Platform): """Check if the platform supports torch.compile with inductor backend.""" raise NotImplementedError + @classmethod + def has_flash_attn_package(cls) -> bool: + """Check if a Flash Attention package is available and usable on this platform.""" + return False + @classmethod def get_torch_device(cls, local_rank: int | None = None) -> torch.device: raise NotImplementedError diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 77ef22510..5fe3fa4bf 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -8,7 +8,6 @@ from vllm.logger import init_logger from vllm_musa.platform import MUSAPlatformBase from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum -from vllm_omni.diffusion.attention.backends.utils.fa import is_mate_available from vllm_omni.platforms.interface import OmniPlatform, OmniPlatformEnum logger = init_logger(__name__) @@ -47,6 +46,12 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): # MUSA uses default runtime preparation return None + @classmethod + def has_flash_attn_package(cls) -> bool: + from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed + + return is_flash_attn_installed() + @classmethod def get_diffusion_attn_backend_cls( cls, @@ -65,21 +70,41 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): Returns: Fully qualified class path of the selected backend. """ + from vllm_omni.diffusion.envs import PACKAGES_CHECKER - flash_attn_available = is_mate_available() + # Check compute capability for Flash Attention support + # Flash Attention requires compute capability >= 3.1 + compute_capability = cls.get_device_capability() + compute_supported = False + if compute_capability is not None: + major, minor = compute_capability + capability = major * 10 + minor + compute_supported = capability >= 31 + + # 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_available: - logger.warning("Flash Attention (mate package) not available. Falling back to TORCH_SDPA backend.") + if backend_upper == "FLASH_ATTN" and not flash_attn_supported: + if not compute_supported: + logger.warning( + "Flash Attention requires MUSA GPU with compute capability >= 3.1. " + "Falling back to TORCH_SDPA backend." + ) + elif not packages_available: + logger.warning("Flash Attention (mate package) 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] logger.info("Using diffusion attention backend '%s'", backend_upper) return backend.get_path() - # Default to FLASH_ATTN if mate is available, otherwise SDPA - if flash_attn_available: + if flash_attn_supported: logger.info("Defaulting to diffusion attention backend FLASH_ATTN") return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path() diff --git a/vllm_omni/platforms/rocm/platform.py b/vllm_omni/platforms/rocm/platform.py index 7b0e09c12..84aca4629 100644 --- a/vllm_omni/platforms/rocm/platform.py +++ b/vllm_omni/platforms/rocm/platform.py @@ -56,6 +56,12 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform): def get_omni_generation_worker_cls(cls) -> str: return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker" + @classmethod + def has_flash_attn_package(cls) -> bool: + from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed + + return is_flash_attn_installed() + @classmethod def get_diffusion_attn_backend_cls( cls,