Fix DSpark CUDA graph replay with MegaMoE TP attention (#34919)

python/sglang/srt/models/deepseek_v4.py

@@ -2387,8 +2387,8 @@ class DeepseekV4DecoderLayer(nn.Module):
         hidden_states: torch.Tensor,
         forward_batch: ForwardBatch,
         *,
-        input_ids: torch.Tensor,
+        input_ids: Optional[torch.Tensor],
-        input_ids_global: torch.Tensor,
+        input_ids_global: Optional[torch.Tensor],
     ) -> torch.Tensor:
         _use_cp = self.dsa_enable_prefill_cp and dsa_use_prefill_cp(forward_batch)
         _use_tp_moe_gather = (
@@ -2475,8 +2475,12 @@ class DeepseekV4DecoderLayer(nn.Module):
             s, r = get_parallel().attn_tp_size, get_parallel().attn_tp_rank
             _a2a_scatter_chunks = list(hidden_states.tensor_split(s))
             hidden_states = _a2a_scatter_chunks[r].contiguous()
-            input_ids = input_ids.tensor_split(s)[r].contiguous()
+            # DSpark next-token layers are not hash-routed and intentionally do not
-            input_ids_global = input_ids_global.tensor_split(s)[r].contiguous()
+            # carry token IDs. Only split IDs for callers that actually provide them.
+            if input_ids is not None:
+                input_ids = input_ids.tensor_split(s)[r].contiguous()
+            if input_ids_global is not None:
+                input_ids_global = input_ids_global.tensor_split(s)[r].contiguous()
         # Skip the MoE-internal post-experts all_reduce when we will do the
         # reduce via reduce_scatterv/reduce_scatter at the combine below
         # (else double-reduce).

python/sglang/srt/speculative/dspark_components/dspark_draft.py

@@ -225,6 +225,11 @@ class DraftBlockProposer:
         # Persistent (bs, gamma) mask-token buffer: only column 0 (the bonus
         # token) changes per step, so avoid a fresh torch.full every decode.
         self._draft_block_ids_buf: Optional[torch.Tensor] = None
+        self._num_token_non_padded = (
+            torch.empty((1,), dtype=torch.int32, device=self.draft_model_runner.device)
+            if enable_num_token_non_padded()
+            else None
+        )
 
     def attach_draft_sampler(self, draft_sampler) -> None:
         self._draft_sampler = draft_sampler
@@ -473,6 +478,12 @@ class DraftBlockProposer:
         # The dense DSpark draft still reuses the target batch's graph tier.
         # Set graph eligibility before the DP-MoE-only metadata early return.
         forward_batch.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph
+        device = self.draft_model_runner.device
+        num_tokens = forward_batch.input_ids.numel()
+        if self._num_token_non_padded is not None:
+            self._num_token_non_padded.fill_(num_tokens)
+            forward_batch.num_token_non_padded = self._num_token_non_padded
+        forward_batch.num_token_non_padded_cpu = num_tokens
         if not self._dp_moe_sync or batch.global_num_tokens is None:
             return
         # Graph bucket selection uses the raw per-rank request counts.  Keep
@@ -484,7 +495,6 @@ class DraftBlockProposer:
             batch.global_num_tokens,
             batch.global_num_tokens_for_logprob,
         )
-        device = self.draft_model_runner.device
         forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens
         num_tokens = forward_batch.input_ids.numel()
         num_token_non_padded = _make_num_token_non_padded(num_tokens, device)