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 5, 09:45 UTC
[Perf] Optimize redundant sync for pooling model, 3.7% Throughput Improvement (#39113)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 9, 2026 at 02:12 UTC
ed2f282bc8cacfbe49a86ab1dbaf7e4471ec6b18
1 file changed
+16
-9
vllm/v1/worker/gpu_model_runner.py
+16
-9
@@ -1046,6 +1046,13 @@ class GPUModelRunner(
1046
def _sync_device(self) -> None:
1047
torch.accelerator.synchronize()
1048
1049
+ def _get_or_create_async_output_copy_stream(self) -> torch.cuda.Stream:
1050
+ stream = self.async_output_copy_stream
1051
+ if stream is None:
1052
+ stream = torch.cuda.Stream()
1053
+ self.async_output_copy_stream = stream
1054
+ return stream
1055
+
1056
def _update_states(self, scheduler_output: "SchedulerOutput") -> Callable | None:
1057
"""Update the cached states and the persistent batch with the scheduler
1058
output.
@@ -3152,21 +3159,21 @@ class GPUModelRunner(
3159
model_runner_output.pooler_output = [None] * num_reqs
3160
return model_runner_output
3161
3155
- if self.use_async_scheduling:
3156
- return AsyncGPUPoolingModelRunnerOutput(
3157
- model_runner_output=model_runner_output,
3162
+ if not current_platform.is_cuda_alike():
3163
+ # cpu/xpu runners cannot use the CUDA stream/event-based wrapper.
3164
+ model_runner_output.pooler_output = _copy_pooler_output_to_cpu(
3165
raw_pooler_output=raw_pooler_output,
3166
finished_mask=finished_mask,
3160
- async_output_copy_stream=self.async_output_copy_stream,
3167
)
3168
+ self._sync_device()
3169
+ return model_runner_output
3170
3163
- model_runner_output.pooler_output = _copy_pooler_output_to_cpu(
3171
+ return AsyncGPUPoolingModelRunnerOutput(
3172
+ model_runner_output=model_runner_output,
3173
raw_pooler_output=raw_pooler_output,
3174
finished_mask=finished_mask,
3175
+ async_output_copy_stream=self._get_or_create_async_output_copy_stream(),
3176
)
3167
- self._sync_device()
3168
-
3169
- return model_runner_output
3177
3178
def _pad_for_sequence_parallelism(self, num_scheduled_tokens: int) -> int:
3179
# Pad tokens to multiple of tensor_parallel_size when
@@ -4333,7 +4340,7 @@ class GPUModelRunner(
4340
sampled_token_ids=sampler_output.sampled_token_ids,
4341
logprobs_tensors=sampler_output.logprobs_tensors,
4342
invalid_req_indices=invalid_req_indices,
4336
- async_output_copy_stream=self.async_output_copy_stream,
4343
+ async_output_copy_stream=self._get_or_create_async_output_copy_stream(),
4344
vocab_size=self.input_batch.vocab_size,
4345
)
4346
with record_function_or_nullcontext(