[MUSA][4/N] Add common device utilities, distributed backend, and custom op wiring

python/sglang/multimodal_gen/runtime/layers/custom_op.py

@@ -53,6 +53,12 @@ def forward_tpu(self, *args, **kwargs) -> Any:
         # NOTE(woosuk): This is a placeholder for future extensions.
         return self.forward_native(*args, **kwargs)
 
+    def forward_musa(self, *args, **kwargs) -> Any:
+        # XXX (MUSA): MUSA kernels follow the CUDA path by default.
+        # At this stage, sgl-kernel support for MUSA is still under active
+        # development, so we fall back to the PyTorch-native implementation.
+        return self.forward_native(*args, **kwargs)
+
     def forward_oot(self, *args, **kwargs) -> Any:
         # By default, we assume that OOT ops are compatible with the
         # PyTorch-native implementation.
@@ -67,6 +73,8 @@ def dispatch_forward(self) -> Callable:
             return self.forward_npu
         elif current_platform.is_xpu():
             return self.forward_xpu
+        elif current_platform.is_musa():
+            return self.forward_musa
         else:
             return self.forward_native
 

python/sglang/srt/configs/device_config.py

@@ -11,7 +11,7 @@ class DeviceConfig:
     gpu_id: Optional[int]
 
     def __init__(self, device: str = "cuda", gpu_id: int = -1) -> None:
-        if device in ["cuda", "xpu", "hpu", "cpu", "npu"]:
+        if device in ["cuda", "xpu", "hpu", "cpu", "npu", "musa"]:
             self.device_type = device
         else:
             raise RuntimeError(f"Not supported device type: {device}")

python/sglang/srt/distributed/parallel_state.py

@@ -49,6 +49,7 @@
     is_cpu,
     is_cuda_alike,
     is_hip,
+    is_musa,
     is_npu,
     is_shm_available,
     is_xpu,
@@ -58,6 +59,7 @@
 _is_npu = is_npu()
 _is_cpu = is_cpu()
 _is_xpu = is_xpu()
+_is_musa = is_musa()
 
 TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])
 
@@ -247,6 +249,8 @@ def __init__(
             self.device = torch.device(f"npu:{local_rank}")
         elif _is_xpu:
             self.device = torch.device(f"xpu:{local_rank}")
+        elif _is_musa:
+            self.device = torch.device(f"musa:{local_rank}")
         else:
             self.device = torch.device("cpu")
         self.device_module = torch.get_device_module(self.device)
@@ -1903,6 +1907,8 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
             torch.xpu.empty_cache()
         elif hasattr(torch, "npu") and torch.npu.is_available():
             torch.npu.empty_cache()
+        elif hasattr(torch, "musa") and torch.musa.is_available():
+            torch.musa.empty_cache()
 
 
 def in_the_same_node_as(pg: ProcessGroup, source_rank: int = 0) -> List[bool]:

python/sglang/srt/layers/rotary_embedding.py

@@ -20,6 +20,7 @@
     is_cpu,
     is_cuda,
     is_hip,
+    is_musa,
     is_npu,
     is_xpu,
 )
@@ -31,6 +32,7 @@
 _is_cpu_amx_available = cpu_has_amx_support()
 _is_cpu = is_cpu()
 _is_xpu = is_xpu()
+_is_musa = is_musa()
 
 if _is_cuda:
     from sgl_kernel import FusedSetKVBufferArg, apply_rope_with_cos_sin_cache_inplace
@@ -119,6 +121,7 @@ def __init__(
             and not (_is_cpu)
             and not (_is_xpu)
             and not (_is_npu)
+            and not (_is_musa)
         ):
             if _is_cuda or _is_hip:
                 from sgl_kernel import rotary_embedding

python/sglang/srt/layers/utils/multi_platform.py

@@ -7,6 +7,7 @@
     is_cpu,
     is_cuda,
     is_hip,
+    is_musa,
     is_npu,
     is_xpu,
 )
@@ -17,6 +18,7 @@
 _is_cpu_amx_available = cpu_has_amx_support()
 _is_npu = is_npu()
 _is_xpu = is_xpu()
+_is_musa = is_musa()
 
 
 class MultiPlatformOp(nn.Module):
@@ -83,6 +85,12 @@ def forward_hip(self, *args, **kwargs):
     def forward_xpu(self, *args, **kwargs):
         return self.forward_native(*args, **kwargs)
 
+    def forward_musa(self, *args, **kwargs):
+        # XXX (MUSA): MUSA kernels follow the CUDA path by default.
+        # At this stage, sgl-kernel support for MUSA is still under active
+        # development, so we fall back to the PyTorch-native implementation.
+        return self.forward_native(*args, **kwargs)
+
     def forward_hpu(self, *args, **kwargs):
         return self.forward_native(*args, **kwargs)
 
@@ -100,5 +108,7 @@ def dispatch_forward(self):
             return self.forward_npu
         elif _is_xpu:
             return self.forward_xpu
+        elif _is_musa:
+            return self.forward_musa
         else:
             return self.forward_native

python/sglang/srt/model_executor/model_runner.py

@@ -575,7 +575,7 @@ def initialize(self, min_per_gpu_memory: float):
         # Init routed experts capturer
         self.init_routed_experts_capturer()
 
-        if self.device == "cuda":
+        if self.device == "cuda" or self.device == "musa":
             self.init_cublas()
             self.init_attention_backend()
             self.kernel_warmup()
@@ -732,6 +732,8 @@ def init_torch_distributed(self):
             backend = "gloo"
         elif self.device == "npu":
             backend = "hccl"
+        elif self.device == "musa":
+            backend = "mccl"
 
         before_avail_memory = get_available_gpu_memory(self.device, self.gpu_id)
         if not self.server_args.enable_p2p_check:
@@ -1090,7 +1092,7 @@ def model_load_weights(model, iter):
         self.server_args.load_format = load_format
         self.load_config = load_config
 
-        if recapture_cuda_graph and self.device == "cuda":
+        if recapture_cuda_graph and (self.device == "cuda" or self.device == "musa"):
             self.init_device_graphs()
 
         logger.info("Update weights end.")

python/sglang/srt/utils/common.py

@@ -519,8 +519,8 @@ def get_available_gpu_memory(
 
         if empty_cache:
             torch.cuda.empty_cache()
-        SHARED_SYSMEM_DEVICE_MEM_SMS = (87, 110, 121)  # Orin, Thor, Spark
-        if get_device_sm() in SHARED_SYSMEM_DEVICE_MEM_SMS:
+        props = torch.cuda.get_device_properties(gpu_id)
+        if props.is_integrated:
             # On these devices, which use sysmem as device mem, torch.cuda.mem_get_info()
             # only reports "free" memory, which can be lower than what is actually
             # available due to not including cache memory. So we use the system available
@@ -574,6 +574,25 @@ def get_available_gpu_memory(
         if empty_cache:
             torch.npu.empty_cache()
         free_gpu_memory, total_gpu_memory = torch.npu.mem_get_info()
+    elif device == "musa":
+        num_gpus = torch.musa.device_count()
+        assert gpu_id < num_gpus
+
+        if torch.musa.current_device() != gpu_id:
+            print(
+                f"WARNING: current device is not {gpu_id}, but {torch.musa.current_device()}, ",
+                "which may cause useless memory allocation for torch MUSA context.",
+            )
+        if empty_cache:
+            torch.musa.empty_cache()
+        props = torch.musa.get_device_properties(gpu_id)
+        if props.is_integrated:
+            # On these devices, which use sysmem as device mem, torch.musa.mem_get_info()
+            # only reports "free" memory, which can be lower than what is actually
+            # available due to not including cache memory. So we use the system available
+            # memory metric instead.
+            free_gpu_memory = psutil.virtual_memory().available
+        free_gpu_memory, total_gpu_memory = torch.musa.mem_get_info()
 
     if distributed:
         tensor = torch.tensor(free_gpu_memory, dtype=torch.float32)
@@ -1227,7 +1246,9 @@ def broadcast_pyobj(
     of dist_group argument).
     """
     device = torch.device(
-        "cuda" if torch.cuda.is_available() and not force_cpu_device else "cpu"
+        "cuda"
+        if torch.cuda.is_available() and not force_cpu_device
+        else "musa" if is_musa() and not force_cpu_device else "cpu"
     )
 
     if rank == src:
@@ -1788,6 +1809,47 @@ def get_xpu_memory_capacity():
         raise RuntimeError("torch.xpu is not available.")
 
 
+def get_mtgpu_memory_capacity():
+    try:
+        # Run mthreads-gmi and capture the output
+        result = subprocess.run(
+            [
+                "mthreads-gmi --query | grep 'FB Memory Usage' -A 2 | grep 'Total' | awk -F':' '{print $2}' | awk '{print $1}' | sed 's/MiB//'"
+            ],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            shell=True,
+            text=True,
+        )
+
+        if result.returncode != 0:
+            raise RuntimeError(f"mthreads-gmi error: {result.stderr.strip()}")
+
+        # Parse the output to extract memory values
+        memory_values = [
+            float(mem)
+            for mem in result.stdout.strip().split("\n")
+            if re.match(r"^\d+(\.\d+)?$", mem.strip())
+        ]
+
+        if not memory_values:
+            # Fallback to torch.musa.mem_get_info() when failed to get memory capacity from mthreads-gmi.
+            if hasattr(torch, "musa") and torch.musa.is_available():
+                logger.warning(
+                    "Failed to get GPU memory capacity from mthreads-gmi, falling back to torch.musa.mem_get_info()."
+                )
+                return torch.musa.mem_get_info()[1] // 1024 // 1024  # unit: MB
+            raise ValueError("No GPU memory values found.")
+
+        # Return the minimum memory value
+        return min(memory_values)
+
+    except FileNotFoundError:
+        raise RuntimeError(
+            "mthreads-gmi not found. Ensure Moore Threads drivers are installed and accessible."
+        )
+
+
 def get_device_memory_capacity(device: str = None):
     if is_cuda():
         gpu_mem = get_nvgpu_memory_capacity()
@@ -1801,6 +1863,8 @@ def get_device_memory_capacity(device: str = None):
         gpu_mem = get_cpu_memory_capacity()
     elif device == "xpu":
         gpu_mem = get_xpu_memory_capacity()
+    elif device == "musa":
+        gpu_mem = get_mtgpu_memory_capacity()
     else:
         # GPU memory is not known yet or no GPU is available.
         gpu_mem = None
@@ -1899,7 +1963,7 @@ def print_info_once(msg: str) -> None:
 
 
 def get_device_name(device_id: int = 0) -> str:
-    if hasattr(torch, "cuda") and torch.cuda.is_available():
+    if (hasattr(torch, "cuda") and torch.cuda.is_available()) or is_musa():
         return torch.cuda.get_device_name(device_id)
 
     if hasattr(torch, "xpu") and torch.xpu.is_available():
@@ -1956,12 +2020,17 @@ def get_device(device_id: Optional[int] = None) -> str:
                 "Habana frameworks detected, but failed to import 'habana_frameworks.torch.hpu'."
             )
 
-    raise RuntimeError("No accelerator (CUDA, XPU, HPU, NPU) is available.")
+    if is_musa():
+        if device_id == None:
+            return "musa"
+        return "musa:{}".format(device_id)
+
+    raise RuntimeError("No accelerator (CUDA, XPU, HPU, NPU, MUSA) is available.")
 
 
 @lru_cache(maxsize=1)
 def get_device_count() -> int:
-    if hasattr(torch, "cuda") and torch.cuda.is_available():
+    if (hasattr(torch, "cuda") and torch.cuda.is_available()) or is_musa():
         try:
             return torch.cuda.device_count()
         except RuntimeError:
@@ -1986,15 +2055,15 @@ def get_device_count() -> int:
 
 
 def get_device_core_count(device_id: int = 0) -> int:
-    if hasattr(torch, "cuda") and torch.cuda.is_available():
+    if (hasattr(torch, "cuda") and torch.cuda.is_available()) or is_musa():
         return torch.cuda.get_device_properties(device_id).multi_processor_count
 
     return 0
 
 
 def get_device_capability(device_id: int = 0) -> Tuple[int, int]:
     major, minor = None, None
-    if hasattr(torch, "cuda") and torch.cuda.is_available():
+    if (hasattr(torch, "cuda") and torch.cuda.is_available()) or is_musa():
         major, minor = torch.cuda.get_device_capability(device_id)
 
     if hasattr(torch, "xpu") and torch.xpu.is_available():