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 10, 19:00 UTC
[Minor] Small pooler output processing optimization (#31667)
Signed-off-by: njhill <nickhill123@gmail.com>
Nick Hill committed
Jan 4, 2026 at 18:33 UTC
da436f868ae738f3c7cf8463d545813244499d0a
1 file changed
+8
-11
vllm/v1/worker/gpu_model_runner.py
+8
-11
@@ -263,7 +263,6 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
263
async_output_copy_stream: torch.cuda.Stream,
264
):
265
self._model_runner_output = model_runner_output
266
- self._finished_mask = finished_mask
266
267
# Event on the copy stream so we can synchronize the non-blocking copy.
268
self.async_copy_ready_event = torch.Event()
@@ -276,11 +275,15 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
275
default_stream = torch.cuda.current_stream()
276
with torch.cuda.stream(async_output_copy_stream):
277
async_output_copy_stream.wait_stream(default_stream)
279
- self._raw_pooler_output_cpu = json_map_leaves(
278
+ raw_pooler_output_cpu = json_map_leaves(
279
lambda x: None if x is None else x.to("cpu", non_blocking=True),
280
self._raw_pooler_output,
281
)
282
self.async_copy_ready_event.record()
283
+ self._model_runner_output.pooler_output = [
284
+ out if include else None
285
+ for out, include in zip(raw_pooler_output_cpu, finished_mask)
286
+ ]
287
288
def get_output(self) -> ModelRunnerOutput:
289
"""Copy the device tensors to the host and return a ModelRunnerOutput.
@@ -290,11 +293,6 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
293
294
# Release the device tensors once the copy has completed.
295
del self._raw_pooler_output
293
-
294
- self._model_runner_output.pooler_output = [
295
- out if include else None
296
- for out, include in zip(self._raw_pooler_output_cpu, self._finished_mask)
297
- ]
296
return self._model_runner_output
297
298
@@ -2537,8 +2535,7 @@ class GPUModelRunner(
2535
2536
model = cast(VllmModelForPooling, self.model)
2537
raw_pooler_output: PoolerOutput = model.pooler(
2540
- hidden_states=hidden_states,
2541
- pooling_metadata=pooling_metadata,
2538
+ hidden_states=hidden_states, pooling_metadata=pooling_metadata
2539
)
2540
2541
finished_mask = [
@@ -2568,12 +2565,12 @@ class GPUModelRunner(
2565
lambda x: None if x is None else x.to("cpu", non_blocking=True),
2566
raw_pooler_output,
2567
)
2571
- self._sync_device()
2572
-
2568
model_runner_output.pooler_output = [
2569
out if include else None
2570
for out, include in zip(raw_pooler_output, finished_mask)
2571
]
2572
+ self._sync_device()
2573
+
2574
return model_runner_output
2575
2576
def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int: