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, 20:15 UTC
[Perf] Optimize pooling model redundant copy, 1.8% throughput improvement (#35127)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Feb 24, 2026 at 07:13 UTC
14561fabfd39030168a3365327d921ae5c49bb0c
1 file changed
+51
-15
vllm/v1/worker/gpu_model_runner.py
+51
-15
@@ -95,7 +95,6 @@ from vllm.sequence import IntermediateTensors
95
from vllm.tasks import GenerationTask, PoolingTask, SupportedTask
96
from vllm.tracing import instrument
97
from vllm.utils import length_from_prompt_token_ids_or_embeds
98
-from vllm.utils.jsontree import json_map_leaves
98
from vllm.utils.math_utils import cdiv, round_up
99
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
100
from vllm.utils.nvtx_pytorch_hooks import PytHooks
@@ -268,6 +267,51 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
267
return output
268
269
270
+def _copy_pooler_output_to_cpu(
271
+ raw_pooler_output: PoolerOutput, finished_mask: list[bool]
272
+) -> list[torch.Tensor | None]:
273
+ num_reqs = len(finished_mask)
274
+
275
+ if isinstance(raw_pooler_output, torch.Tensor):
276
+ if raw_pooler_output.shape[0] != num_reqs:
277
+ raise ValueError(
278
+ "Pooler output batch size does not match finished mask size: "
279
+ f"{raw_pooler_output.shape[0]} != {num_reqs}."
280
+ )
281
+
282
+ num_finished = sum(finished_mask)
283
+ if num_finished == 0:
284
+ return [None] * num_reqs
285
+ if num_finished == num_reqs:
286
+ return list(raw_pooler_output.to("cpu", non_blocking=True))
287
+
288
+ # partial finished
289
+ finished_indices = [i for i, include in enumerate(finished_mask) if include]
290
+ index_tensor = torch.tensor(
291
+ finished_indices, device=raw_pooler_output.device, dtype=torch.long
292
+ )
293
+ finished_outputs = raw_pooler_output.index_select(0, index_tensor).to(
294
+ "cpu", non_blocking=True
295
+ )
296
+ partial_pooler_output: list[torch.Tensor | None] = [None] * num_reqs
297
+ for i, out in zip(finished_indices, finished_outputs):
298
+ partial_pooler_output[i] = out
299
+ return partial_pooler_output
300
+
301
+ assert isinstance(raw_pooler_output, list)
302
+ if len(raw_pooler_output) != num_reqs:
303
+ raise ValueError(
304
+ "Pooler output batch size does not match finished mask size: "
305
+ f"{len(raw_pooler_output)} != {num_reqs}."
306
+ )
307
+
308
+ pooler_output: list[torch.Tensor | None] = [None] * num_reqs
309
+ for i, (out, include) in enumerate(zip(raw_pooler_output, finished_mask)):
310
+ if include and out is not None:
311
+ pooler_output[i] = out.to("cpu", non_blocking=True)
312
+ return pooler_output
313
+
314
+
315
class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
316
def __init__(
317
self,
@@ -289,15 +333,11 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
333
default_stream = torch.cuda.current_stream()
334
with torch.cuda.stream(async_output_copy_stream):
335
async_output_copy_stream.wait_stream(default_stream)
292
- raw_pooler_output_cpu = json_map_leaves(
293
- lambda x: None if x is None else x.to("cpu", non_blocking=True),
294
- self._raw_pooler_output,
336
+ self._model_runner_output.pooler_output = _copy_pooler_output_to_cpu(
337
+ raw_pooler_output=self._raw_pooler_output,
338
+ finished_mask=finished_mask,
339
)
340
self.async_copy_ready_event.record()
297
- self._model_runner_output.pooler_output = [
298
- out if include else None
299
- for out, include in zip(raw_pooler_output_cpu, finished_mask)
300
- ]
341
342
def get_output(self) -> ModelRunnerOutput:
343
"""Copy the device tensors to the host and return a ModelRunnerOutput.
@@ -2705,14 +2745,10 @@ class GPUModelRunner(
2745
async_output_copy_stream=self.async_output_copy_stream,
2746
)
2747
2708
- raw_pooler_output = json_map_leaves(
2709
- lambda x: None if x is None else x.to("cpu", non_blocking=True),
2710
- raw_pooler_output,
2748
+ model_runner_output.pooler_output = _copy_pooler_output_to_cpu(
2749
+ raw_pooler_output=raw_pooler_output,
2750
+ finished_mask=finished_mask,
2751
)
2712
- model_runner_output.pooler_output = [
2713
- out if include else None
2714
- for out, include in zip(raw_pooler_output, finished_mask)
2715
- ]
2752
self._sync_device()
2753
2754
return model_runner_output