[Hardware][MUSA] Complete Omni platform interfaces (#6058)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>

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)