From c2059c4fb2f5f12c112735424b5d3f8c78dc1a1f Mon Sep 17 00:00:00 2001 From: abing Date: Tue, 15 Sep 2026 23:41:20 -0700 Subject: [PATCH] run pass llm cp (cherry picked from commit 1d85394563d96cf22d5d84f15c0e9043582419e4) --- python/sglang/srt/arg_groups/model_hook.py | 6 ++- .../sglang/srt/disaggregation/common/conn.py | 11 ++++- python/sglang/srt/mem_cache/common.py | 4 ++ .../srt/model_executor/runner/eager_runner.py | 14 ++++-- python/sglang/srt/models/deepseek_v4.py | 43 +++++++++++++++---- python/sglang/srt/server_args.py | 1 + 6 files changed, 63 insertions(+), 16 deletions(-) diff --git a/python/sglang/srt/arg_groups/model_hook.py b/python/sglang/srt/arg_groups/model_hook.py index c4505d0ad..f14e182d2 100644 --- a/python/sglang/srt/arg_groups/model_hook.py +++ b/python/sglang/srt/arg_groups/model_hook.py @@ -966,12 +966,14 @@ def handle_language_model_only(server_args: Any): ): if flag: raise ValueError(f"--language-model-only cannot be combined with {name}") - if cfg.disaggregation_mode != "null": + hf_config = model_config_of(server_args).hf_config + # V4.1 text-only workers use the standard PD KV transfer path. + if cfg.disaggregation_mode != "null" and hf_config.model_type != "deepseek_v41": raise ValueError( "--language-model-only is incompatible with --disaggregation-mode " "prefill/decode" ) - architectures = model_config_of(server_args).hf_config.architectures + architectures = hf_config.architectures if not any( a in server_args.LANGUAGE_MODEL_ONLY_ARCHITECTURES for a in architectures ): diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index b4196f657..b129cae84 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -934,9 +934,16 @@ class CommonKVManager(BaseKVManager): "enable DSpark with the same block size and target/draft KV " "layout. Upgrade both servers together." ) - if info.attn_tp_size != self.attn_tp_size: + same_tp_with_prefill_cp = ( + info.attn_cp_size > 1 + and (self.is_mla_backend or self.is_hybrid_mla_backend) + and self.attn_cp_size == 1 + and info.attn_tp_size * info.attn_cp_size == self.attn_tp_size + ) + if info.attn_tp_size != self.attn_tp_size and not same_tp_with_prefill_cp: raise RuntimeError( - "DeepSeek-V4.1 DSpark PD requires the same TP size on both servers" + "DeepSeek-V4.1 DSpark PD requires the same TP size on both " + "servers (including prefill CP ranks for MLA)" ) if self.dcp_size > 1: diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 9640071cc..545f9b65f 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -155,6 +155,10 @@ def free_kv_row_segments( def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs): if getattr(req, "skip_radix_cache_insert", False): + kv_indices = tree_cache.req_to_token_pool.req_to_token[ + req.kv.req_pool_idx, : len(req.get_fill_ids()) + ] + req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return tree_cache.cache_unfinished_req(req, **kwargs) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 54a01852d..33383b7df 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -385,14 +385,22 @@ class EagerRunner(BaseRunner): """ model = self.model_runner.model + input_ids = forward_batch.input_ids input_embeds = kwargs.get("input_embeds") + # Multimodal spans must be embedded in global token order, before CP + # slicing. The model may also normalize image hash IDs for its router. + prepare_inputs = getattr(model, "prepare_language_model_inputs", None) + if prepare_inputs is not None: + input_ids, input_embeds = prepare_inputs( + input_ids, forward_batch, input_embeds + ) if input_embeds is None: - input_embeds = model.get_input_embeddings()(forward_batch.input_ids) + input_embeds = model.get_input_embeddings()(input_ids) with cp_shard_model_inputs( input_embeds, forward_batch.positions, forward_batch, - forward_batch.input_ids, + input_ids, ) as (sharded_input_embeds, sharded_positions, model_input_ids): model_kwargs = {"input_embeds": sharded_input_embeds} if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None: @@ -437,7 +445,7 @@ class EagerRunner(BaseRunner): if aux_hidden_states is None: logits_kwargs["hidden_states_before_norm"] = hidden_states_before_norm return model.logits_processor( - forward_batch.input_ids, + input_ids, hidden_states, model.lm_head, forward_batch, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index a3b8b6102..2cfe15220 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -2650,7 +2650,8 @@ class DeepseekV4DecoderLayer(nn.Module): is_nextn=is_nextn, is_deepseek_v4=True, vl_correction_bias=config.model_type == "deepseek_v41" - and config.vision_n_layers > 0, + and config.vision_n_layers > 0 + and not getattr(config, "language_model_only", False), ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -3872,7 +3873,15 @@ class DeepseekV4DecoderLayer(nn.Module): finally: forward_batch.num_token_non_padded = saved_num_token_non_padded if _use_cp and get_moe_a2a_backend().is_none(): - hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states) + if self.config.model_type == "deepseek_v41": + parallel = get_parallel() + hidden_states = parallel.tp_group.all_reduce(hidden_states) + parallel = get_parallel() + hidden_states = hidden_states.tensor_split(parallel.attn_cp_size)[ + parallel.attn_cp_rank + ].contiguous() + else: + hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states) elif _use_tp_moe_gather: hidden_states, global_hidden_states = ( get_local_dp_buffer(get_parallel().tp_group), @@ -4871,7 +4880,7 @@ class DeepseekV4ForCausalLM(nn.Module): or not get_moe_a2a_backend().is_none() ): raise ValueError( - "V4.1 vision currently supports TP/EP/DP without CP, PP or MoE A2A" + "V4.1 vision supports TP/EP/DP and prefill CP without PP or MoE A2A" ) args = SimpleNamespace(**vars(config), dim=config.hidden_size) @@ -5078,16 +5087,19 @@ class DeepseekV4ForCausalLM(nn.Module): 0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts ) - def forward( + def prepare_language_model_inputs( self, input_ids: torch.Tensor, - positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: Optional[torch.Tensor] = None, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """Prepare full-sequence image embeddings and model IDs before CP splits. + + Scheduler hash IDs stay intact for multimodal cache keys; the language + model uses image_token_id for Engram masking and visual MoE routing. + """ if ( - self.vision is not None + getattr(self, "vision", None) is not None and not forward_batch.forward_mode.is_decode() and not forward_batch.forward_mode.is_target_verify() and forward_batch.mm_inputs is not None @@ -5096,7 +5108,7 @@ class DeepseekV4ForCausalLM(nn.Module): if input_embeds is not None: raise ValueError("Cannot combine input_embeds and image inputs") input_embeds = self._prepare_mm_embeddings(input_ids, forward_batch) - if self.vision is not None and not ( + if getattr(self, "vision", None) is not None and not ( forward_batch.forward_mode.is_decode_or_idle() or forward_batch.forward_mode.is_target_verify() ): @@ -5106,6 +5118,19 @@ class DeepseekV4ForCausalLM(nn.Module): input_ids >= MM_PAD_SHIFT_VALUE, self.config.image_token_id ) + return input_ids, input_embeds + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + input_ids, input_embeds = self.prepare_language_model_inputs( + input_ids, forward_batch, input_embeds + ) with get_attn_tp_context().maybe_input_scattered(forward_batch): hidden_states = self.model.forward( input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1fafc9035..249daf5e2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -331,6 +331,7 @@ class ServerArgs: # ===== END TO BE REFACTORED ==== LANGUAGE_MODEL_ONLY_ARCHITECTURES = ( + "DeepseekV4ForCausalLM", "MuseGlimmerForConditionalGeneration", "Cosmos3ForConditionalGeneration", "Cosmos3EdgeForConditionalGeneration",