From 4a3a787f1e1ff3d1a5efd9eb0b7ffa00eb53d9a8 Mon Sep 17 00:00:00 2001 From: lw9527 <952799980@qq.com> Date: Wed, 25 Feb 2026 15:22:11 +0800 Subject: [PATCH] [Fix] Kimi K2.5 support pp (#18434) Co-authored-by: Ilya Boytsov Co-authored-by: ybyang <10629930+whybeyoung@users.noreply.github.com> Co-authored-by: Shangming Cai --- python/sglang/srt/models/deepseek_v2.py | 23 +++++++++++------------ python/sglang/srt/models/kimi_k25.py | 4 +++- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index f11d665c2..840343136 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -2656,7 +2656,17 @@ class DeepseekV2Model(nn.Module): pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: total_num_layers = self.end_layer - self.start_layer - device = input_embeds.device if input_embeds is not None else input_ids.device + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + device = hidden_states.device zero_allocator = BumpAllocator( buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1), dtype=torch.float32, @@ -2678,17 +2688,6 @@ class DeepseekV2Model(nn.Module): else None ) - if self.pp_group.is_first_rank: - if input_embeds is None: - hidden_states = self.embed_tokens(input_ids) - else: - hidden_states = input_embeds - residual = None - else: - assert pp_proxy_tensors is not None - hidden_states = pp_proxy_tensors["hidden_states"] - residual = pp_proxy_tensors["residual"] - if nsa_use_prefill_cp(forward_batch): if self.pp_group.is_first_rank: hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) diff --git a/python/sglang/srt/models/kimi_k25.py b/python/sglang/srt/models/kimi_k25.py index d8399a691..5e35a0d6f 100644 --- a/python/sglang/srt/models/kimi_k25.py +++ b/python/sglang/srt/models/kimi_k25.py @@ -30,7 +30,7 @@ from sglang.srt.managers.schedule_batch import ( MultimodalDataItem, MultimodalInputs, ) -from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV3ForCausalLM from sglang.srt.models.kimi_vl_moonvit import MLP2 @@ -722,6 +722,7 @@ class KimiK25ForConditionalGeneration(nn.Module): positions: torch.Tensor, forward_batch: ForwardBatch, get_embedding: bool = False, + pp_proxy_tensors: Optional[PPProxyTensors] = None, ): hidden_states = general_mm_embed_routine( input_ids=input_ids, @@ -731,6 +732,7 @@ class KimiK25ForConditionalGeneration(nn.Module): Modality.IMAGE: self.get_image_feature, }, positions=positions, + pp_proxy_tensors=pp_proxy_tensors, ) return hidden_states