[diffusion][MUSA] fix: MUSA platform breakage caused by PR #13662 (#18456)

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

python/sglang/multimodal_gen/runtime/platforms/musa.py

@@ -18,6 +18,7 @@ import torchada  # noqa: F401
 # isort: on
 from typing_extensions import ParamSpec
 
+from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.platforms.interface import (
     AttentionBackendEnum,
     DeviceCapability,
@@ -70,6 +71,10 @@ class MusaPlatformBase(Platform):
     dispatch_key: str = "MUSA"
     device_control_env_var: str = "MUSA_VISIBLE_DEVICES"
 
+    @classmethod
+    def get_local_torch_device(cls) -> torch.device:
+        return torch.device(f"musa:{envs.LOCAL_RANK}")
+
     @classmethod
     def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
         raise NotImplementedError
@@ -120,6 +125,9 @@ class MusaPlatformBase(Platform):
         if empty_cache:
             torch.cuda.empty_cache()
 
+        if torch.distributed.is_initialized():
+            device_id = torch.distributed.get_rank()
+
         device_props = torch.cuda.get_device_properties(device_id)
         if device_props.is_integrated:
             free_gpu_memory = psutil.virtual_memory().available