[PP][XPU]Overlap async-scheduling PP sampled-token broadcast with compute (#51650)
Signed-off-by: yisheng <yi.sheng@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
YiSheng5 committed
Aug 14, 2026 at 17:00 UTC
3c8676aebbd54ccf5d666b9c493d28badd3f970f
1 file changed
+11
-1
vllm/v1/worker/gpu_model_runner.py
+11
-1
@@ -586,6 +586,9 @@ class GPUModelRunner(
586
# Async scheduling
587
self.use_async_scheduling = self.scheduler_config.async_scheduling
588
589
+ # Async PP broadcast of sampled token ids, waited on in _prepare_input_ids.
590
+ self._pp_recv_work: torch.distributed.Work | None = None
591
+
592
# Sampler
593
self.sampler = Sampler(
594
logprobs_mode=self.model_config.logprobs_mode,
@@ -1843,6 +1846,11 @@ class GPUModelRunner(
1846
(-1 for new requests).
1847
"""
1848
1849
+ # Sync the async PP broadcast before reading sampled tokens.
1850
+ if self._pp_recv_work is not None:
1851
+ self._pp_recv_work.wait()
1852
+ self._pp_recv_work = None
1853
+
1854
if self.input_batch.prev_sampled_token_ids is None:
1855
# Normal scheduling case
1856
self.input_ids.copy_to_gpu(total_num_scheduled_tokens)
@@ -4963,7 +4971,9 @@ class GPUModelRunner(
4971
recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device)
4972
# skip for chunked prefill.
4973
if not self._is_all_reqs_chunked_prefill():
4966
- torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
4974
+ self._pp_recv_work = torch.distributed.broadcast(
4975
+ recv, src=pp.last_rank, group=pp.device_group, async_op=True
4976
+ )
4977
self.input_batch.prev_sampled_token_ids = recv
4978
4979
# construct `prev_req_id_to_index` here so `_prepare_input_ids`