diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py index c0d135f50..21ded56d7 100644 --- a/python/sglang/srt/models/deepseek_nextn.py +++ b/python/sglang/srt/models/deepseek_nextn.py @@ -179,8 +179,9 @@ class DeepseekModelNextN(nn.Module): use_cp = nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp) use_cp_local_draft = use_cp and envs.SGLANG_CP_DRAFT_SHARED_KV.get() if use_cp_local_draft: - local_input_ids = cp_split_and_rebuild_1d(forward_batch, input_ids) - local_num_tokens = local_input_ids.shape[0] + local_num_tokens = cp_split_and_rebuild_1d( + forward_batch, input_ids + ).shape[0] local_positions = cp_split_and_rebuild_position(forward_batch, positions) spec_hidden_states = self._get_cp_local_spec_hidden_states( forward_batch, @@ -193,7 +194,11 @@ class DeepseekModelNextN(nn.Module): else: positions = local_positions if input_embeds is None: - hidden_states = self.embed_tokens(local_input_ids) + # Embed full input first so all ranks see the same tensor + # shape in the TP all-reduce, then CP-split the result. + hidden_states = cp_split_and_rebuild_data( + forward_batch, self.embed_tokens(input_ids) + ) elif input_embeds.shape[0] == local_num_tokens: hidden_states = input_embeds elif input_embeds.shape[0] == input_ids.shape[0]: