A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 01:00 UTC
[Bugfix][Core] Fix stuck chunked pipeline parallelism with async scheduling (#38726)
Signed-off-by: Jing Wang <jingwang96@qq.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Jing Wang committed
Apr 18, 2026 at 00:02 UTC
747256bb5d645f28579bdb1a913f3584b6a246f6
1 file changed
+18
-4
vllm/v1/worker/gpu_model_runner.py
+18
-4
@@ -3758,6 +3758,15 @@ class GPUModelRunner(
3758
3759
return slot_mappings_by_gid, slot_mappings_by_layer
3760
3761
+ def _is_all_reqs_chunked_prefill(self) -> bool:
3762
+ """Check if all scheduled requests are marked to discard sampled tokens.
3763
+
3764
+ This is true when `discard_request_mask` is set for every scheduled
3765
+ request (e.g., for chunked prefill requests that are not the last
3766
+ prefill chunk)."""
3767
+ num_reqs = self.input_batch.num_reqs
3768
+ return bool(self.discard_request_mask.np[:num_reqs].all())
3769
+
3770
@torch.inference_mode()
3771
def execute_model(
3772
self,
@@ -4361,9 +4370,12 @@ class GPUModelRunner(
4370
assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, (
4371
"PP+async expects sampled_token_ids to have shape [num_reqs, 1]"
4372
)
4364
- torch.distributed.broadcast(
4365
- sampled_token_ids, src=pp.rank, group=pp.device_group
4366
- )
4373
+ # Skip for chunked prefill: sampled tokens are dummy
4374
+ # and will be discarded, no need to broadcast.
4375
+ if not self._is_all_reqs_chunked_prefill():
4376
+ torch.distributed.broadcast(
4377
+ sampled_token_ids, src=pp.rank, group=pp.device_group
4378
+ )
4379
4380
def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None:
4381
"""Receive sampled token ids broadcast from last PP stage"""
@@ -4372,7 +4384,9 @@ class GPUModelRunner(
4384
num_reqs = self.input_batch.num_reqs
4385
# `prev_sampled_token_ids` is expected to have shape [num_reqs, 1].
4386
recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device)
4375
- torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
4387
+ # skip for chunked prefill.
4388
+ if not self._is_all_reqs_chunked_prefill():
4389
+ torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
4390
self.input_batch.prev_sampled_token_ids = recv
4391
4392
# construct `prev_req_id_to_index` here so `_prepare_input_ids`