diff --git a/3rdparty/amd/wheel/sglang/pyproject.toml b/3rdparty/amd/wheel/sglang/pyproject.toml index a96f7049d..bac322ab3 100644 --- a/3rdparty/amd/wheel/sglang/pyproject.toml +++ b/3rdparty/amd/wheel/sglang/pyproject.toml @@ -123,7 +123,7 @@ srt_musa = [ "sglang[runtime_common]", "torch", "torch_musa", - "torchada>=0.1.25", + "torchada>=0.1.45", "mthreads-ml-py", "numpy<2.0", ] diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index 4348f673d..ba19cbc13 100755 --- a/python/pyproject_other.toml +++ b/python/pyproject_other.toml @@ -114,7 +114,7 @@ srt_musa = [ "sglang[runtime_common]", "torch", "torch_musa", - "torchada>=0.1.25", + "torchada>=0.1.45", "mthreads-ml-py", "numpy<2.0", ] diff --git a/python/sglang/multimodal_gen/runtime/platforms/musa.py b/python/sglang/multimodal_gen/runtime/platforms/musa.py index 7d443be6b..a368be696 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/musa.py +++ b/python/sglang/multimodal_gen/runtime/platforms/musa.py @@ -150,10 +150,61 @@ class MusaPlatformBase(Platform): head_size: int, dtype: torch.dtype, ) -> str: - logger.info("Using Torch SDPA backend.") - return ( - "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" - ) + target_backend: AttentionBackendEnum | None = None + + if selected_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend") + return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" + elif selected_backend in [ + AttentionBackendEnum.FA, + ]: + target_backend = AttentionBackendEnum.FA + elif selected_backend: + raise ValueError(f"Invalid attention backend for {cls.device_name}") + else: + target_backend = AttentionBackendEnum.FA + + # Ensure we have a target backend selected before validation/fallback. + if target_backend is None: + target_backend = AttentionBackendEnum.FA + + if dtype not in (torch.float16, torch.bfloat16): + logger.info( + "Cannot use FlashAttention backend for dtype other than " + "torch.float16 or torch.bfloat16." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + # FlashAttn is valid for the model, checking if the package is + # installed. + if target_backend == AttentionBackendEnum.FA: + try: + from sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn import ( # noqa: F401 + FlashAttentionBackend, + ) + + supported_sizes = FlashAttentionBackend.get_supported_head_sizes() + if head_size not in supported_sizes: + logger.info( + "Cannot use FlashAttention backend for head size %d.", + head_size, + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + except ImportError: + logger.info( + "Cannot use FlashAttention backend because the " + "flash_attn package is not found. " + "Make sure that flash_attn was built and installed " + "(on by default)." + ) + target_backend = AttentionBackendEnum.TORCH_SDPA + + if target_backend == AttentionBackendEnum.TORCH_SDPA: + logger.info("Using Torch SDPA backend") + return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" + + logger.info("Using FlashAttention (FA3) backend on MUSA") + return "sglang.multimodal_gen.runtime.layers.attention.backends.flash_attn.FlashAttentionBackend" @classmethod def get_device_communicator_cls(cls) -> str: diff --git a/sgl-kernel/pyproject_musa.toml b/sgl-kernel/pyproject_musa.toml index 669927be3..8120e7de6 100644 --- a/sgl-kernel/pyproject_musa.toml +++ b/sgl-kernel/pyproject_musa.toml @@ -3,7 +3,7 @@ requires = [ "setuptools>=75.0", "scikit-build-core>=0.10", "torch", - "torchada>=0.1.14", + "torchada>=0.1.45", "wheel", ] build-backend = "setuptools.build_meta"