[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