[Feat] Add MUSA flash attention support via mate package (#2451)

Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -209,3 +209,13 @@ class FlashAttentionImpl(AttentionImpl):
             layout="BNSD",
         )
         return output
+
+    def forward_musa(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata = None,
+    ) -> torch.Tensor:
+        # XXX (MUSA): MUSA uses the same implementation as XPU (mate only provides flash_attn_varlen_func)
+        return self.forward_xpu(query, key, value, attn_metadata)

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_flash_attention_utils.py
+from functools import lru_cache
+
 import torch
 import torch.nn.functional as F
 
@@ -38,8 +40,10 @@ elif current_omni_platform.is_xpu():
     except (ImportError, ModuleNotFoundError):
         pass
 elif current_omni_platform.is_musa():
-    # XXX (MUSA): Add MUSA-specific Flash Attention when available
-    pass
+    try:
+        from mate import flash_attn_varlen_func  # noqa: F401
+    except (ImportError, ModuleNotFoundError):
+        pass
 else:
     # CUDA: try FA3 -> FA2 fallback chain
     # Try FA3 from fa3-fwd PyPI package
@@ -76,6 +80,12 @@ else:
 HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
 
 
+@lru_cache(maxsize=1)
+def is_mate_available() -> bool:
+    """Check if MATE (MUSA Flash Attention) is available."""
+    return current_omni_platform.is_musa() and flash_attn_varlen_func is not None
+
+
 def _index_first_axis(tensor, indices):
     """
     A local implementation of the PyTorch indexing operation `tensor[indices]` on the first axis,

vllm_omni/diffusion/envs.py

@@ -7,6 +7,7 @@ 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:
@@ -52,6 +53,10 @@ class PackagesEnvChecker:
         """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

vllm_omni/diffusion/layers/rope.py

@@ -145,6 +145,14 @@ class RotaryEmbedding(CustomOp):
     ) -> torch.Tensor:
         return self.forward_native(x, cos, sin)
 
+    def forward_musa(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        return self.forward_native(x, cos, sin)
+
     def forward_native(
         self,
         x: torch.Tensor,

vllm_omni/platforms/musa/platform.py

@@ -8,6 +8,7 @@ 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__)
@@ -54,9 +55,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
     ) -> str:
         """Get the diffusion attention backend class path for MUSA platform.
 
-        MUSA currently supports SDPA (Scaled Dot Product Attention) as the
-        primary backend. Flash Attention support may be added in future
-        when MUSA-specific implementations are available.
+        MUSA supports FLASH_ATTN via the mate package, and SDPA as fallback.
 
         Args:
             selected_backend: User-selected backend name (e.g., "FLASH_ATTN",
@@ -66,13 +65,24 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase):
         Returns:
             Fully qualified class path of the selected backend.
         """
+
+        flash_attn_available = is_mate_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.")
+                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 SDPA for MUSA as it's the most compatible backend
+        # Default to FLASH_ATTN if mate is available, otherwise SDPA
+        if flash_attn_available:
+            logger.info("Defaulting to diffusion attention backend FLASH_ATTN")
+            return DiffusionAttentionBackendEnum.FLASH_ATTN.get_path()
+
         logger.info("Defaulting to diffusion attention backend SDPA")
         return DiffusionAttentionBackendEnum.TORCH_SDPA.get_path()