diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index 5fe3fa4bf..acb27c468 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -1,10 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from typing import Any import torch from vllm.logger import init_logger +from vllm.platforms.interface import DeviceCapability from vllm_musa.platform import MUSAPlatformBase from vllm_omni.diffusion.attention.backends.registry import DiffusionAttentionBackendEnum @@ -34,18 +34,6 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): def get_default_stage_config_path(cls) -> str: return "vllm_omni/model_executor/stage_configs" - @classmethod - def get_diffusion_model_impl_qualname(cls, op_name: str) -> str: - # MUSA uses default implementations for diffusion ops - if op_name == "hunyuan_fused_moe": - return "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault" - return super().get_diffusion_model_impl_qualname(op_name) - - @classmethod - def prepare_diffusion_op_runtime(cls, op_name: str, **kwargs: Any) -> None: - # 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 @@ -135,11 +123,22 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): return torch.device("musa") return torch.device("musa", local_rank) + @classmethod + def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None: + """Get the compute capability of the MUSA device.""" + major, minor = torch.cuda.get_device_capability(device_id) + return DeviceCapability(major=major, minor=minor) + @classmethod def get_device_count(cls) -> int: """Get the number of available MUSA devices.""" return torch.musa.device_count() + @classmethod + def get_device_version(cls) -> str | None: + """Get the MUSA runtime version.""" + return torch.version.musa + @classmethod def synchronize(cls) -> None: """Synchronize all MUSA operations."""