diff --git a/vllm/v1/worker/gpu/async_utils.py b/vllm/v1/worker/gpu/async_utils.py index b3d6f5e4d..e4659104f 100644 --- a/vllm/v1/worker/gpu/async_utils.py +++ b/vllm/v1/worker/gpu/async_utils.py @@ -24,7 +24,8 @@ class AsyncOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.sampler_output = sampler_output self.num_sampled_tokens = num_sampled_tokens - self.copy_event = torch.cuda.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) @@ -81,7 +82,8 @@ class AsyncPoolingOutput(AsyncModelRunnerOutput): self.model_runner_output = model_runner_output self.pooler_output = pooler_output self.is_valid = is_valid - self.copy_event = torch.cuda.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) with stream(copy_stream, main_stream): copy_stream.wait_stream(main_stream) diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 4ab45b2ae..8e73ec8ab 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -12,7 +12,8 @@ class DraftTokensHandler: def __init__(self, device: torch.device | None = None): self.device = device self.copy_stream = torch.cuda.Stream(device) - self.copy_event = torch.cuda.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.copy_event = torch.cuda.Event(blocking=True) self.req_ids: list[str] = [] self.draft_tokens_np: np.ndarray | None = None diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 74938a823..8756c97a1 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -254,7 +254,8 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): self._invalid_req_indices = invalid_req_indices # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) # Keep a reference to the device tensor to avoid it being # deallocated until we finish copying it to the host. @@ -375,7 +376,8 @@ class AsyncGPUPoolingModelRunnerOutput(AsyncModelRunnerOutput): self._model_runner_output = model_runner_output # Event on the copy stream so we can synchronize the non-blocking copy. - self.async_copy_ready_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. + self.async_copy_ready_event = torch.cuda.Event(blocking=True) # Keep a reference to the device tensors to avoid them being # deallocated until we finish copying it to the host. @@ -697,7 +699,9 @@ class GPUModelRunner( self.prepare_inputs_event: torch.Event | None = None if self.use_async_scheduling: self.async_output_copy_stream = torch.cuda.Stream() - self.prepare_inputs_event = torch.Event() + # Blocking (sleep) event to avoid busy-polling the CUDA driver lock; + # under TP contention that spin can balloon and make the rank a straggler. + self.prepare_inputs_event = torch.cuda.Event(blocking=True) # self.cudagraph_batch_sizes sorts in ascending order. if (