[MUSA][7/N] Enhance CUDA / PyNccl wrapper to support MTLink connectivity detection (#17499)

Signed-off-by: jingzhi.xue <jingzhi.xue@mthreads.com>
Co-authored-by: jingzhi.xue <jingzhi.xue@mthreads.com>

python/pyproject_other.toml

@@ -108,7 +108,7 @@ srt_musa = [
     "sglang[runtime_common]",
     "torch",
     "torch_musa",
-    "torchada>=0.1.15",
+    "torchada>=0.1.25",
     "mthreads-ml-py",
     "numpy<2.0",
 ]

python/sglang/srt/distributed/device_communicators/cuda_wrapper.py

@@ -13,6 +13,10 @@ from typing import Any, Dict, List, Optional
 # this line makes it possible to directly load `libcudart.so` using `ctypes`
 import torch  # noqa
 
+from sglang.srt.utils import is_musa
+
+_is_musa = is_musa()
+
 logger = logging.getLogger(__name__)
 
 # === export types and functions from cudart to Python ===
@@ -113,7 +117,7 @@ class CudaRTLibrary:
 
     def __init__(self, so_file: Optional[str] = None):
         if so_file is None:
-            so_file = find_loaded_library("libcudart")
+            so_file = find_loaded_library("libcudart" if not _is_musa else "libmusart")
             assert so_file is not None, "libcudart is not loaded in the current process"
         if so_file not in CudaRTLibrary.path_to_library_cache:
             lib = ctypes.CDLL(so_file)

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -19,10 +19,17 @@ from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import
 )
 from sglang.srt.distributed.parallel_state import in_the_same_node_as
 from sglang.srt.environ import envs
-from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, log_info_on_rank0
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_cuda,
+    is_hip,
+    is_musa,
+    log_info_on_rank0,
+)
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
+_is_musa = is_musa()
 
 logger = logging.getLogger(__name__)
 
@@ -47,6 +54,9 @@ class CustomAllreduce:
     if _is_hip:
         # crossover is at 16MB buffer size for ROCm
         _MAX_CAR_SIZE = 2 * 8192 * 1024
+    if _is_musa:
+        # crossover is at 128MB buffer size for MUSA
+        _MAX_CAR_SIZE = 16 * 8196 * 1024
 
     # max_size: max supported allreduce size
     def __init__(
@@ -129,7 +139,7 @@ class CustomAllreduce:
         # test nvlink first, this will filter out most of the cases
         # where custom allreduce is not supported
         # this checks hardware and driver support for NVLink
-        if _is_cuda or _is_hip:
+        if _is_cuda or _is_hip or _is_musa:
             full_nvlink = is_full_nvlink(physical_device_ids, world_size)
 
         if world_size > 2 and not full_nvlink:
@@ -210,6 +220,8 @@ class CustomAllreduce:
         """
         lib = CudaRTLibrary()
         pointer = lib.cudaMalloc(size_in_bytes)
+        if _is_musa:
+            lib.cudaMemset(pointer, 0, size_in_bytes)
         handle = lib.cudaIpcGetMemHandle(pointer)
         world_size = dist.get_world_size(group=group)
         rank = dist.get_rank(group=group)
@@ -433,7 +445,7 @@ def dispatch_custom_allreduce():
     On AMD with 1-stage AR enabled, use sglang's CustomAllreduce (has deterministic_all_reduce method).
     Otherwise use AiterCustomAllreduce if available.
     """
-    if _is_cuda:
+    if _is_cuda or _is_musa:
         return CustomAllreduce
 
     assert _is_hip

python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py

@@ -4,14 +4,15 @@ from typing import List, Optional, Tuple
 
 import torch
 
-from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils import is_cuda, is_hip, is_musa
 
 logger = logging.getLogger(__name__)
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
+_is_musa = is_musa()
 
-IS_CUSTOM_AR_AVAILABLE = _is_cuda or _is_hip
+IS_CUSTOM_AR_AVAILABLE = _is_cuda or _is_hip or _is_musa
 IS_QUICK_AR_AVAILABLE = _is_hip
 # TODO(zyksir): mscclpp is untested on AMD and therefore disabled.
 IS_MSCCLPP_AR_AVAILABLE = _is_cuda
@@ -30,7 +31,7 @@ except ImportError as e:
 if not IS_CUSTOM_AR_AVAILABLE:
     pass
 
-elif _is_cuda:
+elif _is_cuda or _is_musa:
     # CUDA custom allreduce
 
     def init_custom_ar(

python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py

@@ -18,12 +18,13 @@ import torch.multiprocessing as mp
 from typing_extensions import ParamSpec
 
 from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
-from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils import is_cuda, is_hip, is_musa
 
 logger = logging.getLogger(__name__)
 
 _is_cuda = is_cuda()
 _is_hip = is_hip()
+_is_musa = is_musa()
 
 if _is_cuda:
     try:
@@ -31,6 +32,12 @@ if _is_cuda:
     except ImportError as e:
         logger.warning("Failed to import pynvml with %r", e)
 
+if _is_musa:
+    try:
+        import pymtml as pynvml
+    except ImportError as e:
+        logger.warning("Failed to import pymtml with %r", e)
+
 if _is_hip:
     try:
         from amdsmi import (

python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py

@@ -38,8 +38,8 @@ def find_nccl_library() -> str:
     """
     We either use the library file specified by the `SGLANG_NCCL_SO_PATH`
     environment variable, or we find the library file brought by PyTorch.
-    After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be
+    After importing `torch`, `libnccl.so.2`, `librccl.so.1` or `libmccl.so.2`
-    found by `ctypes` automatically.
+    can be found by `ctypes` automatically.
     """
 
     # so_file can be set to None in sglang
@@ -55,8 +55,10 @@ def find_nccl_library() -> str:
             so_file = "libnccl.so.2"
         elif torch.version.hip is not None:
             so_file = "librccl.so.1"
+        elif hasattr(torch.version, "musa") and torch.version.musa is not None:
+            so_file = "libmccl.so.2"
         else:
-            raise ValueError("NCCL only supports CUDA and ROCm backends.")
+            raise ValueError("NCCL only supports CUDA, ROCm and MUSA backends.")
         logger.debug("Found nccl from library %s", so_file)
     return so_file
 
@@ -341,7 +343,7 @@ class NCCLLibrary:
         except Exception as e:
             logger.error(
                 "Failed to load NCCL library from %s . "
-                "It is expected if you are not running on NVIDIA/AMD GPUs. "
+                "It is expected if you are not running on NVIDIA/AMD/MTHREADS GPUs. "
                 "Otherwise, the nccl library might not exist, be corrupted "
                 "or it does not support the current platform %s. "
                 "If you already have the library, please set the "