[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>

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*