[Spec][PD] Enable fused TopK for GLM-5.2 MTP IndexShare (#31477)

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

@@ -64,21 +64,42 @@ INDEXER_K_CACHE_PRESHUFFLE_TILE = 16
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.server_args import ServerArgs
 
 
 def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int):
     return original_seq_lens.clamp(max=dsa_index_topk)
 
 
+def should_remap_pd_dsa_seed_to_local_slots(server_args: "ServerArgs") -> bool:
+    """Whether a PD seed should enter the allocator-local fused TopK domain."""
+    return (
+        is_cuda()
+        and envs.SGLANG_DSA_FUSE_TOPK.get()
+        and server_args.disaggregation_mode == "decode"
+        and not server_args.enable_hisparse
+        and server_args.dcp_size == 1
+    )
+
+
 def should_use_dsa_fused_topk(
-    server_args, seed_dsa_topk_from_draft_extend: bool
+    server_args: "ServerArgs", seed_dsa_topk_from_draft_extend: bool
 ) -> bool:
+    """Select fused TopK for PD IndexShare.
+
+    PD Prefill worker:
+    - Target prefill: fused TopK enabled.
+    - Draft extend: fused TopK disabled.
+
+    PD Decode worker:
+    - Draft decode / target verify / draft extend: fused TopK enabled.
+    """
     pd_index_share_seed = (
         server_args.disaggregation_mode != "null" and seed_dsa_topk_from_draft_extend
     )
-    # TODO(kpham-sgl): Transfer request-relative IndexShare seeds and remap them
-    # to decode-local KV slots so fused top-k can remain enabled under PD.
-    return envs.SGLANG_DSA_FUSE_TOPK.get() and not pd_index_share_seed
+    return envs.SGLANG_DSA_FUSE_TOPK.get() and (
+        not pd_index_share_seed or should_remap_pd_dsa_seed_to_local_slots(server_args)
+    )
 
 
 def is_dsa_enable_prefill_cp():

python/sglang/srt/speculative/eagle_disaggregation.py

@@ -4,6 +4,9 @@ from typing import TYPE_CHECKING
 
 import torch
 
+from sglang.srt.layers.attention.dsa.utils import (
+    should_remap_pd_dsa_seed_to_local_slots,
+)
 from sglang.srt.managers.overlap_utils import RelayPayload
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
 from sglang.srt.speculative.eagle_info import EagleDraftInput
@@ -55,6 +58,30 @@ def build_eagle_disagg_draft_input(
     dsa_indices_list = [req.output_dsa_topk_indices for req in batch.reqs]
     if dsa_indices_list and all(t is not None for t in dsa_indices_list):
         dsa_topk_indices = torch.stack(dsa_indices_list, dim=0).to(batch.device)
+        if should_remap_pd_dsa_seed_to_local_slots(server_args):
+            # PD sends request-relative positions; fused TopK consumes
+            # decode-local physical slots. Remap once before the draft loop/graph.
+            req_to_token = batch.req_to_token_pool.req_to_token
+            table_width = req_to_token.shape[1]
+            valid_positions = dsa_topk_indices >= 0
+            gather_positions = dsa_topk_indices.clamp(min=0, max=table_width - 1).to(
+                torch.int64
+            )
+            local_slots = req_to_token[
+                batch.req_pool_indices[:, None], gather_positions
+            ]
+            invalid_rows = torch.any(
+                (dsa_topk_indices < -1)
+                | (dsa_topk_indices >= batch.seq_lens[:, None])
+                | (dsa_topk_indices >= table_width)
+                # Slot 0 is the reserved padding sink; real KV allocations
+                # start at 1, and untouched req-to-token entries remain 0.
+                | (valid_positions & (local_slots <= 0)),
+                dim=1,
+            )
+            local_slots.masked_fill_(~valid_positions, -1)
+            local_slots.masked_fill_(invalid_rows[:, None], -1)
+            dsa_topk_indices = local_slots
         if torch.any(torch.all(dsa_topk_indices < 0, dim=1)).item():
             dsa_topk_indices = None
 

test/registered/unit/disaggregation/test_disaggregation_wire.py

@@ -1,5 +1,6 @@
 import unittest
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import numpy as np
 import torch
@@ -17,6 +18,8 @@ from sglang.srt.disaggregation.utils import (
     get_dsv4_c128_state_indices,
     setup_state_kv_args,
 )
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.dsa.utils import should_use_dsa_fused_topk
 from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.speculative.eagle_disaggregation import (
@@ -150,6 +153,7 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
             speculative_eagle_topk=1,
             speculative_num_steps=5,
             enable_multi_layer_eagle=False,
+            disaggregation_mode="null",
         )
         last_tokens = torch.tensor([11, 12], dtype=torch.int64)
 
@@ -168,6 +172,54 @@ class TestEagleDsaSeedTransfer(unittest.TestCase):
             )
             self.assertIsNone(draft_input.dsa_topk_indices)
 
+    def test_pd_decode_fused_topk_remaps_wire_positions_to_local_slots(self):
+        wire_positions = (
+            torch.tensor([2, 0, -1], dtype=torch.int32),
+            torch.tensor([1, 3, -1], dtype=torch.int32),
+        )
+        req_to_token = torch.tensor(
+            [
+                [0, 0, 0, 0],
+                [700, 801, 902, 990],
+                [410, 420, 430, 440],
+                [101, 205, 309, 450],
+            ],
+            dtype=torch.int32,
+        )
+        batch = SimpleNamespace(
+            reqs=[self._make_req(seed) for seed in wire_positions],
+            device="cpu",
+            enable_overlap=False,
+            req_pool_indices=torch.tensor([3, 1], dtype=torch.int64),
+            req_to_token_pool=SimpleNamespace(req_to_token=req_to_token),
+            seq_lens=torch.tensor([4, 4], dtype=torch.int32),
+        )
+        server_args = SimpleNamespace(
+            speculative_eagle_topk=1,
+            speculative_num_steps=5,
+            enable_multi_layer_eagle=False,
+            disaggregation_mode="decode",
+            enable_hisparse=False,
+            dcp_size=1,
+        )
+
+        with envs.SGLANG_DSA_FUSE_TOPK.override(True), patch(
+            "sglang.srt.layers.attention.dsa.utils.is_cuda", return_value=True
+        ):
+            self.assertTrue(
+                should_use_dsa_fused_topk(
+                    server_args, seed_dsa_topk_from_draft_extend=True
+                )
+            )
+            draft_input = build_eagle_disagg_draft_input(
+                batch, server_args, torch.tensor([11, 12], dtype=torch.int64), None
+            )
+
+        self.assertEqual(
+            draft_input.dsa_topk_indices.tolist(),
+            [[309, 101, -1], [801, 990, -1]],
+        )
+
     def test_future_map_initializes_seed_buffer_after_seedless_payload(self):
         future_map = object.__new__(FutureMap)
         future_map.dsa_topk_indices_buf = None