From d3cdee0a040bebf0bf651e7c9d13449741731a50 Mon Sep 17 00:00:00 2001 From: R0CKSTAR Date: Thu, 29 Jan 2026 15:13:24 +0800 Subject: [PATCH] [MUSA][4/N] Add common device utilities, distributed backend, and custom op wiring (#17246) Signed-off-by: Xiaodong Ye --- .../runtime/layers/custom_op.py | 8 ++ python/sglang/srt/configs/device_config.py | 2 +- .../sglang/srt/distributed/parallel_state.py | 6 ++ python/sglang/srt/layers/rotary_embedding.py | 3 + .../sglang/srt/layers/utils/multi_platform.py | 10 +++ .../sglang/srt/model_executor/model_runner.py | 6 +- python/sglang/srt/utils/common.py | 85 +++++++++++++++++-- 7 files changed, 109 insertions(+), 11 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/python/sglang/multimodal_gen/runtime/layers/custom_op.py index f6c797c64..373261183 100644 --- a/python/sglang/multimodal_gen/runtime/layers/custom_op.py +++ b/python/sglang/multimodal_gen/runtime/layers/custom_op.py @@ -53,6 +53,12 @@ class CustomOp(nn.Module): # 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 @@ class CustomOp(nn.Module): 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 diff --git a/python/sglang/srt/configs/device_config.py b/python/sglang/srt/configs/device_config.py index 20b9af9be..8ddcfd108 100644 --- a/python/sglang/srt/configs/device_config.py +++ b/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}") diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index b01595526..3070178b6 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -49,6 +49,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda_alike, is_hip, + is_musa, is_npu, is_shm_available, is_xpu, @@ -58,6 +59,7 @@ from sglang.srt.utils.custom_op import register_custom_op _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 @@ class GroupCoordinator: 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]: diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index b38c980f1..f85031ac4 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -20,6 +20,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_musa, is_npu, is_xpu, ) @@ -31,6 +32,7 @@ _is_npu = is_npu() _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 @@ class RotaryEmbedding(MultiPlatformOp): 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 diff --git a/python/sglang/srt/layers/utils/multi_platform.py b/python/sglang/srt/layers/utils/multi_platform.py index 33033d4c8..954728edb 100644 --- a/python/sglang/srt/layers/utils/multi_platform.py +++ b/python/sglang/srt/layers/utils/multi_platform.py @@ -7,6 +7,7 @@ from sglang.srt.utils import ( is_cpu, is_cuda, is_hip, + is_musa, is_npu, is_xpu, ) @@ -17,6 +18,7 @@ _is_cpu = is_cpu() _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 @@ class MultiPlatformOp(nn.Module): 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 @@ class MultiPlatformOp(nn.Module): return self.forward_npu elif _is_xpu: return self.forward_xpu + elif _is_musa: + return self.forward_musa else: return self.forward_native diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 4490956c2..4dcf981fa 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -575,7 +575,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): # 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 @@ class ModelRunner(ModelRunnerKVCacheMixin): 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 @@ class ModelRunner(ModelRunnerKVCacheMixin): 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.") diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 2b560fcc3..17a3360d2 100644 --- a/python/sglang/srt/utils/common.py +++ b/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,7 +2055,7 @@ 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 @@ -1994,7 +2063,7 @@ def get_device_core_count(device_id: int = 0) -> int: 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():