diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 7df042c37..2a7f4a2d8 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/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(): diff --git a/python/sglang/srt/speculative/eagle_disaggregation.py b/python/sglang/srt/speculative/eagle_disaggregation.py index 57742557c..0c7b84c1c 100644 --- a/python/sglang/srt/speculative/eagle_disaggregation.py +++ b/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 diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index b96982b0c..1fe8f0a47 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/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