From fb3c6bd9b131479b9f7c76afed47bf8d87ca9718 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Wed, 8 Apr 2026 16:01:30 +0800 Subject: [PATCH] [Feat] Add MUSA flash attention support via mate package (#2451) Signed-off-by: Xiaodong Ye --- .../diffusion/attention/backends/flash_attn.py | 10 ++++++++++ .../diffusion/attention/backends/utils/fa.py | 14 ++++++++++++-- vllm_omni/diffusion/envs.py | 5 +++++ vllm_omni/diffusion/layers/rope.py | 8 ++++++++ vllm_omni/platforms/musa/platform.py | 18 ++++++++++++++---- 5 files changed, 49 insertions(+), 6 deletions(-) diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 5c586c063..b6ab3a57a 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/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) diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index 77596a103..fe6051f8b 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/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, diff --git a/vllm_omni/diffusion/envs.py b/vllm_omni/diffusion/envs.py index a71dc2e8e..ea7b2c24c 100644 --- a/vllm_omni/diffusion/envs.py +++ b/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 diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 65d37d0b0..61ddb4d84 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/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, diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 3bd520c61..fe1ccc6d0 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/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()