[MRV2] Run pooling post processing on non-final PP ranks (#56666)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Taneem Ibrahim committed
Sep 14, 2026 at 18:11 UTC
ba2ae9f23961ac67bc5c055da8c26fbd660989c6
1 file changed
+2
vllm/v1/worker/gpu_worker.py
+2
@@ -1238,6 +1238,8 @@ class Worker(WorkerBase):
1238
)
1239
self._pp_send_work = handles[1:]
1240
1241
+ if self.use_v2_model_runner and self.model_runner.is_pooling_model:
1242
+ return self.model_runner.pool() # type: ignore
1243
return None
1244
1245
def take_draft_token_ids(self) -> DraftTokenIds | None: