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

[MUSA][Feature] Enable Piecewise CUDA Graph support for MUSA platform (#20758)

Browse Source 
Signed-off-by: yafeng.li <yafeng.li@mthreads.com>
Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
This commit is contained in:
MARATRIX
2026-03-25 21:00:28 -07:00
committed by GitHub
co-authored by R0CKSTAR
parent abf4f1a47a
commit f420b9b4a5
2 changed files with 4 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/compilation/weak_ref_tensor.py
+2 -2
View File
@@ -2,9 +2,9 @@ from typing import Any, Union
import torch
from sglang.srt.utils.common import is_cuda, is_hip, is_npu
from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu
if is_cuda() or is_hip():
if is_cuda() or is_hip() or is_musa():
from sgl_kernel import weak_ref_tensor
elif is_npu():
from torch_npu._C import _weak_ref_tensor as weak_ref_tensor
python/sglang/srt/server_args.py
+2 -3
View File
@@ -51,7 +51,6 @@ from sglang.srt.utils.common import (
is_hip,
is_hopper_with_cuda_12_3,
is_mps,
is_musa,
is_no_spec_infer_or_topk_one,
is_npu,
is_remote_url,
@@ -1075,8 +1074,8 @@ class ServerArgs:
# 5. Pipeline parallelism
if self.pp_size > 1:
self.disable_piecewise_cuda_graph = True
# 6. Non-CUDA hardware (AMD, NPU, CPU, MPS, MUSA, XPU, etc.)
if is_hip() or is_npu() or is_cpu() or is_mps() or is_musa() or is_xpu():
# 6. Non-CUDA hardware (AMD, NPU, CPU, MPS, XPU, etc.)
if is_hip() or is_npu() or is_cpu() or is_mps() or is_xpu():
self.disable_piecewise_cuda_graph = True
# 7. MoE A2A backend
if self.moe_a2a_backend != "none":