From be7c13af07cd93cf1967bdac61d8aa484af7e9e2 Mon Sep 17 00:00:00 2001 From: Zheng Wengang Date: Sat, 25 Jul 2026 03:19:02 +0800 Subject: [PATCH] [BugFix] Fix DS/Kimi crash on non-first PP ranks when resolving input length (#31752) --- python/sglang/srt/models/deepseek_v2.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index ba3b1c504..bc177ef24 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -2872,10 +2872,14 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin): input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: - # Minor fix for multi-modal model: input_ids is None - len_input_ids = ( - input_ids.shape[0] if input_ids is not None else input_embeds.shape[0] - ) + # Multi-modal: input_ids may be None (use input_embeds). + # Non-first PP ranks: both are None (activations via pp_proxy_tensors). + 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 can_dsa_cp_split( len_input_ids, self.cp_size, self.use_dsa, forward_batch