[Bugfix][V2] record_stream idx_mapping in the PP draft broadcast (#55745)
Signed-off-by: Chris Eastwood <chris.eastwood@pwn4g3.dev> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Chris Eastwood committed
Sep 9, 2026 at 15:55 UTC
e8064a96d02db70ebc1ca922bc9aba967a654483
1 file changed
+3
-1
vllm/v1/worker/gpu/pp_utils.py
+3
-1
@@ -170,10 +170,12 @@ class PPHandler:
170
with torch.cuda.stream(self.broadcast_stream):
171
self.broadcast_stream.wait_stream(self.main_stream)
172
send = draft_tokens[input_batch.idx_mapping].contiguous()
173
+ # Must record the idx_mapping tensor since it was allocated
174
+ # on the main stream.
175
+ input_batch.idx_mapping.record_stream(self.broadcast_stream)
176
torch.distributed.broadcast(
177
send, src=self.last_rank, group=self.broadcast_group
178
)
176
- send.record_stream(self.broadcast_stream)
179
180
def receive(self, input_batch: InputBatch) -> bool:
181
"""Returns True iff sampled tokens need to be gathered from *all*