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
[Core] Pipeline Parallel support for Model Runner V2 (#33960)
Signed-off-by: Zhanqiu Hu <zh338@cornell.edu>
zhanqiuhu committed
Feb 16, 2026 at 20:48 UTC
9a8853f781d0d5c0964eba5ea61ac25d9eca0185
2 files changed
+221
-17
vllm/v1/worker/gpu/model_runner.py
+102
-17
@@ -3,7 +3,6 @@
3
import gc
4
import time
5
from copy import deepcopy
6
-from typing import Any
6
7
import numpy as np
8
import torch
@@ -11,11 +10,15 @@ import torch.nn as nn
10
11
from vllm.config import VllmConfig
12
from vllm.config.compilation import CUDAGraphMode
14
-from vllm.distributed.parallel_state import prepare_communication_buffer_for_model
13
+from vllm.distributed.parallel_state import (
14
+ get_pp_group,
15
+ prepare_communication_buffer_for_model,
16
+)
17
from vllm.forward_context import set_forward_context
18
from vllm.logger import init_logger
19
from vllm.model_executor.model_loader import get_model_loader
20
from vllm.multimodal import MULTIMODAL_REGISTRY
21
+from vllm.sequence import IntermediateTensors
22
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
23
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
24
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -54,6 +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
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
@@ -178,6 +182,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
182
# KV Connector if configured.
183
self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
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
+ )
190
+
191
def update_max_model_len(self, max_model_len: int) -> None:
192
self.max_model_len = max_model_len
193
self.req_states.max_model_len = max_model_len
@@ -290,7 +300,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
300
@torch.inference_mode()
301
def _dummy_run(
302
self, num_tokens: int, *args, skip_attn: bool = True, **kwargs
293
- ) -> tuple[torch.Tensor, torch.Tensor]:
303
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
304
# Create a dummy scheduler output.
305
num_reqs = min(num_tokens, self.max_num_reqs)
306
num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
@@ -306,13 +316,31 @@ class GPUModelRunner(LoRAModelRunnerMixin):
316
# Disable any use of KVConnector for dummy runs.
317
self.kv_connector.set_disabled(True)
318
319
+ # For non-first PP ranks, create dummy intermediate_tensors.
320
+ intermediate_tensors = None
321
+ if self.use_pp and not get_pp_group().is_first_rank:
322
+ intermediate_tensors = self.model.make_empty_intermediate_tensors(
323
+ batch_size=num_tokens,
324
+ dtype=self.model_config.dtype,
325
+ device=self.device,
326
+ )
327
+
328
# Execute the model.
329
self.execute_model(
311
- dummy_scheduler_output, dummy_run=True, skip_attn_for_dummy_run=skip_attn
330
+ dummy_scheduler_output,
331
+ intermediate_tensors=intermediate_tensors,
332
+ dummy_run=True,
333
+ skip_attn_for_dummy_run=skip_attn,
334
)
335
self.kv_connector.set_disabled(False)
336
+
337
+ # Non-last PP ranks don't produce output for sampling.
338
+ if self.use_pp and not get_pp_group().is_last_rank:
339
+ return None, None
340
+
341
assert self.execute_model_state is not None
342
hidden_states, input_batch, _ = self.execute_model_state
343
+ assert hidden_states is not None # Last PP rank always has hidden_states
344
sample_hidden_states = hidden_states[input_batch.logits_indices]
345
return hidden_states, sample_hidden_states
346
@@ -345,7 +373,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
373
hidden_states, sample_hidden_states = self._dummy_run(
374
self.max_num_tokens, skip_attn=True
375
)
348
- self._dummy_sampler_run(sample_hidden_states)
376
+ # 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:
378
+ assert sample_hidden_states is not None
379
+ 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
@@ -381,6 +412,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
412
)
413
return 0
414
415
+ # TODO (zhanqiu): support CUDA graph for PP.
416
+ if self.use_pp:
417
+ logger.warning_once(
418
+ "Skipping CUDA graph capture because pipeline parallel is "
419
+ "enabled. Pipeline parallel is currently eager-only.",
420
+ )
421
+ return 0
422
+
423
start_time = time.perf_counter()
424
gc.collect()
425
torch.cuda.empty_cache()
@@ -801,11 +840,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
840
def execute_model(
841
self,
842
scheduler_output: SchedulerOutput,
804
- intermediate_tensors: Any | None = None,
843
+ intermediate_tensors: IntermediateTensors | None = None,
844
dummy_run: bool = False,
845
skip_attn_for_dummy_run: bool = False,
807
- ) -> ModelRunnerOutput | None:
808
- assert intermediate_tensors is None
846
+ ) -> ModelRunnerOutput | IntermediateTensors | None:
847
if not dummy_run:
848
# Update the request states.
849
self.finish_requests(scheduler_output)
@@ -851,8 +889,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
889
)
890
self._set_active_loras(*lora_inputs)
891
854
- if self.supports_mm_inputs:
855
- # Execute the multimodal encoder.
892
+ # 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
+ ):
896
mm_embeds, is_mm_embed = self.get_mm_embeddings(
897
scheduler_output.scheduled_encoder_inputs, input_batch
898
)
@@ -894,6 +934,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
934
if self.uses_mrope:
935
assert input_batch.mrope_positions is not None
936
positions = input_batch.mrope_positions
937
+
938
with set_forward_context(
939
input_batch.attn_metadata,
940
self.vllm_config,
@@ -904,27 +945,71 @@ class GPUModelRunner(LoRAModelRunnerMixin):
945
slot_mapping=input_batch.slot_mappings,
946
):
947
self.kv_connector.pre_forward(scheduler_output)
907
- hidden_states = self.model(
908
- input_ids=input_batch.input_ids,
909
- positions=positions,
910
- inputs_embeds=input_batch.inputs_embeds,
911
- )
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
964
kv_connector_output = self.kv_connector.post_forward(scheduler_output)
914
- self.execute_model_state = hidden_states, input_batch, kv_connector_output
965
+
966
+ if self.use_pp and not get_pp_group().is_last_rank:
967
+ # Non-last PP rank: return IntermediateTensors for sending.
968
+ assert isinstance(hidden_states, IntermediateTensors)
969
+ hidden_states.kv_connector_output = kv_connector_output
970
+ self.execute_model_state = (None, input_batch, kv_connector_output)
971
+ return hidden_states
972
+
973
+ assert isinstance(hidden_states, torch.Tensor)
974
+ # Last rank (or no PP): hidden_states is a tensor for sampling.
975
+ self.execute_model_state = (hidden_states, input_batch, kv_connector_output)
976
return None
977
978
@torch.inference_mode()
979
def sample_tokens(
980
self, grammar_output: GrammarOutput | None
920
- ) -> AsyncOutput | ModelRunnerOutput:
981
+ ) -> AsyncOutput | ModelRunnerOutput | None:
982
assert self.execute_model_state is not None
983
hidden_states, input_batch, kv_connector_output = self.execute_model_state
984
self.execute_model_state = None # type: ignore
985
986
+ # Non-last PP rank: hidden_states is None because this rank produced
987
+ # IntermediateTensors instead of final hidden states. Receive the
988
+ # sampled tokens broadcast by the last rank and update local state.
989
+ if self.use_pp and not get_pp_group().is_last_rank:
990
+ assert self.pp_handler is not None
991
+ 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,
995
+ )
996
+ if received is not None:
997
+ sampled, num_sampled, num_rejected = received
998
+ self.postprocess(input_batch, sampled, num_sampled, num_rejected)
999
+ return None
1000
+
1001
+ # Last rank: sample tokens
1002
sampler_output, num_sampled, num_rejected = self.sample(
1003
hidden_states, input_batch, grammar_output
1004
)
1005
+
1006
+ # Broadcast to non-last PP ranks (handles spec decode multi-token).
1007
+ if self.use_pp:
1008
+ assert self.pp_handler is not None
1009
+ self.pp_handler.maybe_broadcast_sampled_tokens(
1010
+ sampler_output, num_sampled, num_rejected
1011
+ )
1012
+
1013
prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
1014
self.model.compute_logits,
1015
hidden_states,
vllm/v1/worker/gpu/pp_handler.py
new
+119
@@ -0,0 +1,119 @@
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 maybe_broadcast_sampled_tokens(
19
+ self,
20
+ sampler_output: SamplerOutput,
21
+ num_sampled: torch.Tensor,
22
+ num_rejected: torch.Tensor,
23
+ ) -> None:
24
+ """Broadcast sampled tokens from the last PP rank to all other ranks.
25
+
26
+ No-ops if this is not the last rank.
27
+
28
+ Broadcasts sampled_token_ids [num_reqs, max_sample_len], num_sampled
29
+ [num_reqs], and num_rejected [num_reqs] to support both regular decode
30
+ and speculative decoding.
31
+
32
+ Args:
33
+ sampler_output: SamplerOutput from sampling.
34
+ num_sampled: Number of accepted tokens per request.
35
+ num_rejected: Number of rejected tokens per request.
36
+ """
37
+ pp = get_pp_group()
38
+ if not pp.is_last_rank:
39
+ return
40
+
41
+ torch.distributed.broadcast(
42
+ sampler_output.sampled_token_ids.contiguous(),
43
+ src=pp.last_rank,
44
+ group=pp.device_group,
45
+ )
46
+ # NOTE: num_sampled/num_rejected are only needed
47
+ # for speculative decoding.
48
+ torch.distributed.broadcast(
49
+ num_sampled.contiguous(),
50
+ src=pp.last_rank,
51
+ group=pp.device_group,
52
+ )
53
+ torch.distributed.broadcast(
54
+ num_rejected.contiguous(),
55
+ src=pp.last_rank,
56
+ group=pp.device_group,
57
+ )
58
+
59
+ def maybe_receive_sampled_tokens(
60
+ self,
61
+ num_reqs: int,
62
+ device: torch.device,
63
+ max_sample_len: int = 1,
64
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None:
65
+ """Receive sampled tokens broadcast by the last PP rank.
66
+
67
+ Returns None if this is the last rank (which samples, not receives).
68
+
69
+ Args:
70
+ num_reqs: Number of requests in the batch.
71
+ device: Device to create tensors on.
72
+ max_sample_len: Maximum number of tokens sampled per request
73
+ (1 for regular decode, >1 for speculative decoding).
74
+
75
+ Returns:
76
+ None if called on last rank.
77
+ Otherwise, tuple of (sampled_tokens, num_sampled, num_rejected):
78
+ - sampled_tokens: shape [num_reqs, max_sample_len]
79
+ - num_sampled: shape [num_reqs]
80
+ - num_rejected: shape [num_reqs]
81
+ """
82
+ pp = get_pp_group()
83
+ if pp.is_last_rank:
84
+ return None
85
+
86
+ sampled_tokens = torch.empty(
87
+ num_reqs, max_sample_len, dtype=torch.int64, device=device
88
+ )
89
+ torch.distributed.broadcast(
90
+ sampled_tokens,
91
+ src=pp.last_rank,
92
+ group=pp.device_group,
93
+ )
94
+ # NOTE: num_sampled/num_rejected are only needed
95
+ # for speculative decoding.
96
+ num_sampled = torch.empty(num_reqs, dtype=torch.int32, device=device)
97
+ torch.distributed.broadcast(
98
+ num_sampled,
99
+ src=pp.last_rank,
100
+ group=pp.device_group,
101
+ )
102
+ num_rejected = torch.empty(num_reqs, dtype=torch.int32, device=device)
103
+ torch.distributed.broadcast(
104
+ num_rejected,
105
+ src=pp.last_rank,
106
+ group=pp.device_group,
107
+ )
108
+ 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()