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 7, 06:30 UTC
[Model Runner V2] Fix pooling (#36019)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Mar 4, 2026 at 10:53 UTC
417fd28fb125cbb166ef3ada187d06d0c8dd0d30
2 files changed
+2
-2
vllm/v1/worker/gpu/async_utils.py
+1
-1
@@ -95,8 +95,8 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):
95
self.copy_event.record(copy_stream)
96
97
def get_output(self) -> ModelRunnerOutput:
98
+ pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
99
self.copy_event.synchronize()
99
- pooler_output = self.pooler_output_cpu.unbind(dim=0)
100
if self.is_valid_cpu is not None:
101
is_valid_cpu = self.is_valid_cpu.tolist()
102
for i, is_valid in enumerate(is_valid_cpu):
vllm/v1/worker/gpu/model_runner.py
+1
-1
@@ -1117,7 +1117,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1117
# The prior execute_model call must have failed.
1118
return None
1119
1120
- input_batch, _, _, _, hidden_states, _, kv_connector_output = (
1120
+ input_batch, _, _, _, hidden_states, _, kv_connector_output, _ = (
1121
self.execute_model_state
1122
)
1123
self.execute_model_state = None