From ac19ce7efb37679031669791eb86726d18ca4a76 Mon Sep 17 00:00:00 2001 From: Xuchun Shang Date: Mon, 10 Nov 2025 14:58:19 +0800 Subject: [PATCH] [PP] put pp assert in model runner (#12934) Signed-off-by: Xuchun Shang --- python/sglang/srt/managers/scheduler_pp_mixin.py | 2 -- python/sglang/srt/model_executor/model_runner.py | 11 ++++++++--- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py index d14a31b85..cfa4e2369 100644 --- a/python/sglang/srt/managers/scheduler_pp_mixin.py +++ b/python/sglang/srt/managers/scheduler_pp_mixin.py @@ -122,8 +122,6 @@ class SchedulerPPMixin: # send out proxy tensors to the next stage if self.cur_batch: - # FIXME(lsyin): remove this assert - assert result.pp_hidden_states_proxy_tensors.tensors is not None self.pp_group.send_tensor_dict( result.pp_hidden_states_proxy_tensors.tensors, all_gather_group=self.attn_tp_group, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 8433b0247..19e029d60 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -327,6 +327,11 @@ class ModelRunner: "pp_proxy_tensors" in inspect.signature(self.model.forward).parameters ) + if self.pp_size > 1: + assert ( + self.support_pp + ), "Pipeline Parallel is not compatible with this model." + # For weight updates self._model_update_group = {} self._weights_send_group = {} @@ -2056,7 +2061,7 @@ class ModelRunner: forward_batch: ForwardBatch, skip_attn_backend_init: bool = False, pp_proxy_tensors=None, - ) -> LogitsProcessorOutput: + ) -> Union[LogitsProcessorOutput, PPProxyTensors]: if not skip_attn_backend_init: if self.server_args.enable_pdmux: self.decode_attn_backend.init_forward_metadata(forward_batch) @@ -2079,7 +2084,7 @@ class ModelRunner: forward_batch: ForwardBatch, skip_attn_backend_init: bool = False, pp_proxy_tensors=None, - ) -> LogitsProcessorOutput: + ) -> Union[LogitsProcessorOutput, PPProxyTensors]: kwargs = {} if self.support_pp: kwargs["pp_proxy_tensors"] = pp_proxy_tensors @@ -2106,7 +2111,7 @@ class ModelRunner: def forward_idle( self, forward_batch: ForwardBatch, pp_proxy_tensors=None - ) -> LogitsProcessorOutput: + ) -> Union[LogitsProcessorOutput, PPProxyTensors]: kwargs = {} if self.support_pp: kwargs["pp_proxy_tensors"] = pp_proxy_tensors