Fix MUSA detection in compiled prefill path (#39061)

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

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)

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

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`