From 3c8676aebbd54ccf5d666b9c493d28badd3f970f Mon Sep 17 00:00:00 2001 From: YiSheng5 Date: Fri, 14 Aug 2026 17:00:23 +0800 Subject: [PATCH] [PP][XPU]Overlap async-scheduling PP sampled-token broadcast with compute (#51650) Signed-off-by: yisheng Co-authored-by: Kunshang Ji --- vllm/v1/worker/gpu_model_runner.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index e81913f928..976d5df3bc 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -586,6 +586,9 @@ class GPUModelRunner( # Async scheduling self.use_async_scheduling = self.scheduler_config.async_scheduling + # Async PP broadcast of sampled token ids, waited on in _prepare_input_ids. + self._pp_recv_work: torch.distributed.Work | None = None + # Sampler self.sampler = Sampler( logprobs_mode=self.model_config.logprobs_mode, @@ -1843,6 +1846,11 @@ class GPUModelRunner( (-1 for new requests). """ + # Sync the async PP broadcast before reading sampled tokens. + if self._pp_recv_work is not None: + self._pp_recv_work.wait() + self._pp_recv_work = None + if self.input_batch.prev_sampled_token_ids is None: # Normal scheduling case self.input_ids.copy_to_gpu(total_num_scheduled_tokens) @@ -4963,7 +4971,9 @@ class GPUModelRunner( recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device) # skip for chunked prefill. if not self._is_all_reqs_chunked_prefill(): - torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group) + self._pp_recv_work = torch.distributed.broadcast( + recv, src=pp.last_rank, group=pp.device_group, async_op=True + ) self.input_batch.prev_sampled_token_ids = recv # construct `prev_req_id_to_index` here so `_prepare_input_ids`