From 67db2ac3e74ef9ed9a38dc7e280c2503faf3051c Mon Sep 17 00:00:00 2001
From: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Date: Thu, 18 Jun 2026 02:23:47 -0700
Subject: [PATCH] refactor(runner): unify pp_proxy_tensors forward kwarg into
 one helper (#28382)

---
 .../sglang/srt/model_executor/model_runner.py | 20 ++++++++++---------
 1 file changed, 11 insertions(+), 9 deletions(-)

diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index c32a1c06f..63ffcb03a 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/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 = {}
-        if self.support_pp:
-            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
+        kwargs = self._pp_kwargs(pp_proxy_tensors)
 
         # Launch forward
         ctx = (
@@ -3350,9 +3356,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput], bool
     ]:
         # Setup extra arguments
-        kwargs = {}
-        if self.support_pp:
-            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
+        kwargs = self._pp_kwargs(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 = {}
-        if self.support_pp:
-            kwargs["pp_proxy_tensors"] = pp_proxy_tensors
+        kwargs = self._pp_kwargs(pp_proxy_tensors)
         ctx = (
             self.device_timer.wrap(metadata={"category": "idle"})
             if self.device_timer