[ROCm] Define the DSA head-gate graph helpers on HIP (#37118)

Co-authored-by: Zhang, Jiejing <jiejing.zhang@amd.com>

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -132,16 +132,19 @@ if TYPE_CHECKING:
 DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0
 
 
+if _is_cuda or _is_hip:
+    # Plain-torch graph helpers: usable wherever the split-op surface is.
+    from sglang.srt.layers.attention.dsa.dsa_prefill_cuda_graph import (
+        logits_head_gate_graph,
+        scale_head_gate_graph,
+    )
+
 if _is_cuda:
     from sglang.kernels.ops.attention.dsv4 import fused_q_indexer_rope_first_quant
     from sglang.kernels.ops.quantization.dsv32 import (
         fused_k_indexer_norm_rope,
         fused_k_indexer_norm_rope_store,
     )
-    from sglang.srt.layers.attention.dsa.dsa_prefill_cuda_graph import (
-        logits_head_gate_graph,
-        scale_head_gate_graph,
-    )
 
     @register_custom_op(mutates_args=["topk_indices"])
     @register_split_op()

python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py

@@ -14,10 +14,11 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
     get_tc_piecewise_forward_context,
     is_in_tc_piecewise_cuda_graph,
 )
-from sglang.srt.utils import is_cuda
+from sglang.srt.utils import is_cuda, is_hip
 from sglang.srt.utils.custom_op import register_custom_op
 
 _is_cuda = is_cuda()
+_is_hip = is_hip()
 
 GRAPH_WEIGHTS_PROJ_LORA_ERROR = (
     "DSA indexer weights_proj LoRA is incompatible with "
@@ -31,7 +32,7 @@ def _is_in_piecewise_or_breakable_cuda_graph() -> bool:
     return is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph()
 
 
-if _is_cuda:
+if _is_cuda or _is_hip:
 
     def _scale_head_gate_graph_fake_impl(
         weights_raw: torch.Tensor,

test/registered/unit/layers/attention/test_dsa_head_gate_guard.py

@@ -0,0 +1,42 @@
+import unittest
+
+from sglang.srt.utils import is_cuda, is_hip
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+HELPERS = ("logits_head_gate_graph", "scale_head_gate_graph")
+
+
+class TestDsaHeadGateGuard(CustomTestCase):
+    """The head-gate helpers are defined in one module and imported in another.
+
+    Both sides carry their own platform condition, so the two can drift: widen
+    the import without widening the definition and every DSA model fails to
+    import on the platform in between. That is a startup failure far from its
+    cause, and no CUDA runner can see it.
+    """
+
+    def test_definition_and_import_agree(self):
+        from sglang.srt.layers.attention.dsa import dsa_indexer, dsa_prefill_cuda_graph
+
+        for name in HELPERS:
+            with self.subTest(helper=name):
+                self.assertEqual(
+                    hasattr(dsa_prefill_cuda_graph, name),
+                    hasattr(dsa_indexer, name),
+                    f"{name} is defined on one side of the guard but not the other",
+                )
+
+    def test_helpers_exist_on_cuda_and_hip(self):
+        from sglang.srt.layers.attention.dsa import dsa_prefill_cuda_graph
+
+        expected = is_cuda() or is_hip()
+        for name in HELPERS:
+            with self.subTest(helper=name):
+                self.assertEqual(hasattr(dsa_prefill_cuda_graph, name), expected)
+
+
+if __name__ == "__main__":
+    unittest.main()