From e8064a96d02db70ebc1ca922bc9aba967a654483 Mon Sep 17 00:00:00 2001 From: Chris Eastwood <106503529+eastwood-c@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:55:56 -0400 Subject: [PATCH] [Bugfix][V2] record_stream idx_mapping in the PP draft broadcast (#55745) Signed-off-by: Chris Eastwood Signed-off-by: Nick Hill Co-authored-by: Nick Hill --- vllm/v1/worker/gpu/pp_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/pp_utils.py b/vllm/v1/worker/gpu/pp_utils.py index 33386c9222..8c65a23ef7 100644 --- a/vllm/v1/worker/gpu/pp_utils.py +++ b/vllm/v1/worker/gpu/pp_utils.py @@ -170,10 +170,12 @@ class PPHandler: with torch.cuda.stream(self.broadcast_stream): self.broadcast_stream.wait_stream(self.main_stream) send = draft_tokens[input_batch.idx_mapping].contiguous() + # Must record the idx_mapping tensor since it was allocated + # on the main stream. + input_batch.idx_mapping.record_stream(self.broadcast_stream) torch.distributed.broadcast( send, src=self.last_rank, group=self.broadcast_group ) - send.record_stream(self.broadcast_stream) def receive(self, input_batch: InputBatch) -> bool: """Returns True iff sampled tokens need to be gathered from *all*