diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index fe65ef542..22626081c 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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_global: torch.Tensor, + input_ids: Optional[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() - input_ids_global = input_ids_global.tensor_split(s)[r].contiguous() + # DSpark next-token layers are not hash-routed and intentionally do not + # 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). diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 39b106aca..e6d683067 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/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)