A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 02:00 UTC
[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>
Nick Hill committed
Nov 18, 2025 at 00:20 UTC
439368496db48d8f992ba8c606a0c0b1eebbfa69
2 files changed
+6
-5
vllm/v1/engine/core.py
+2
-1
@@ -184,6 +184,7 @@ class EngineCore:
184
vllm_config.ec_transfer_config is not None
185
and vllm_config.ec_transfer_config.is_ec_producer
186
)
187
+ self.is_pooling_model = vllm_config.model_config.runner_type == "pooling"
188
189
self.request_block_hasher: Callable[[Request], list[BlockHash]] | None = None
190
if vllm_config.cache_config.enable_prefix_caching or kv_connector is not None:
@@ -392,7 +393,7 @@ class EngineCore:
393
if not self.ec_producer:
394
model_executed = scheduler_output.total_num_scheduled_tokens > 0
395
395
- if not model_executed:
396
+ if self.is_pooling_model or not model_executed:
397
# No sampling required (no requests scheduled).
398
future = cast(Future[ModelRunnerOutput], exec_future)
399
else:
vllm/v1/executor/ray_executor.py
+4
-4
@@ -99,9 +99,9 @@ class RayDistributedExecutor(Executor):
99
# KV connector setup
100
self.has_connector = self.vllm_config.kv_transfer_config is not None
101
102
- self.ec_producer = (
103
- self.vllm_config.ec_transfer_config is not None
104
- and self.vllm_config.ec_transfer_config.is_ec_producer
102
+ self.uses_sampler = self.vllm_config.model_config.runner_type != "pooling" and (
103
+ self.vllm_config.ec_transfer_config is None
104
+ or not self.vllm_config.ec_transfer_config.is_ec_producer
105
)
106
107
self.scheduler_output: SchedulerOutput | None = None
@@ -401,7 +401,7 @@ class RayDistributedExecutor(Executor):
401
"after execute_model() returns None."
402
)
403
404
- if self.ec_producer or not scheduler_output.total_num_scheduled_tokens:
404
+ if not self.uses_sampler or not scheduler_output.total_num_scheduled_tokens:
405
# Model will not execute, call model runner immediately.
406
return self._execute_dag(scheduler_output, None, non_block)
407