refactor(runner): unify pp_proxy_tensors forward kwarg into one helper (#28382)

python/sglang/srt/model_executor/model_runner.py

@@ -3292,6 +3292,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             forward_batch_template=forward_batch,
         )
 
+    def _pp_kwargs(self, pp_proxy_tensors) -> dict:
+        """Build the pp_proxy_tensors forward kwarg, in one place.
+
+        Pipeline-parallel proxy tensors are threaded into model.forward only
+        when the model accepts them (``support_pp``).
+        """
+        return {"pp_proxy_tensors": pp_proxy_tensors} if self.support_pp else {}
+
     def forward_decode(
         self,
         forward_batch: ForwardBatch,
@@ -3315,9 +3323,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             else:
                 self.attn_backend.init_forward_metadata(forward_batch)
         # FIXME: add pp_proxy_tensors arg to all models
-        kwargs = {}
+        kwargs = self._pp_kwargs(pp_proxy_tensors)
-        if self.support_pp:
-            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
 
         # Launch forward
         ctx = (
@@ -3350,9 +3356,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput], bool
     ]:
         # Setup extra arguments
-        kwargs = {}
+        kwargs = self._pp_kwargs(pp_proxy_tensors)
-        if self.support_pp:
-            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
         if forward_batch.input_embeds is not None:
             kwargs["input_embeds"] = forward_batch.input_embeds.bfloat16()
         if (
@@ -3454,9 +3458,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         else:
             self.attn_backend.forward_metadata = None
 
-        kwargs = {}
+        kwargs = self._pp_kwargs(pp_proxy_tensors)
-        if self.support_pp:
-            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
         ctx = (
             self.device_timer.wrap(metadata={"category": "idle"})
             if self.device_timer