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 9, 11:15 UTC
[Feature] Fully support for async scheduling + PP, 30.8% E2E throughput improvement, 31.8% TPOT improvement (#32618)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>
Wentao Ye committed
Jan 28, 2026 at 15:30 UTC
3e440786afe763e892e12125ee7529f95f141c54
4 files changed
+103
-23
tests/test_config.py
+16
@@ -42,6 +42,22 @@ def test_compile_config_repr_succeeds():
42
assert "inductor_passes" in val
43
44
45
+def test_async_scheduling_with_pipeline_parallelism_is_allowed():
46
+ cfg = VllmConfig(
47
+ scheduler_config=SchedulerConfig(
48
+ max_model_len=8192,
49
+ is_encoder_decoder=False,
50
+ async_scheduling=True,
51
+ ),
52
+ parallel_config=ParallelConfig(
53
+ pipeline_parallel_size=2,
54
+ distributed_executor_backend="mp",
55
+ nnodes=2,
56
+ ),
57
+ )
58
+ assert cfg.scheduler_config.async_scheduling is True
59
+
60
+
61
@dataclass
62
class _TestConfigFields:
63
a: int
tests/v1/core/test_scheduler.py
+15
@@ -127,6 +127,21 @@ def test_schedule_multimodal_requests():
127
assert len(encoder_input) == 1
128
129
130
+def test_async_scheduling_pp_allows_rescheduling_with_output_placeholders():
131
+ """Async scheduling + PP: allow multi-step in-flight scheduling per request"""
132
+ scheduler = create_scheduler(async_scheduling=True, pipeline_parallel_size=2)
133
+ (req,) = create_requests(num_requests=1, num_tokens=8)
134
+ scheduler.add_request(req)
135
+
136
+ _ = scheduler.schedule()
137
+ assert req.num_output_placeholders > 0
138
+
139
+ # before any update_from_output, we still expect the request can be
140
+ # scheduled again (multi-step in-flight).
141
+ output = scheduler.schedule()
142
+ assert req.request_id in output.num_scheduled_tokens
143
+
144
+
145
def test_schedule_partial_requests():
146
"""Test scheduling behavior with partial requests.
147
vllm/v1/core/sched/scheduler.py
+4
-8
@@ -344,13 +344,6 @@ class Scheduler(SchedulerInterface):
344
while req_index < len(self.running) and token_budget > 0:
345
request = self.running[req_index]
346
347
- # do not schedule another step for the same request while it still has
348
- # output placeholders for PP.
349
- # TODO: support PP + async scheduling without this limit
350
- if self.use_pp and request.num_output_placeholders > 0:
351
- req_index += 1
352
- continue
353
-
347
if (
348
request.num_output_placeholders > 0
349
# This is (num_computed_tokens + 1) - (num_output_placeholders - 1).
@@ -1003,7 +996,10 @@ class Scheduler(SchedulerInterface):
996
for idx, req in enumerate(itertools.chain(running_reqs, resumed_reqs)):
997
req_id = req.request_id
998
req_ids.append(req_id)
1006
- if self.use_pp:
999
+ # NOTE: In PP+async scheduling, we consume token ids via a direct GPU
1000
+ # broadcast path (`input_batch.prev_sampled_token_ids`), so we can
1001
+ # omit this payload.
1002
+ if self.use_pp and not self.scheduler_config.async_scheduling:
1003
# When using PP, the scheduler sends the sampled tokens back,
1004
# because there's no direct communication between the first-
1005
# stage worker and the last-stage worker. Otherwise, we don't
vllm/v1/worker/gpu_model_runner.py
+68
-15
@@ -1010,20 +1010,26 @@ class GPUModelRunner(
1010
req_state.num_computed_tokens = num_computed_tokens
1011
1012
if not is_last_rank:
1013
- # When using PP, the scheduler sends the sampled tokens back,
1014
- # because there's no direct communication between the first-
1015
- # stage worker and the last-stage worker.
1016
- new_token_ids = req_data.new_token_ids[i]
1017
- # Add the sampled token(s) from the previous step (if any).
1018
- # This doesn't include "unverified" tokens like spec tokens.
1019
- num_new_tokens = (
1020
- num_computed_tokens + len(new_token_ids) - req_state.num_tokens
1021
- )
1022
- if num_new_tokens == 1:
1023
- # Avoid slicing list in most common case.
1024
- req_state.output_token_ids.append(new_token_ids[-1])
1025
- elif num_new_tokens > 0:
1026
- req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:])
1013
+ if not req_data.new_token_ids:
1014
+ # Async scheduled PP: Sampled tokens propagated via GPU broadcast.
1015
+ new_token_ids: list[int] = []
1016
+ else:
1017
+ # Non-async scheduling with PP: The scheduler sends
1018
+ # sampled token ids back because there's no direct communication
1019
+ # between the first-stage worker and the last-stage worker.
1020
+ new_token_ids = req_data.new_token_ids[i]
1021
+ # Add the sampled token(s) from the previous step (if any).
1022
+ # This doesn't include "unverified" tokens like spec tokens.
1023
+ num_new_tokens = (
1024
+ num_computed_tokens + len(new_token_ids) - req_state.num_tokens
1025
+ )
1026
+ if num_new_tokens == 1:
1027
+ # Avoid slicing list in most common case.
1028
+ req_state.output_token_ids.append(new_token_ids[-1])
1029
+ elif num_new_tokens > 0:
1030
+ req_state.output_token_ids.extend(
1031
+ new_token_ids[-num_new_tokens:]
1032
+ )
1033
elif num_output_tokens < len(req_state.output_token_ids):
1034
# Some output tokens were discarded due to a sync-KV-load
1035
# failure. Align the cached state.
@@ -3577,7 +3583,9 @@ class GPUModelRunner(
3583
self.kv_connector_output = None
3584
3585
if self.execute_model_state is None:
3580
- # Nothing to do (PP non-final rank case), output isn't used.
3586
+ # receive sampled token ids from the last PP rank.
3587
+ if self.use_async_scheduling and get_pp_group().world_size > 1:
3588
+ self._pp_receive_prev_sampled_token_ids_to_input_batch()
3589
if not kv_connector_output:
3590
return None # type: ignore[return-value]
3591
@@ -3618,6 +3626,12 @@ class GPUModelRunner(
3626
self._update_states_after_model_execute(
3627
sampler_output.sampled_token_ids, scheduler_output
3628
)
3629
+ if self.use_async_scheduling:
3630
+ pp = get_pp_group()
3631
+ if pp.world_size > 1 and pp.is_last_rank:
3632
+ self._pp_broadcast_prev_sampled_token_ids(
3633
+ sampler_output.sampled_token_ids
3634
+ )
3635
3636
self._draft_token_ids = None
3637
self._draft_token_req_ids = None
@@ -3753,6 +3767,45 @@ class GPUModelRunner(
3767
3768
return async_output
3769
3770
+ def _pp_broadcast_prev_sampled_token_ids(
3771
+ self, sampled_token_ids: torch.Tensor
3772
+ ) -> None:
3773
+ """Broadcast sampled token ids (GPU) from last PP stage"""
3774
+ pp = get_pp_group()
3775
+ assert pp.is_last_rank
3776
+ # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1].
3777
+ assert sampled_token_ids.dim() == 2 and sampled_token_ids.shape[-1] == 1, (
3778
+ "PP+async expects sampled_token_ids to have shape [num_reqs, 1]"
3779
+ )
3780
+ torch.distributed.broadcast(
3781
+ sampled_token_ids, src=pp.rank, group=pp.device_group
3782
+ )
3783
+
3784
+ def _pp_receive_prev_sampled_token_ids_to_input_batch(self) -> None:
3785
+ """Receive sampled token ids broadcast from last PP stage"""
3786
+ pp = get_pp_group()
3787
+ assert not pp.is_last_rank
3788
+ num_reqs = self.input_batch.num_reqs
3789
+ # `prev_sampled_token_ids` is expected to have shape [num_reqs, 1].
3790
+ recv = torch.empty((num_reqs, 1), dtype=torch.int32, device=self.device)
3791
+ torch.distributed.broadcast(recv, src=pp.last_rank, group=pp.device_group)
3792
+ self.input_batch.prev_sampled_token_ids = recv
3793
+
3794
+ # construct `prev_req_id_to_index` here so `_prepare_input_ids`
3795
+ # can map req_id -> previous batch row
3796
+ discard_req_indices = np.nonzero(self.discard_request_mask.np[:num_reqs])[0]
3797
+ discard_req_indices_set = set(discard_req_indices)
3798
+ prev_req_id_to_index: dict[str, int] = {}
3799
+ for i, req_id in enumerate(self.input_batch.req_ids):
3800
+ if i in discard_req_indices_set:
3801
+ continue
3802
+ prev_req_id_to_index[req_id] = i
3803
+ # PP+async scheduling: advance per-request local cached output length by
3804
+ # appending a placeholder (-1) token id.
3805
+ if (req_state := self.requests.get(req_id)) is not None:
3806
+ req_state.output_token_ids.append(-1)
3807
+ self.input_batch.prev_req_id_to_index = prev_req_id_to_index
3808
+
3809
def take_draft_token_ids(self) -> DraftTokenIds | None:
3810
if not self.num_spec_tokens or not self._draft_token_req_ids:
3811
return None