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

[BugFix] Fix DS/Kimi crash on non-first PP ranks when resolving input length (#31752)

Browse Source
This commit is contained in:
Zheng Wengang
2026-07-24 12:19:02 -07:00
committed by GitHub
parent 71015f3fea
commit be7c13af07
1 changed files with 8 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/deepseek_v2.py
+8 -4
View File
@@ -2872,10 +2872,14 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
input_embeds: torch.Tensor = None, input_embeds: torch.Tensor = None,
pp_proxy_tensors: Optional[PPProxyTensors] = None, pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> torch.Tensor: ) -> torch.Tensor:
# Minor fix for multi-modal model: input_ids is None # Multi-modal: input_ids may be None (use input_embeds).
len_input_ids = ( # Non-first PP ranks: both are None (activations via pp_proxy_tensors).
input_ids.shape[0] if input_ids is not None else input_embeds.shape[0] if input_ids is not None:
) len_input_ids = input_ids.shape[0]
elif input_embeds is not None:
len_input_ids = input_embeds.shape[0]
else:
len_input_ids = pp_proxy_tensors["hidden_states"].shape[0]
if self.dsa_enable_prefill_cp: if self.dsa_enable_prefill_cp:
if can_dsa_cp_split( if can_dsa_cp_split(
len_input_ids, self.cp_size, self.use_dsa, forward_batch len_input_ids, self.cp_size, self.use_dsa, forward_batch