[AMD][Spec][PD] Enable the PD DSA fused-TopK seed remap on ROCm (#36714)

Co-authored-by: Tianxiao Jiang <tianxiaojiang@Tianxiaos-MacBook-Pro.local>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

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

@@ -78,7 +78,7 @@ def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int):
 def should_remap_pd_dsa_seed_to_local_slots() -> bool:
     """Whether a PD seed should enter the allocator-local fused TopK domain."""
     return (
-        is_cuda()
+        (is_cuda() or is_hip())
         and envs.SGLANG_DSA_FUSE_TOPK.get()
         and get_disagg().disaggregation_mode == "decode"
         and not get_memory().enable_hisparse

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -399,20 +399,32 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
         override.install()
         self.addCleanup(override.restore)
 
-        with envs.SGLANG_DSA_FUSE_TOPK.override(True), patch(
-            "sglang.srt.layers.attention.dsa.utils.is_cuda", return_value=True
+        local_slots = [[309, 101, -1], [801, 990, -1]]
+        unremapped = [[2, 0, -1], [1, 3, -1]]
+        for platform, cuda, hip, fused, expected in (
+            ("cuda", True, False, True, local_slots),
+            ("hip", False, True, True, local_slots),
+            # Everything that is neither CUDA nor ROCm -- NPU in particular --
+            # still declines the seed, so fusion stays off and the wire
+            # positions are passed through unremapped.
+            ("other", False, False, False, unremapped),
         ):
-            self.assertTrue(
-                should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend=True)
-            )
-            draft_input = build_eagle_disagg_draft_input(
-                batch, torch.tensor([11, 12], dtype=torch.int64), None
-            )
+            with self.subTest(platform=platform), envs.SGLANG_DSA_FUSE_TOPK.override(
+                True
+            ), patch(
+                "sglang.srt.layers.attention.dsa.utils.is_cuda", return_value=cuda
+            ), patch(
+                "sglang.srt.layers.attention.dsa.utils.is_hip", return_value=hip
+            ):
+                self.assertEqual(
+                    should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend=True),
+                    fused,
+                )
+                draft_input = build_eagle_disagg_draft_input(
+                    batch, torch.tensor([11, 12], dtype=torch.int64), None
+                )
 
-        self.assertEqual(
-            draft_input.dsa_topk_indices.tolist(),
-            [[309, 101, -1], [801, 990, -1]],
-        )
+                self.assertEqual(draft_input.dsa_topk_indices.tolist(), expected)
 
     def test_future_map_initializes_seed_buffer_after_seedless_payload(self):
         future_map = object.__new__(FutureMap)