[bug-fix] Stabilize GLM-5.2 MTP IndexShare across PD and CUDA graph replay (#30839)

Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

python/sglang/srt/debug_utils/pr_fix_toggle.py

@@ -27,12 +27,8 @@ patches:
   - target: sglang.srt.speculative.eagle_draft_cuda_graph_runner.EAGLEDraftCudaGraphRunner.capture_one_shape
     edits:
       - match: |
-          forward_batch.spec_info.hidden_states = hidden_states_backup
           forward_batch.positions.sub_(self.eagle_worker.speculative_num_steps - 1)
-          return ret
-        replacement: |
-          forward_batch.spec_info.hidden_states = hidden_states_backup
-          return ret
+        replacement: ""
 """
 
 

python/sglang/srt/disaggregation/decode.py

@@ -1640,6 +1640,7 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin):
             output_topk_p,
             output_topk_index,
             output_hidden_states,
+            output_dsa_topk_indices,
             output_bootstrap_room,
         ) = self.metadata_buffers.get_buf(idx)
 
@@ -1733,6 +1734,12 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin):
             decode_req.req.output_topk_p = output_topk_p
             decode_req.req.output_topk_index = output_topk_index
             decode_req.req.hidden_states_tensor = output_hidden_states
+            if (
+                output_dsa_topk_indices is not None
+                and torch.all(output_dsa_topk_indices < 0).item()
+            ):
+                output_dsa_topk_indices = None
+            decode_req.req.output_dsa_topk_indices = output_dsa_topk_indices
 
         if decode_req.req.return_logprob and not replayed_boundary:
             decode_req.req.logprob.output_token_logprobs_val.append(

python/sglang/srt/disaggregation/prefill.py

@@ -657,8 +657,14 @@ class SchedulerDisaggregationPrefillMixin:
                     req.hidden_states_tensor = (
                         batch.spec_info.hidden_states[i].cpu().clone()
                     )
+                    dsa_topk_indices = batch.spec_info.dsa_topk_indices
+                    if dsa_topk_indices is not None:
+                        req.output_dsa_topk_indices = dsa_topk_indices[i].cpu().clone()
+                    else:
+                        req.output_dsa_topk_indices = None
                 else:
                     req.hidden_states_tensor = None
+                    req.output_dsa_topk_indices = None
                 if req.return_logprob:
                     assert extend_logprob_start_len_per_req is not None
                     assert extend_input_len_per_req is not None
@@ -1158,6 +1164,7 @@ class SchedulerDisaggregationPrefillMixin:
         req.start_send_idx = 0
         req.tmp_end_idx = -1
         req.hidden_states_tensor = None
+        req.output_dsa_topk_indices = None
         req.pending_bootstrap = True
         req.time_stats.reset_prefill_retry_time()
         if req.prefill_attempt_count >= max_attempts:

python/sglang/srt/disaggregation/utils.py

@@ -11,6 +11,7 @@ import numpy as np
 import torch
 import torch.distributed as dist
 
+from sglang.srt.configs.model_config import get_dsa_index_topk
 from sglang.srt.disaggregation.base import KVPoll
 from sglang.srt.environ import envs
 from sglang.srt.utils import is_hip, is_npu
@@ -33,6 +34,13 @@ FAKE_BOOTSTRAP_HOST = "2.2.2.2"
 _IS_HIP = is_hip()
 
 
+def get_dsa_seed_metadata_dim(hf_config) -> int:
+    """Return the model-defined PD seed width, independent of local spec mode."""
+    if not getattr(hf_config, "index_share_for_mtp_iteration", False):
+        return 0
+    return get_dsa_index_topk(hf_config)
+
+
 def is_dsv4_c128_online_enabled() -> bool:
     """Return whether DSV4 C128 uses request-scoped online state."""
     return not _IS_HIP and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
@@ -228,8 +236,10 @@ class MetadataBuffers:
         max_top_logprobs_num: int = 128,
         max_sampling_mask_tokens: Optional[int] = None,
         custom_mem_pool: torch.cuda.MemPool = None,
+        output_dsa_topk_indices_dim: int = 0,
     ):
         self.custom_mem_pool = custom_mem_pool
+        self.output_dsa_topk_indices_dim = output_dsa_topk_indices_dim
         if max_sampling_mask_tokens is None:
             max_sampling_mask_tokens = (
                 envs.SGLANG_DISAGGREGATION_SAMPLING_MASK_MAX_TOKENS.get()
@@ -295,6 +305,15 @@ class MetadataBuffers:
             self.output_hidden_states = torch.zeros(
                 (size, hidden_size), dtype=hidden_states_dtype, device=device
             )
+            if self.output_dsa_topk_indices_dim > 0:
+                self.output_dsa_topk_indices = torch.full(
+                    (size, self.output_dsa_topk_indices_dim),
+                    -1,
+                    dtype=torch.int32,
+                    device=device,
+                )
+            else:
+                self.output_dsa_topk_indices = None
             # Request validation: store bootstrap_room to detect metadata corruption
             self.bootstrap_room = torch.zeros(
                 (size, 8), dtype=bootstrap_room_dtype, device=device
@@ -322,9 +341,11 @@ class MetadataBuffers:
                 self.output_topk_p,
                 self.output_topk_index,
                 self.output_hidden_states,
-                self.bootstrap_room,
             ]
         )
+        if self.output_dsa_topk_indices is not None:
+            bufs.append(self.output_dsa_topk_indices)
+        bufs.append(self.bootstrap_room)
         ptrs = [buf.data_ptr() for buf in bufs]
         data_lens = [buf.nbytes for buf in bufs]
         item_lens = [buf[0].nbytes for buf in bufs]
@@ -351,6 +372,11 @@ class MetadataBuffers:
             self.output_topk_p[idx].clone(),
             self.output_topk_index[idx].clone(),
             self.output_hidden_states[idx].clone(),
+            (
+                self.output_dsa_topk_indices[idx].clone()
+                if self.output_dsa_topk_indices is not None
+                else None
+            ),
             self.bootstrap_room[idx].clone(),
         )
 
@@ -461,6 +487,14 @@ class MetadataBuffers:
             self.output_hidden_states[req.metadata_buffer_index].copy_(
                 req.hidden_states_tensor
             )
+            if self.output_dsa_topk_indices is not None:
+                dsa_topk_indices = req.output_dsa_topk_indices
+                if dsa_topk_indices is not None:
+                    self.output_dsa_topk_indices[req.metadata_buffer_index].copy_(
+                        dsa_topk_indices
+                    )
+                else:
+                    self.output_dsa_topk_indices[req.metadata_buffer_index].fill_(-1)
         # Store bootstrap_room for validation on decode side
         self.bootstrap_room[req.metadata_buffer_index, 0] = (
             req.bootstrap_room if req.bootstrap_room is not None else 0

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

@@ -67,6 +67,17 @@ def compute_dsa_seqlens(original_seq_lens, dsa_index_topk: int):
     return original_seq_lens.clamp(max=dsa_index_topk)
 
 
+def should_use_dsa_fused_topk(
+    server_args, seed_dsa_topk_from_draft_extend: bool
+) -> bool:
+    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
+
+
 def is_dsa_enable_prefill_cp():
     return get_server_args().enable_dsa_prefill_context_parallel
 

python/sglang/srt/layers/attention/dsa_backend.py

@@ -45,6 +45,7 @@ from sglang.srt.layers.attention.dsa.utils import (
     is_dsa_enable_prefill_cp,
     is_dsa_prefill_cp_in_seq_split,
     pad_dsa_cache_seqlens,
+    should_use_dsa_fused_topk,
 )
 from sglang.srt.layers.attention.utils import (
     concat_mla_absorb_q_general,
@@ -63,6 +64,7 @@ from sglang.srt.utils import (
     is_gfx95_supported,
     is_hip,
     is_sm100_supported,
+    print_warning_once,
 )
 
 # Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton
@@ -340,6 +342,7 @@ class DeepseekSparseAttnBackend(
         speculative_step_id=0,
         topk=0,
         speculative_num_steps=0,
+        seed_dsa_topk_from_draft_extend: bool = False,
     ):
         super().__init__()
         self.forward_metadata: DSAMetadata
@@ -433,6 +436,13 @@ class DeepseekSparseAttnBackend(
             model_runner.server_args.speculative_num_draft_tokens
         )
         self.speculative_step_id = speculative_step_id
+        self.use_fused_topk = should_use_dsa_fused_topk(
+            model_runner.server_args, seed_dsa_topk_from_draft_extend
+        )
+        if envs.SGLANG_DSA_FUSE_TOPK.get() and not self.use_fused_topk:
+            print_warning_once(
+                "Disabling fused DSA top-k for IndexShare under PD disaggregation."
+            )
 
         self.device_capability = torch.cuda.get_device_capability()
         self.device_sm_major = self.device_capability[0]
@@ -1107,7 +1117,7 @@ class DeepseekSparseAttnBackend(
             and self.real_page_size > 1
             and self.hisparse_coordinator is None
             and not self.speculative_num_draft_tokens
-            and envs.SGLANG_DSA_FUSE_TOPK.get()
+            and self.use_fused_topk
             and envs.SGLANG_OPT_USE_TOPK_V2.get()
             and self.dsa_index_topk is not None
             and self.dsa_index_topk <= 2048
@@ -1897,7 +1907,7 @@ class DeepseekSparseAttnBackend(
         topk_transform_method = self.get_topk_transform_method(
             forward_batch.forward_mode
         )
-        if envs.SGLANG_DSA_FUSE_TOPK.get():
+        if self.use_fused_topk:
             page_table_1 = self._get_fused_topk_page_table(topk_indices)
         else:
             if topk_transform_method == TopkTransformMethod.RAGGED:
@@ -2112,7 +2122,7 @@ class DeepseekSparseAttnBackend(
                 topk_indices,
                 layer.layer_id,
             )
-        elif envs.SGLANG_DSA_FUSE_TOPK.get():
+        elif self.use_fused_topk:
             page_table_1 = self._get_fused_topk_page_table(topk_indices)
         else:
             page_table_1 = transform_index_page_table_decode(
@@ -2669,7 +2679,7 @@ class DeepseekSparseAttnBackend(
         if topk_indices is not None:
             topk_indices = self._pad_topk_indices(topk_indices, q.shape[0])
 
-        if envs.SGLANG_DSA_FUSE_TOPK.get():
+        if self.use_fused_topk:
             page_table_1 = self._get_fused_topk_page_table(topk_indices)
         elif is_prefill:
             page_table_1 = transform_index_page_table_prefill(
@@ -2835,7 +2845,7 @@ class DeepseekSparseAttnBackend(
     def get_indexer_metadata(
         self, layer_id: int, forward_batch: ForwardBatch
     ) -> DSAIndexerMetadata:
-        force_unfused = (
+        force_unfused = not self.use_fused_topk or (
             self.hisparse_coordinator is not None
             and forward_batch.forward_mode.is_decode_or_idle()
         )
@@ -2879,7 +2889,11 @@ class DeepseekSparseAttnMultiStepBackend:
     needs_cpu_seq_lens: bool = False
 
     def __init__(
-        self, model_runner: ModelRunner, topk: int, speculative_num_steps: int
+        self,
+        model_runner: ModelRunner,
+        topk: int,
+        speculative_num_steps: int,
+        seed_dsa_topk_from_draft_extend: bool = False,
     ):
         self.topk = topk
         self.speculative_num_steps = speculative_num_steps
@@ -2891,6 +2905,7 @@ class DeepseekSparseAttnMultiStepBackend:
                     speculative_step_id=i,
                     topk=self.topk,
                     speculative_num_steps=self.speculative_num_steps,
+                    seed_dsa_topk_from_draft_extend=seed_dsa_topk_from_draft_extend,
                 )
             )
 

python/sglang/srt/managers/overlap_utils.py

@@ -146,7 +146,7 @@ class RelayPayload:
             topk_index=draft_input.topk_index,
             hidden_states=draft_input.hidden_states,
             draft_probs=getattr(draft_input, "draft_probs", None),
-            dsa_topk_indices=getattr(draft_input, "dsa_topk_indices", None),
+            dsa_topk_indices=draft_input.dsa_topk_indices,
         )
 
 
@@ -281,6 +281,7 @@ class FutureMap:
             self.fwd_prepare_d2h_stream = None
         # Lazy-inited on the first non-empty stash (peeks tensor shapes); non-spec's is a no-op.
         self._forward_buf_initialized = False
+        self.dsa_topk_indices_buf = None
 
         self.publish_ready = None  # lazy device.Event(); only spec_v2 needs it
         # Debug consume-once state: armed by a recording publish, consumed by
@@ -338,14 +339,15 @@ class FutureMap:
                 device=self.device,
             )
 
-        self.dsa_topk_indices_buf = None
-        if payload.dsa_topk_indices is not None:
-            seed0 = payload.dsa_topk_indices[0]
-            self.dsa_topk_indices_buf = torch.empty(
-                (self.req_pool_size, *seed0.shape),
-                dtype=payload.dsa_topk_indices.dtype,
-                device=self.device,
-            )
+    def _maybe_init_dsa_topk_indices_buf(self, payload: RelayPayload) -> None:
+        if self.dsa_topk_indices_buf is not None or payload.dsa_topk_indices is None:
+            return
+        seed0 = payload.dsa_topk_indices[0]
+        self.dsa_topk_indices_buf = torch.empty(
+            (self.req_pool_size, *seed0.shape),
+            dtype=payload.dsa_topk_indices.dtype,
+            device=self.device,
+        )
 
     def resolve_confidence_cpu(
         self, batch: ScheduleBatch
@@ -397,8 +399,11 @@ class FutureMap:
             draft_input.bonus_tokens = self.output_tokens_buf[indices]
         if self.need_hidden_states and not self.need_topk:
             draft_input.hidden_states = self.hidden_states_buf[indices]
-        if self.dsa_topk_indices_buf is not None:
+        if draft_input.future_dsa_topk_indices_available:
+            assert self.dsa_topk_indices_buf is not None
             draft_input.dsa_topk_indices = self.dsa_topk_indices_buf[indices]
+        else:
+            draft_input.dsa_topk_indices = None
         if _DEBUG_ASSERT:
             _assert_nonneg_and_invalidate(
                 draft_input.bonus_tokens, self.output_tokens_buf, indices
@@ -503,6 +508,7 @@ class FutureMap:
             return
         if not self._forward_buf_initialized:
             self._lazy_init_forward_buf(payload)
+        self._maybe_init_dsa_topk_indices_buf(payload)
         self.output_tokens_buf[indices] = payload.bonus_tokens.to(
             self.output_tokens_buf.dtype
         )

python/sglang/srt/managers/schedule_batch.py

@@ -933,6 +933,7 @@ class Req(ReqDllmMixin):
         self.hidden_states_tensor = None  # Note: use tensor instead of list to transfer hidden_states when PD + MTP
         self.output_topk_p = None
         self.output_topk_index = None
+        self.output_dsa_topk_indices = None
 
         # capture routed experts
         self.return_routed_experts = return_routed_experts

python/sglang/srt/managers/scheduler.py

@@ -61,6 +61,7 @@ from sglang.srt.disaggregation.utils import (
     MetadataBuffers,
     ReqToMetadataIdxAllocator,
     TransferBackend,
+    get_dsa_seed_metadata_dim,
     prepare_abort,
 )
 from sglang.srt.distributed import get_pp_group, get_world_group
@@ -72,9 +73,7 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r
 from sglang.srt.layers.attention.mamba.ops import (
     initialize_mamba_selective_state_update_backend,
 )
-from sglang.srt.layers.dp_attention import (
-    compute_dp_attention_world_info,
-)
+from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config
 from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
@@ -1139,6 +1138,12 @@ class Scheduler(
             disagg_hidden_size = 16  # minimal padding size for RDMA
             disagg_hidden_states_dtype = torch.float32
 
+        # The PD metadata wire schema must match on P and D even when only D
+        # enables spec decoding; a seedless prefill writes the invalid sentinel.
+        output_dsa_topk_indices_dim = get_dsa_seed_metadata_dim(
+            self.model_config.hf_config
+        )
+
         if (
             self.disaggregation_mode == DisaggregationMode.DECODE
         ):  # *8 headroom for MiniMax-M3; *2 for other models.
@@ -1154,6 +1159,7 @@ class Scheduler(
                 hidden_size=disagg_hidden_size,
                 hidden_states_dtype=disagg_hidden_states_dtype,
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
+                output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
             )
 
             # The decode requests polling kv cache
@@ -1199,6 +1205,7 @@ class Scheduler(
                 hidden_size=disagg_hidden_size,
                 hidden_states_dtype=disagg_hidden_states_dtype,
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
+                output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
             )
 
             self.disagg_prefill_bootstrap_queue = PrefillBootstrapQueue(
@@ -3390,6 +3397,9 @@ class Scheduler(
 
                 if not batch.spec_algorithm.is_none():
                     batch.spec_info = batch_result.next_draft_input
+                    batch.spec_info.future_dsa_topk_indices_available = (
+                        batch.spec_info.dsa_topk_indices is not None
+                    )
                     batch.spec_info.future_indices = future_indices
             elif self.enable_pdmux and batch.forward_mode.is_split_prefill():
                 resolve_forward_inputs(batch, self.future_map)

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -61,6 +61,36 @@ def _resolve_attn_backend(forward_batch: ForwardBatch):
     return backend
 
 
+def _forward_dsa_indexer_for_mha(
+    indexer,
+    *,
+    hidden_states: torch.Tensor,
+    q_lora: torch.Tensor,
+    positions: torch.Tensor,
+    forward_batch: ForwardBatch,
+    layer_id: int,
+) -> None:
+    """Fill the indexer K cache and publish an MTP seed when requested."""
+    spec_info = forward_batch.spec_info
+    seed_buf = spec_info.dsa_seed_topk_capture if spec_info is not None else None
+    topk_indices = indexer(
+        x=hidden_states,
+        q_lora=q_lora,
+        positions=positions,
+        forward_batch=forward_batch,
+        layer_id=layer_id,
+        return_indices=seed_buf is not None,
+    )
+    if seed_buf is None:
+        return
+    if topk_indices is None:
+        raise RuntimeError("DSA MHA indexer did not produce the requested MTP seed")
+
+    select = spec_info.dsa_seed_topk_select
+    src = topk_indices if select is None else topk_indices[select]
+    seed_buf[: src.shape[0]].copy_(src)
+
+
 # Configs for DeepSeek-V3:
 # num_local_heads = 128
 # qk_nope_head_dim = 128
@@ -174,13 +204,13 @@ class DeepseekMHAForwardMixin:
                         -1, self.num_local_heads, self.qk_head_dim
                     )
                 if self.should_run_indexer():
-                    _ = self.indexer(
-                        x=hidden_states,
+                    _forward_dsa_indexer_for_mha(
+                        self.indexer,
+                        hidden_states=hidden_states,
                         q_lora=q_lora,
                         positions=positions,
                         forward_batch=forward_batch,
                         layer_id=self.layer_id,
-                        return_indices=False,
                     )
             elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
                 # MXFP4: fused RMSNorm + quant

python/sglang/srt/speculative/draft_utils.py

@@ -20,11 +20,13 @@ class DraftBackendFactory:
         draft_model_runner,
         topk: int,
         speculative_num_steps: int,
+        seed_dsa_topk_from_draft_extend: bool = False,
     ):
         self.server_args = server_args
         self.draft_model_runner = draft_model_runner
         self.topk = topk
         self.speculative_num_steps = speculative_num_steps
+        self.seed_dsa_topk_from_draft_extend = seed_dsa_topk_from_draft_extend
         self.draft_attn_backend = server_args.speculative_draft_attention_backend
 
     def _create_backend(
@@ -110,13 +112,20 @@ class DraftBackendFactory:
         )
 
         return DeepseekSparseAttnMultiStepBackend(
-            self.draft_model_runner, self.topk, self.speculative_num_steps
+            self.draft_model_runner,
+            self.topk,
+            self.speculative_num_steps,
+            seed_dsa_topk_from_draft_extend=self.seed_dsa_topk_from_draft_extend,
         )
 
     def _create_dsa_prefill_backend(self):
         from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend
 
-        return DeepseekSparseAttnBackend(self.draft_model_runner, skip_prefill=False)
+        return DeepseekSparseAttnBackend(
+            self.draft_model_runner,
+            skip_prefill=False,
+            seed_dsa_topk_from_draft_extend=self.seed_dsa_topk_from_draft_extend,
+        )
 
     def _create_flashinfer_decode_backend(self):
         if not self.draft_model_runner.use_mla_backend:

python/sglang/srt/speculative/eagle_disaggregation.py

@@ -51,15 +51,24 @@ def build_eagle_disagg_draft_input(
         [req.hidden_states_tensor for req in batch.reqs], dim=0
     ).to(batch.device)
 
+    dsa_topk_indices = None
+    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 torch.any(torch.all(dsa_topk_indices < 0, dim=1)).item():
+            dsa_topk_indices = None
+
     spec_info = EagleDraftInput(
         topk_p=topk_p,
         topk_index=topk_index,
         hidden_states=hidden_states,
         bonus_tokens=last_tokens_tensor,
+        dsa_topk_indices=dsa_topk_indices,
     )
     spec_info.capture_hidden_mode = CaptureHiddenMode.LAST
 
     if batch.enable_overlap:
+        spec_info.future_dsa_topk_indices_available = dsa_topk_indices is not None
         spec_info.future_indices = batch.req_pool_indices
         # Seed the relay buf with the known seq_lens; publish's chained record
         # keeps the in-flight forward's fence intact (see FutureMap.publish).

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -442,11 +442,13 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
 
             output_cache_loc_backup = forward_batch.out_cache_loc
             hidden_states_backup = forward_batch.spec_info.hidden_states
+            dsa_topk_indices_backup = forward_batch.spec_info.dsa_topk_indices
 
             ret = self.eagle_worker.draft_forward(forward_batch)
 
             forward_batch.out_cache_loc = output_cache_loc_backup
             forward_batch.spec_info.hidden_states = hidden_states_backup
+            forward_batch.spec_info.dsa_topk_indices = dsa_topk_indices_backup
             forward_batch.positions.sub_(self.eagle_worker.speculative_num_steps - 1)
             return ret
 
@@ -647,6 +649,8 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
         )
         with timer_ctx:
             out = self._replay_graph(shape_key, forward_batch)
+        if self.buffers.dsa_seed_topk is not None:
+            forward_batch.spec_info.dsa_topk_indices = None
 
         if bs != raw_bs:
             out = self._postprocess_output_to_raw_bs(out, raw_bs)

python/sglang/srt/speculative/eagle_info.py

@@ -178,6 +178,7 @@ class EagleDraftInput(SpecInput):
 
     # V2 overlap worker only: req_pool_indices used as buf slot keys.
     future_indices: Optional[torch.Tensor] = None
+    future_dsa_topk_indices_available: bool = False
 
     def __post_init__(self):
         super().__init__(SpecInputType.EAGLE_DRAFT)
@@ -255,6 +256,10 @@ class EagleDraftInput(SpecInput):
             self.future_indices = torch.cat(
                 [self.future_indices, spec_info.future_indices]
             )
+            self.future_dsa_topk_indices_available = (
+                self.future_dsa_topk_indices_available
+                and spec_info.future_dsa_topk_indices_available
+            )
             return
 
         # Detect idle stub by `topk_index` length (idle inputs have

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -373,6 +373,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             self.draft_runner,
             self.topk,
             self.speculative_num_steps,
+            seed_dsa_topk_from_draft_extend=self.seed_dsa_topk_from_draft_extend,
         )
 
         # Initialize decode attention backend
@@ -519,6 +520,13 @@ class EagleDraftWorker(EagleDraftWorkerBase):
             self.topk,
             self.speculative_num_steps,
         )
+        if (
+            can_cuda_graph
+            and not forward_batch.forward_mode.is_idle()
+            and self.seed_dsa_topk_from_draft_extend
+            and draft_input.dsa_topk_indices is None
+        ):
+            can_cuda_graph = False
 
         n_inner = self.speculative_num_steps - 1
         canary_outside_ctx = (

python/sglang/srt/speculative/spec_info.py

@@ -281,9 +281,7 @@ class SpeculativeAlgorithm(Enum):
 
             return EAGLEWorkerV2
         elif self.is_standalone():
-            from sglang.srt.speculative.standalone_worker_v2 import (
-                StandaloneWorkerV2,
-            )
+            from sglang.srt.speculative.standalone_worker_v2 import StandaloneWorkerV2
 
             return StandaloneWorkerV2
         elif self.is_ngram():
@@ -316,6 +314,13 @@ class SpecInput(ABC):
     # assignment, so an init-time default would clobber the passed layout.
     ragged_verify_layout: Optional[RaggedVerifyLayout] = None
 
+    # DSA MTP IndexShare seed relay. Class-level defaults (same rationale as
+    # ragged_verify_layout) so scheduler/relay/attention code reads them
+    # uniformly on any SpecInput; only the EAGLE-family inputs override them.
+    dsa_topk_indices: Optional[torch.Tensor] = None
+    future_dsa_topk_indices_available: bool = False
+    dsa_seed_topk_capture: Optional[torch.Tensor] = None
+
     def __init__(self, spec_input_type: SpecInputType):
         self.spec_input_type = spec_input_type