chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
R0CKSTAR
2026-01-29 15:13:24 +08:00
committed by GitHub
parent 9409c43593
commit d3cdee0a04
7 changed files with 109 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/multimodal_gen/runtime/layers/custom_op.py
View File

@@ -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

2
python/sglang/srt/configs/device_config.py
View File

@@ -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}")

6
python/sglang/srt/distributed/parallel_state.py
View File

@@ -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]:

3
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -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

10
python/sglang/srt/layers/utils/multi_platform.py
View File

@@ -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

6
python/sglang/srt/model_executor/model_runner.py
View File

@@ -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.")

85
python/sglang/srt/utils/common.py
View File

@@ -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():