From 24c9251ac52ada1660f372922c72c1d3af722247 Mon Sep 17 00:00:00 2001 From: tianxiaojiang4 Date: Fri, 28 Aug 2026 22:34:23 -0700 Subject: [PATCH] [AMD][Spec][PD] Enable the PD DSA fused-TopK seed remap on ROCm (#36714) Co-authored-by: Tianxiao Jiang Co-authored-by: Claude Opus 5 (1M context) --- .../sglang/srt/layers/attention/dsa/utils.py | 2 +- .../test_disaggregation_wire.py | 36 ++++++++++++------- 2 files changed, 25 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 1b1b34df6..e469a6116 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/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 diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index c311a673c..3fbaea191 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/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)