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 8, 03:45 UTC
[Model Runner V2] Further simplification for PP (#34724)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 17, 2026 at 15:18 UTC
be3af2d29e2507f32b2190fe015cd6609b348caa
3 files changed
+46
-118
vllm/v1/worker/gpu/model_runner.py
+3
-9
@@ -57,7 +57,7 @@ from vllm.v1.worker.gpu.kv_connector import (
57
from vllm.v1.worker.gpu.lora_utils import LoraState
58
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
59
from vllm.v1.worker.gpu.mm.mrope_utils import MRopeState
60
-from vllm.v1.worker.gpu.pp_handler import PPHandler
60
+from vllm.v1.worker.gpu.pp_utils import pp_broadcast, pp_receive
61
from vllm.v1.worker.gpu.sample.output import SamplerOutput
62
from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
63
from vllm.v1.worker.gpu.sample.sampler import Sampler
@@ -185,11 +185,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
185
if self.use_pp:
186
self.is_first_pp_rank = get_pp_group().is_first_rank
187
self.is_last_pp_rank = get_pp_group().is_last_rank
188
- self.pp_handler: PPHandler | None = PPHandler(self.device)
188
else:
189
self.is_first_pp_rank = True
190
self.is_last_pp_rank = True
192
- self.pp_handler = None
191
192
def update_max_model_len(self, max_model_len: int) -> None:
193
self.max_model_len = max_model_len
@@ -987,8 +985,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
985
# IntermediateTensors instead of final hidden states. Receive the
986
# sampled tokens broadcast by the last rank and update local state.
987
if not self.is_last_pp_rank:
990
- assert self.pp_handler is not None
991
- received = self.pp_handler.maybe_receive_sampled_tokens(
988
+ received = pp_receive(
989
input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
990
)
991
assert received is not None
@@ -1003,10 +1000,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1000
1001
# Broadcast to non-last PP ranks (handles spec decode multi-token).
1002
if self.use_pp:
1006
- assert self.pp_handler is not None
1007
- self.pp_handler.maybe_broadcast_sampled_tokens(
1008
- sampler_output, num_sampled, num_rejected
1009
- )
1003
+ pp_broadcast(sampler_output.sampled_token_ids, num_sampled, num_rejected)
1004
1005
prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
1006
self.model.compute_logits,
vllm/v1/worker/gpu/pp_handler.py
deleted
-109
@@ -1,109 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Pipeline Parallelism handler for V2 Model Runner."""
4
-
5
-import torch
6
-
7
-from vllm.distributed.parallel_state import get_pp_group
8
-from vllm.v1.worker.gpu.sample.output import SamplerOutput
9
-
10
-
11
-class PPHandler:
12
- """Pipeline parallelism handler for Model Runner V2.
13
-
14
- Manages sampled token synchronization between PP ranks.
15
- Only instantiated when PP is enabled (pp_size > 1).
16
- """
17
-
18
- def __init__(self, device: torch.device):
19
- self.device = device
20
-
21
- def maybe_broadcast_sampled_tokens(
22
- self,
23
- sampler_output: SamplerOutput,
24
- num_sampled: torch.Tensor,
25
- num_rejected: torch.Tensor,
26
- ) -> None:
27
- """Broadcast sampled tokens from the last PP rank to all other ranks.
28
-
29
- No-ops if this is not the last rank.
30
-
31
- Broadcasts sampled_token_ids [num_reqs, max_sample_len], num_sampled
32
- [num_reqs], and num_rejected [num_reqs] to support both regular decode
33
- and speculative decoding.
34
-
35
- Args:
36
- sampler_output: SamplerOutput from sampling.
37
- num_sampled: Number of accepted tokens per request.
38
- num_rejected: Number of rejected tokens per request.
39
- """
40
- pp = get_pp_group()
41
- if not pp.is_last_rank:
42
- return
43
-
44
- torch.distributed.broadcast(
45
- sampler_output.sampled_token_ids.contiguous(),
46
- src=pp.last_rank,
47
- group=pp.device_group,
48
- )
49
- # NOTE: num_sampled/num_rejected are only needed
50
- # for speculative decoding.
51
- torch.distributed.broadcast(
52
- num_sampled.contiguous(),
53
- src=pp.last_rank,
54
- group=pp.device_group,
55
- )
56
- torch.distributed.broadcast(
57
- num_rejected.contiguous(),
58
- src=pp.last_rank,
59
- group=pp.device_group,
60
- )
61
-
62
- def maybe_receive_sampled_tokens(
63
- self,
64
- num_reqs: int,
65
- max_sample_len: int = 1,
66
- ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
67
- """Receive sampled tokens broadcast by the last PP rank.
68
-
69
- Returns None if this is the last rank (which samples, not receives).
70
-
71
- Args:
72
- num_reqs: Number of requests in the batch.
73
- max_sample_len: Maximum number of tokens sampled per request
74
- (1 for regular decode, >1 for speculative decoding).
75
-
76
- Returns:
77
- None if called on last rank.
78
- Otherwise, tuple of (sampled_tokens, num_sampled, num_rejected):
79
- - sampled_tokens: shape [num_reqs, max_sample_len]
80
- - num_sampled: shape [num_reqs]
81
- - num_rejected: shape [num_reqs]
82
- """
83
- pp = get_pp_group()
84
- if pp.is_last_rank:
85
- return None
86
-
87
- sampled_tokens = torch.empty(
88
- num_reqs, max_sample_len, dtype=torch.int64, device=self.device
89
- )
90
- torch.distributed.broadcast(
91
- sampled_tokens,
92
- src=pp.last_rank,
93
- group=pp.device_group,
94
- )
95
- # NOTE: num_sampled/num_rejected are only needed
96
- # for speculative decoding.
97
- num_sampled = torch.empty(num_reqs, dtype=torch.int32, device=self.device)
98
- torch.distributed.broadcast(
99
- num_sampled,
100
- src=pp.last_rank,
101
- group=pp.device_group,
102
- )
103
- num_rejected = torch.empty(num_reqs, dtype=torch.int32, device=self.device)
104
- torch.distributed.broadcast(
105
- num_rejected,
106
- src=pp.last_rank,
107
- group=pp.device_group,
108
- )
109
- return sampled_tokens, num_sampled, num_rejected
vllm/v1/worker/gpu/pp_utils.py
new
+43
@@ -0,0 +1,43 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Pipeline Parallelism utils for V2 Model Runner."""
4
+
5
+import torch
6
+
7
+from vllm.distributed.parallel_state import get_pp_group
8
+
9
+
10
+def pp_broadcast(
11
+ sampled_token_ids: torch.Tensor,
12
+ num_sampled: torch.Tensor,
13
+ num_rejected: torch.Tensor,
14
+) -> None:
15
+ pp = get_pp_group()
16
+ if not pp.is_last_rank:
17
+ return
18
+
19
+ assert sampled_token_ids.dtype == torch.int64
20
+ torch.distributed.broadcast(
21
+ sampled_token_ids.contiguous(), src=pp.last_rank, group=pp.device_group
22
+ )
23
+
24
+ combined = torch.stack((num_sampled, num_rejected), dim=0)
25
+ torch.distributed.broadcast(combined, src=pp.last_rank, group=pp.device_group)
26
+
27
+
28
+def pp_receive(
29
+ num_reqs: int, max_sample_len: int = 1
30
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
31
+ pp = get_pp_group()
32
+ if pp.is_last_rank:
33
+ return None
34
+
35
+ sampled_tokens = torch.empty(
36
+ num_reqs, max_sample_len, dtype=torch.int64, device=pp.device
37
+ )
38
+ torch.distributed.broadcast(sampled_tokens, src=pp.last_rank, group=pp.device_group)
39
+
40
+ combined = torch.empty(2, num_reqs, dtype=torch.int32, device=pp.device)
41
+ torch.distributed.broadcast(combined, src=pp.last_rank, group=pp.device_group)
42
+ num_sampled, num_rejected = combined.unbind(dim=0)
43
+ return sampled_tokens, num_sampled, num_rejected