[Perf] Use blocking CUDA events to avoid busy polling cuda driver lock (#47081)
Signed-off-by: Jingyi Yang <girasoleyang@gmail.com>
Summer Yang committed
Jul 7, 2026 at 02:36 UTC
d3e69fd6714e9d1bb6e8e4f03157090dc32e7960
3 files changed
+13
-6
vllm/v1/worker/gpu/async_utils.py
+4
-2
@@ -24,7 +24,8 @@ class AsyncOutput(AsyncModelRunnerOutput):
24
self.model_runner_output = model_runner_output
25
self.sampler_output = sampler_output
26
self.num_sampled_tokens = num_sampled_tokens
27
- self.copy_event = torch.cuda.Event()
27
+ # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
28
+ self.copy_event = torch.cuda.Event(blocking=True)
29
30
with stream(copy_stream, main_stream):
31
copy_stream.wait_stream(main_stream)
@@ -81,7 +82,8 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput):
82
self.model_runner_output = model_runner_output
83
self.pooler_output = pooler_output
84
self.is_valid = is_valid
84
- self.copy_event = torch.cuda.Event()
85
+ # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
86
+ self.copy_event = torch.cuda.Event(blocking=True)
87
88
with stream(copy_stream, main_stream):
89
copy_stream.wait_stream(main_stream)
vllm/v1/worker/gpu/spec_decode/utils.py
+2
-1
@@ -12,7 +12,8 @@ class DraftTokensHandler:
12
def __init__(self, device: torch.device | None = None):
13
self.device = device
14
self.copy_stream = torch.cuda.Stream(device)
15
- self.copy_event = torch.cuda.Event()
15
+ # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
16
+ self.copy_event = torch.cuda.Event(blocking=True)
17
18
self.req_ids: list[str] = []
19
self.draft_tokens_np: np.ndarray | None = None
vllm/v1/worker/gpu_model_runner.py
+7
-3
@@ -259,7 +259,8 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
259
self._invalid_req_indices = invalid_req_indices
260
261
# Event on the copy stream so we can synchronize the non-blocking copy.
262
- self.async_copy_ready_event = torch.Event()
262
+ # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
263
+ self.async_copy_ready_event = torch.cuda.Event(blocking=True)
264
265
# Keep a reference to the device tensor to avoid it being
266
# deallocated until we finish copying it to the host.
@@ -392,7 +393,8 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput):
393
self._model_runner_output = model_runner_output
394
395
# Event on the copy stream so we can synchronize the non-blocking copy.
395
- self.async_copy_ready_event = torch.Event()
396
+ # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
397
+ self.async_copy_ready_event = torch.cuda.Event(blocking=True)
398
399
# Keep a reference to the device tensors to avoid them being
400
# deallocated until we finish copying it to the host.
@@ -718,7 +720,9 @@ class GPUModelRunner(
720
self.prepare_inputs_event: torch.Event | None = None
721
if self.use_async_scheduling:
722
self.async_output_copy_stream = torch.cuda.Stream()
721
- self.prepare_inputs_event = torch.Event()
723
+ # Blocking (sleep) event to avoid busy-polling the CUDA driver lock;
724
+ # under TP contention that spin can balloon and make the rank a straggler.
725
+ self.prepare_inputs_event = torch.cuda.Event(blocking=True)
726
727
# self.cudagraph_batch_sizes sorts in ascending order.
728
if (