From 60f6f034090946e0c4832ce62ee3a7364bb55321 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Bernl=C3=B6hr?= <jbernloehr@nvidia.com>
Date: Mon, 14 Sep 2026 06:23:51 +0200
Subject: [PATCH] Fix MUSA detection in compiled prefill path (#39061)

---
 .../sglang/srt/layers/attention/dsa/utils.py   |  5 +++--
 python/sglang/srt/utils/common.py              | 14 +++++++++-----
 test/registered/cp/test_cp_strategy_unit.py    | 18 ++++++++++++++++++
 test/registered/unit/utils/test_common.py      | 15 +++++++++++++++
 4 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py
index db99b5b20..f685334ab 100644
--- a/python/sglang/srt/layers/attention/dsa/utils.py
+++ b/python/sglang/srt/layers/attention/dsa/utils.py
@@ -115,13 +115,14 @@ def should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend: bool) -> bool:
 
 
 def is_dsa_enable_prefill_cp():
+    if get_parallel().attn_cp_size <= 1:
+        return False
+
     if is_hip() or is_npu() or is_musa():
         return False
 
     # Generic prefill CP derives activation from the runtime topology and model
     # architecture.
-    if get_parallel().attn_cp_size <= 1:
-        return False
     from sglang.srt.configs.model_config import is_deepseek_dsa, is_deepseek_v4
 
     hf_config = process_model_config().hf_config
diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py
index 62ad5159f..cb6efd8da 100644
--- a/python/sglang/srt/utils/common.py
+++ b/python/sglang/srt/utils/common.py
@@ -221,13 +221,17 @@ def is_cpu() -> bool:
     return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" and is_host_cpu_supported
 
 
+try:
+    import torchada  # noqa: F401
+except ImportError:
+    _IS_MUSA = False
+else:
+    _IS_MUSA = hasattr(torch.version, "musa") and torch.version.musa is not None
+
+
 @lru_cache(maxsize=1)
 def is_musa() -> bool:
-    try:
-        import torchada  # noqa: F401
-    except ImportError:
-        return False
-    return hasattr(torch.version, "musa") and torch.version.musa is not None
+    return _IS_MUSA
 
 
 @lru_cache(maxsize=1)
diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py
index 1bc6a537d..23b652bd6 100644
--- a/test/registered/cp/test_cp_strategy_unit.py
+++ b/test/registered/cp/test_cp_strategy_unit.py
@@ -121,6 +121,24 @@ class TestCPStrategyUnit(CustomTestCase):
         ):
             self.assertFalse(is_dsa_enable_prefill_cp())
 
+    def test_disabled_dsa_cp_skips_platform_probes(self):
+        parallel = SimpleNamespace(attn_cp_size=1)
+
+        with (
+            patch(
+                "sglang.srt.layers.attention.dsa.utils.get_parallel",
+                return_value=parallel,
+            ),
+            patch("sglang.srt.layers.attention.dsa.utils.is_hip") as mock_is_hip,
+            patch("sglang.srt.layers.attention.dsa.utils.is_npu") as mock_is_npu,
+            patch("sglang.srt.layers.attention.dsa.utils.is_musa") as mock_is_musa,
+        ):
+            self.assertFalse(is_dsa_enable_prefill_cp())
+
+        mock_is_hip.assert_not_called()
+        mock_is_npu.assert_not_called()
+        mock_is_musa.assert_not_called()
+
 
 class TestPrefillCPBCGReplay(CustomTestCase):
     def tearDown(self):
diff --git a/test/registered/unit/utils/test_common.py b/test/registered/unit/utils/test_common.py
index f93a65ecd..4d772b9a0 100644
--- a/test/registered/unit/utils/test_common.py
+++ b/test/registered/unit/utils/test_common.py
@@ -7,6 +7,7 @@ from sglang.srt.utils.common import (
     flatten_arrays_to_int64_tensor,
     get_device_sm_nvidia_smi,
     get_nvidia_driver_version_str,
+    is_musa,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
@@ -15,6 +16,20 @@ register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
 register_amd_ci(est_time=5, stage="stage-b", runner_config="1-gpu-small-amd")
 
 
+class TestMusaDetection(CustomTestCase):
+    def test_is_musa_is_torch_compile_safe(self):
+        is_musa.cache_clear()
+
+        @torch.compile(backend="eager", fullgraph=True)
+        def add_platform_offset(value):
+            return value + 1 if is_musa() else value - 1
+
+        value = torch.zeros(1)
+        actual = add_platform_offset(value)
+        expected = torch.ones(1) if is_musa() else -torch.ones(1)
+        torch.testing.assert_close(actual, expected)
+
+
 @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
 class TestFlattenArraysToInt64Tensor(CustomTestCase):
     """`flatten_arrays_to_int64_tensor` is invoked by `prepare_for_extend`