From 030d7e7e9b073667ce2f6adc8293363a5a0ff149 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Thu, 3 Sep 2026 13:55:45 +0800 Subject: [PATCH] [ROCm] Define the DSA head-gate graph helpers on HIP (#37118) Co-authored-by: Zhang, Jiejing --- .../srt/layers/attention/dsa/dsa_indexer.py | 11 +++-- .../attention/dsa/dsa_prefill_cuda_graph.py | 5 ++- .../attention/test_dsa_head_gate_guard.py | 42 +++++++++++++++++++ 3 files changed, 52 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/layers/attention/test_dsa_head_gate_guard.py diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 3969dfbee..37b1bfa4e 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/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() diff --git a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py b/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py index f90e7ba95..92a661a18 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py +++ b/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, diff --git a/test/registered/unit/layers/attention/test_dsa_head_gate_guard.py b/test/registered/unit/layers/attention/test_dsa_head_gate_guard.py new file mode 100644 index 000000000..3565ff5db --- /dev/null +++ b/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()