From 0c4f98ed4e325e4db4ca97bb5c012137a97db9e4 Mon Sep 17 00:00:00 2001
From: R0CKSTAR <xiaodong.ye@mthreads.com>
Date: Sat, 7 Mar 2026 08:14:41 +0800
Subject: [PATCH] [diffusion] hardware: add set_musa_arch on MUSA (misc, 15/N)
 (#19381)

Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
---
 .../sglang/multimodal_gen/runtime/managers/gpu_worker.py | 4 +++-
 python/sglang/multimodal_gen/runtime/utils/common.py     | 9 +++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
index ca51f21fb..7b357330f 100644
--- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
@@ -43,7 +43,7 @@ from sglang.multimodal_gen.runtime.pipelines_core import (
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
-from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch
+from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch
 from sglang.multimodal_gen.runtime.utils.layerwise_offload import (
     OffloadableDiTMixin,
     iter_materialized_weights,
@@ -474,6 +474,8 @@ def run_scheduler_process(
     globally_suppress_loggers()
     if current_platform.is_cuda():
         set_cuda_arch()
+    elif current_platform.is_musa():
+        set_musa_arch()
 
     port_args = PortArgs.from_server_args(server_args)
 
diff --git a/python/sglang/multimodal_gen/runtime/utils/common.py b/python/sglang/multimodal_gen/runtime/utils/common.py
index f9f5df924..b5e3dc098 100644
--- a/python/sglang/multimodal_gen/runtime/utils/common.py
+++ b/python/sglang/multimodal_gen/runtime/utils/common.py
@@ -261,6 +261,15 @@ def set_cuda_arch():
     os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}"
 
 
+# musa
+
+
+def set_musa_arch():
+    capability = torch.cuda.get_device_capability()
+    arch = f"{capability[0]}{capability[1]}"
+    os.environ["TORCH_MUSA_ARCH_LIST"] = f"{arch}"
+
+
 # env var managements
 
 _warned_bool_env_var_keys = set()