From 22dd96ec5c0fa25667ef186c31e199ba7809ae50 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Tue, 11 Aug 2026 20:46:42 +0800 Subject: [PATCH] [Hardware][MUSA] Complete Omni platform interfaces (#6058) Signed-off-by: Xiaodong Ye --- vllm_omni/platforms/musa/platform.py | 35 ++++++++++++++++++++++++---- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index df52905c9..d667c9bf1 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -3,6 +3,8 @@ import torch +from vllm.config import VllmConfig +from vllm.config.kernel import IrOpPriorityConfig from vllm.logger import init_logger from vllm.platforms.interface import DeviceCapability from vllm_musa.platform import MUSAPlatformBase @@ -30,10 +32,6 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): def get_omni_generation_worker_cls(cls) -> str: return "vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker" - @classmethod - def get_default_stage_config_path(cls) -> str: - return "vllm_omni/deploy" - @classmethod def has_flash_attn_package(cls) -> bool: from vllm_omni.diffusion.attention.backends.utils.fa import is_flash_attn_installed @@ -115,6 +113,10 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): """MUSA supports torch.compile with inductor backend.""" return True + @classmethod + def get_default_stage_config_path(cls) -> str: + return "vllm_omni/deploy" + @classmethod def supports_talker_mtp_graph_capture(cls) -> bool: """MUSA keeps Qwen3 talker MTP outside its dedicated FULL graph.""" @@ -142,7 +144,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): @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) + major, minor = torch.musa.get_device_capability(device_id) return DeviceCapability(major=major, minor=minor) @classmethod @@ -183,3 +185,26 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): @classmethod def get_device_name(cls, device_id: int = 0) -> str: return torch.musa.get_device_name(device_id) + + @classmethod + def set_device_control_env_var(cls, devices: str | int | None) -> None: + import os + + os.environ["MUSA_VISIBLE_DEVICES"] = devices + + @classmethod + def unset_device_control_env_var(cls) -> None: + import os + + os.environ.pop("MUSA_VISIBLE_DEVICES", None) + + @classmethod + def get_default_ir_op_priority(cls, vllm_config: VllmConfig) -> IrOpPriorityConfig: + """Prefer native ops while compiling, otherwise prefer vLLM kernels.""" + from vllm.config.compilation import CompilationMode + + cc = vllm_config.compilation_config + using_inductor = cc.backend == "inductor" and cc.mode != CompilationMode.NONE + default = ["native"] if using_inductor else ["vllm_c", "native"] + + return IrOpPriorityConfig.with_default(default)