chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

fix: embed before CP split in nextn to prevent TP all_reduce shape mismatch

Browse Source 
With CP=8 and dp_size=1, enable_dp_attention gets reset to False, so
VocabParallelEmbedding uses tensor_model_parallel_all_reduce (tp_size=8).
The CP local draft path was splitting tokens before embedding, giving each
rank a different local_tokens count. This caused an NCCL all_reduce shape
mismatch and a collective hang.

Move embed_tokens() before the CP split: embed on the full input (all ranks
see the same shape), then cp_split_and_rebuild_data the result. The decoder
layer still runs on CP-local tokens, preserving the CP performance benefit.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-05-22 11:21:09 +00:00
parent 8ee249f953
commit bd6e28f8ce
1 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/models/deepseek_nextn.py
View File

@@ -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]: