[Refactor] Extract has_flash_attn_pkg in OmniPlatform (#3068)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>

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):

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."""

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,

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

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()
 

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,