[BugFix] Fix PP/async scheduling with pooling models (#28899)

Signed-off-by: Nick Hill <nhill@redhat.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
This commit is contained in:
Nick Hill
2025-11-18 00:20:45 -08:00
committed by GitHub
parent 896e41ae04
commit 439368496d
2 changed files with 6 additions and 5 deletions

View File

@@ -99,9 +99,9 @@ class RayDistributedExecutor(Executor):
# KV connector setup
self.has_connector = self.vllm_config.kv_transfer_config is not None
self.ec_producer = (
self.vllm_config.ec_transfer_config is not None
and self.vllm_config.ec_transfer_config.is_ec_producer
self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
self.vllm_config.ec_transfer_config is None
or not self.vllm_config.ec_transfer_config.is_ec_producer
)
self.scheduler_output: SchedulerOutput | None = None
@@ -401,7 +401,7 @@ class RayDistributedExecutor(Executor):
"after execute_model() returns None."
)
if self.ec_producer or not scheduler_output.total_num_scheduled_tokens:
if not self.uses_sampler or not scheduler_output.total_num_scheduled_tokens:
# Model will not execute, call model runner immediately.
return self._execute_dag(scheduler_output, None, non_block)