diff --git a/python/pyproject_other.toml b/python/pyproject_other.toml index e1312a46e..0c6d842cd 100755 --- a/python/pyproject_other.toml +++ b/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", ] diff --git a/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py b/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py index c902f3141..f1dc99bd2 100644 --- a/python/sglang/srt/distributed/device_communicators/cuda_wrapper.py +++ b/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) diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index e71f93ebc..6e7ea07e7 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/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 diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py index c312e8f2c..55879e55f 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_ops.py +++ b/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( diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py index c7baac845..c74cd58d6 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py +++ b/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 ( diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py b/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py index 6b12f2922..59775ad19 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_wrapper.py +++ b/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 - found by `ctypes` automatically. + After importing `torch`, `libnccl.so.2`, `librccl.so.1` or `libmccl.so.2` + 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 "