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, 06:30 UTC
[Bugfix] Fix race in non-blocking num_accepted_tokens GPU->CPU copy (#36013)
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Thomas Parnell committed
Mar 4, 2026 at 22:52 UTC
be0a3f7570726ca49cc9b53f9b48175418bddda0
1 file changed
+6
vllm/v1/worker/gpu_model_runner.py
+6
@@ -727,8 +727,10 @@ class GPUModelRunner(
727
self.draft_token_ids_copy_stream: torch.cuda.Stream | None = None
728
self.valid_sampled_token_count_cpu: torch.Tensor | None = None
729
self.draft_token_ids_cpu: torch.Tensor | None = None
730
+ self.num_accepted_tokens_event: torch.Event | None = None
731
if self.num_spec_tokens:
732
self.draft_token_ids_event = torch.Event()
733
+ self.num_accepted_tokens_event = torch.Event()
734
self.draft_token_ids_copy_stream = torch.cuda.Stream()
735
self.draft_token_ids_cpu = torch.empty(
736
(self.max_num_reqs, self.num_spec_tokens),
@@ -1229,6 +1231,8 @@ class GPUModelRunner(
1231
self.input_batch.num_accepted_tokens_cpu_tensor[:num_reqs].copy_(
1232
self.num_accepted_tokens.gpu[:num_reqs], non_blocking=True
1233
)
1234
+ assert self.num_accepted_tokens_event is not None
1235
+ self.num_accepted_tokens_event.record()
1236
1237
def _update_streaming_request(
1238
self, req_id: str, new_req_data: NewRequestData
@@ -1773,6 +1777,8 @@ class GPUModelRunner(
1777
max_seq_len = self.seq_lens.np[:num_reqs].max().item()
1778
1779
if use_spec_decode:
1780
+ if self.num_accepted_tokens_event is not None:
1781
+ self.num_accepted_tokens_event.synchronize()
1782
self.num_accepted_tokens.np[:num_reqs] = (
1783
self.input_batch.num_accepted_tokens_cpu[:num_reqs]
1784
)