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, 05:00 UTC
[Model Runner V2] Minor cleanup for PP (#34666)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Feb 16, 2026 at 19:15 UTC
04925b2202efccb51e47fa944660dc97a8e86444
2 files changed
+53
-58
vllm/v1/worker/gpu/model_runner.py
+47
-43
@@ -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, get_pp_handler
60
+from vllm.v1.worker.gpu.pp_handler import PPHandler
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
@@ -184,9 +184,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
184
185
# Pipeline parallelism.
186
self.use_pp = self.parallel_config.pipeline_parallel_size > 1
187
- self.pp_handler: PPHandler | None = (
188
- get_pp_handler(self.parallel_config) if self.use_pp else None
189
- )
187
+ if self.use_pp:
188
+ self.is_first_pp_rank = get_pp_group().is_first_rank
189
+ self.is_last_pp_rank = get_pp_group().is_last_rank
190
+ self.pp_handler: PPHandler | None = PPHandler(self.device)
191
+ else:
192
+ self.is_first_pp_rank = True
193
+ self.is_last_pp_rank = True
194
+ self.pp_handler = None
195
196
def update_max_model_len(self, max_model_len: int) -> None:
197
self.max_model_len = max_model_len
@@ -318,7 +323,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
323
324
# For non-first PP ranks, create dummy intermediate_tensors.
325
intermediate_tensors = None
321
- if self.use_pp and not get_pp_group().is_first_rank:
326
+ if not self.is_first_pp_rank:
327
intermediate_tensors = self.model.make_empty_intermediate_tensors(
328
batch_size=num_tokens,
329
dtype=self.model_config.dtype,
@@ -335,7 +340,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
340
self.kv_connector.set_disabled(False)
341
342
# Non-last PP ranks don't produce output for sampling.
338
- if self.use_pp and not get_pp_group().is_last_rank:
343
+ if not self.is_last_pp_rank:
344
return None, None
345
346
assert self.execute_model_state is not None
@@ -373,20 +378,23 @@ class GPUModelRunner(LoRAModelRunnerMixin):
378
hidden_states, sample_hidden_states = self._dummy_run(
379
self.max_num_tokens, skip_attn=True
380
)
381
+
382
# Only run sampler on last PP rank (non-last ranks return None).
377
- if not self.use_pp or get_pp_group().is_last_rank:
383
+ if self.is_last_pp_rank:
384
assert sample_hidden_states is not None
385
self._dummy_sampler_run(sample_hidden_states)
380
- if self.do_spec_decode:
381
- num_tokens_across_dp = make_num_tokens_across_dp(
382
- self.parallel_config.data_parallel_size, self.max_num_tokens
383
- )
384
- self.speculator.run_model(
385
- self.max_num_tokens,
386
- attn_metadata=None,
387
- slot_mappings=None,
388
- num_tokens_across_dp=num_tokens_across_dp,
389
- )
386
+
387
+ if self.do_spec_decode:
388
+ num_tokens_across_dp = make_num_tokens_across_dp(
389
+ self.parallel_config.data_parallel_size, self.max_num_tokens
390
+ )
391
+ self.speculator.run_model(
392
+ self.max_num_tokens,
393
+ attn_metadata=None,
394
+ slot_mappings=None,
395
+ num_tokens_across_dp=num_tokens_across_dp,
396
+ )
397
+
398
torch.cuda.synchronize()
399
del hidden_states, sample_hidden_states
400
gc.collect()
@@ -890,9 +898,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
898
self._set_active_loras(*lora_inputs)
899
900
# Only first PP rank prepares multimodal embeddings.
893
- if self.supports_mm_inputs and (
894
- not self.use_pp or get_pp_group().is_first_rank
895
- ):
901
+ if self.supports_mm_inputs and self.is_first_pp_rank:
902
mm_embeds, is_mm_embed = self.get_mm_embeddings(
903
scheduler_output.scheduled_encoder_inputs, input_batch
904
)
@@ -935,6 +941,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
941
assert input_batch.mrope_positions is not None
942
positions = input_batch.mrope_positions
943
944
+ if self.is_first_pp_rank:
945
+ input_ids = input_batch.input_ids
946
+ inputs_embeds = input_batch.inputs_embeds
947
+ assert intermediate_tensors is None
948
+ else:
949
+ input_ids = None
950
+ inputs_embeds = None
951
+ assert intermediate_tensors is not None
952
+
953
with set_forward_context(
954
input_batch.attn_metadata,
955
self.vllm_config,
@@ -945,25 +960,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
960
slot_mapping=input_batch.slot_mappings,
961
):
962
self.kv_connector.pre_forward(scheduler_output)
948
- if self.use_pp and not get_pp_group().is_first_rank:
949
- # Non-first PP rank: forward with intermediate tensors.
950
- assert intermediate_tensors is not None
951
- hidden_states = self.model(
952
- input_ids=None,
953
- positions=positions,
954
- inputs_embeds=None,
955
- intermediate_tensors=intermediate_tensors,
956
- )
957
- else:
958
- hidden_states = self.model(
959
- input_ids=input_batch.input_ids,
960
- positions=positions,
961
- inputs_embeds=input_batch.inputs_embeds,
962
- )
963
+ hidden_states = self.model(
964
+ input_ids=input_ids,
965
+ positions=positions,
966
+ inputs_embeds=inputs_embeds,
967
+ intermediate_tensors=intermediate_tensors,
968
+ )
969
970
kv_connector_output = self.kv_connector.post_forward(scheduler_output)
971
966
- if self.use_pp and not get_pp_group().is_last_rank:
972
+ if not self.is_last_pp_rank:
973
# Non-last PP rank: return IntermediateTensors for sending.
974
assert isinstance(hidden_states, IntermediateTensors)
975
hidden_states.kv_connector_output = kv_connector_output
@@ -986,16 +992,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
992
# Non-last PP rank: hidden_states is None because this rank produced
993
# IntermediateTensors instead of final hidden states. Receive the
994
# sampled tokens broadcast by the last rank and update local state.
989
- if self.use_pp and not get_pp_group().is_last_rank:
995
+ if not self.is_last_pp_rank:
996
assert self.pp_handler is not None
997
received = self.pp_handler.maybe_receive_sampled_tokens(
992
- input_batch.num_reqs,
993
- self.device,
994
- max_sample_len=self.num_speculative_steps + 1,
998
+ input_batch.num_reqs, max_sample_len=self.num_speculative_steps + 1
999
)
996
- if received is not None:
997
- sampled, num_sampled, num_rejected = received
998
- self.postprocess(input_batch, sampled, num_sampled, num_rejected)
1000
+ assert received is not None
1001
+ sampled, num_sampled, num_rejected = received
1002
+ self.postprocess(input_batch, sampled, num_sampled, num_rejected)
1003
return None
1004
1005
# Last rank: sample tokens
vllm/v1/worker/gpu/pp_handler.py
+6
-15
@@ -15,6 +15,9 @@ class PPHandler:
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,
@@ -59,7 +62,6 @@ class PPHandler:
62
def maybe_receive_sampled_tokens(
63
self,
64
num_reqs: int,
62
- device: torch.device,
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.
@@ -84,7 +86,7 @@ class PPHandler:
86
return None
87
88
sampled_tokens = torch.empty(
87
- num_reqs, max_sample_len, dtype=torch.int64, device=device
89
+ num_reqs, max_sample_len, dtype=torch.int64, device=self.device
90
)
91
torch.distributed.broadcast(
92
sampled_tokens,
@@ -93,27 +95,16 @@ class PPHandler:
95
)
96
# NOTE: num_sampled/num_rejected are only needed
97
# for speculative decoding.
96
- num_sampled = torch.empty(num_reqs, dtype=torch.int32, device=device)
98
+ num_sampled = torch.empty(num_reqs, dtype=torch.int32, device=self.device)
99
torch.distributed.broadcast(
100
num_sampled,
101
src=pp.last_rank,
102
group=pp.device_group,
103
)
102
- num_rejected = torch.empty(num_reqs, dtype=torch.int32, device=device)
104
+ num_rejected = torch.empty(num_reqs, dtype=torch.int32, device=self.device)
105
torch.distributed.broadcast(
106
num_rejected,
107
src=pp.last_rank,
108
group=pp.device_group,
109
)
110
return sampled_tokens, num_sampled, num_rejected
109
-
110
-
111
-def get_pp_handler(parallel_config) -> PPHandler:
112
- """Factory function to create PPHandler.
113
-
114
- Must only be called when PP is enabled (pp_size > 1).
115
- """
116
- assert parallel_config.pipeline_parallel_size > 1, (
117
- "PPHandler should not be created when pipeline parallelism is disabled."
118
- )
119
- return PPHandler()