[Fast Start] Support PP (#55477)
Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Xun Sun <UNIDY2002@outlook.com>
siyu committed
Sep 29, 2026 at 17:37 UTC
af5b4857e1353c01fd6bf41bc3cb9f84dc82dd89
5 files changed
+100
-52
tests/model_executor/model_loader/test_weight_cache.py
+26
-10
@@ -250,12 +250,12 @@ def test_daemon_places_tp_and_dp_ranks_on_local_gpus():
250
from vllm.model_executor.model_loader.weight_cache.daemon import plan_local_ranks
251
252
tp_second_node = _parallel(tensor_parallel_size=8, nnodes=2, node_rank=1)
253
- assert plan_local_ranks(tp_second_node) == [(i, 0, 4 + i) for i in range(4)]
253
+ assert plan_local_ranks(tp_second_node) == [(i, 0, 0, 4 + i) for i in range(4)]
254
255
dp_third_node = _parallel(
256
data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=8
257
)
258
- assert plan_local_ranks(dp_third_node) == [(i, 8 + i, 0) for i in range(4)]
258
+ assert plan_local_ranks(dp_third_node) == [(i, 8 + i, 0, 0) for i in range(4)]
259
260
dp_tp = _parallel(
261
tensor_parallel_size=2,
@@ -263,29 +263,43 @@ def test_daemon_places_tp_and_dp_ranks_on_local_gpus():
263
data_parallel_size_local=2,
264
data_parallel_rank=2,
265
)
266
- assert plan_local_ranks(dp_tp) == [(0, 2, 0), (1, 2, 1), (2, 3, 0), (3, 3, 1)]
266
+ assert plan_local_ranks(dp_tp) == [
267
+ (0, 2, 0, 0),
268
+ (1, 2, 0, 1),
269
+ (2, 3, 0, 0),
270
+ (3, 3, 0, 1),
271
+ ]
272
273
# DP + nnodes: node-local TP replicas
274
dp_tp_node0 = _parallel(tensor_parallel_size=8, data_parallel_size=2, nnodes=2)
270
- assert plan_local_ranks(dp_tp_node0) == [(i, 0, i) for i in range(8)]
275
+ assert plan_local_ranks(dp_tp_node0) == [(i, 0, 0, i) for i in range(8)]
276
dp_tp_node1 = _parallel(
277
tensor_parallel_size=8, data_parallel_size=2, nnodes=2, node_rank=1
278
)
274
- assert plan_local_ranks(dp_tp_node1) == [(i, 1, i) for i in range(8)]
279
+ assert plan_local_ranks(dp_tp_node1) == [(i, 1, 0, i) for i in range(8)]
280
281
# DP + nnodes: TP group spans nodes (TP8 x DP2 on 4 nodes)
282
tp_span_node1 = _parallel(
283
tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=1
284
)
280
- assert plan_local_ranks(tp_span_node1) == [(i, 0, 4 + i) for i in range(4)]
285
+ assert plan_local_ranks(tp_span_node1) == [(i, 0, 0, 4 + i) for i in range(4)]
286
tp_span_node2 = _parallel(
287
tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=2
288
)
284
- assert plan_local_ranks(tp_span_node2) == [(i, 1, i) for i in range(4)]
289
+ assert plan_local_ranks(tp_span_node2) == [(i, 1, 0, i) for i in range(4)]
290
291
# DP + nnodes: several DP replicas per node (TP4 x DP4 on 2 nodes)
292
dp_multi_node0 = _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2)
288
- assert plan_local_ranks(dp_multi_node0) == [(i, i // 4, i % 4) for i in range(8)]
293
+ assert plan_local_ranks(dp_multi_node0) == [(i, i // 4, 0, i % 4) for i in range(8)]
294
+
295
+ # PP adds one daemon placement per stage while preserving TP rank order.
296
+ pp_tp = _parallel(tensor_parallel_size=2, pipeline_parallel_size=2)
297
+ assert plan_local_ranks(pp_tp) == [
298
+ (0, 0, 0, 0),
299
+ (1, 0, 0, 1),
300
+ (2, 0, 1, 0),
301
+ (3, 0, 1, 1),
302
+ ]
303
304
305
def test_daemon_rejects_unmappable_parallelism():
@@ -302,8 +316,7 @@ def test_daemon_rejects_unmappable_parallelism():
316
_reject_unsupported_parallelism(
317
_parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2)
318
)
305
- with pytest.raises(ValueError, match="pipeline"):
306
- _reject_unsupported_parallelism(_parallel(pipeline_parallel_size=2))
319
+ _reject_unsupported_parallelism(_parallel(pipeline_parallel_size=2))
320
with pytest.raises(ValueError, match="evenly divide"):
321
_reject_unsupported_parallelism(_parallel(tensor_parallel_size=3, nnodes=2))
322
with pytest.raises(ValueError, match="evenly divide"):
@@ -335,6 +348,9 @@ def test_weight_cache_key_distinguishes_dp_ranks():
348
vllm_version="v",
349
dp_size=16,
350
dp_rank=3,
351
+ pp_size=2,
352
+ pp_rank=1,
353
)
354
assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"]
355
assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"]
356
+ assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"]
vllm/entrypoints/cli/preload.py
+16
-6
@@ -170,6 +170,7 @@ class PreloadSubcommand(CLISubcommand):
170
parallel_config = vllm_config.parallel_config
171
_reject_unsupported_parallelism(parallel_config)
172
tp_size = parallel_config.tensor_parallel_size
173
+ pp_size = parallel_config.pipeline_parallel_size
174
dp_size = parallel_config.data_parallel_size
175
176
placements = plan_local_ranks(parallel_config)
@@ -230,16 +231,21 @@ class PreloadSubcommand(CLISubcommand):
231
232
ctx = multiprocessing.get_context("spawn")
233
ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue()
233
- # Local index == device index; global rank enumerates DP then TP.
234
+ # Local index == device index; global rank enumerates DP, PP, then TP.
235
expected_ready = {
235
- (format_daemon_role(is_draft), dp_rank * tp_size + tp_rank)
236
- for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements)
236
+ (
237
+ format_daemon_role(is_draft),
238
+ dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank,
239
+ )
240
+ for (is_draft, _, _), (_, dp_rank, pp_rank, tp_rank) in product(
241
+ groups, placements
242
+ )
243
}
244
procs = [
245
ctx.Process(
246
target=_run_daemon,
247
args=(
242
- tp_rank,
248
+ dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank,
249
local_rank,
250
config,
251
init_method,
@@ -247,13 +253,17 @@ class PreloadSubcommand(CLISubcommand):
253
ready_queue,
254
is_draft,
255
dp_rank,
256
+ pp_rank,
257
+ ),
258
+ name=(
259
+ f"vllm-weight-cache-{format_daemon_role(is_draft)}-"
260
+ f"{(dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank)}"
261
),
251
- name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-"
252
- f"{dp_rank * tp_size + tp_rank}",
262
)
263
for (is_draft, config, init_method), (
264
local_rank,
265
dp_rank,
266
+ pp_rank,
267
tp_rank,
268
) in product(groups, placements)
269
]
vllm/model_executor/model_loader/weight_cache/daemon.py
+48
-36
@@ -17,8 +17,7 @@ Engines then load from the daemons with:
17
vllm serve /path/to/model --tensor-parallel-size 4 \\
18
--load-format ipc_cache
19
20
-Tensor, expert and data parallelism are supported; pipeline parallelism is
21
-rejected at launch.
20
+Tensor, expert, data and pipeline parallelism are supported.
21
22
For multi-node tensor parallelism, run one launcher per node with a shared
23
rendezvous so the global TP group forms across nodes (CUDA IPC handles are
@@ -40,7 +39,8 @@ node-local, so each node serves only its local GPUs' shards). Reuse the same
39
The global TP rank of local GPU ``i`` on node ``r`` is
40
``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank
41
assignment, so each engine worker maps its shard from the daemon on its own
43
-node.
42
+node. With pipeline parallelism each node additionally runs one daemon per PP
43
+stage; the global rank is ``pp_rank * tp_size + tp_rank``.
44
45
For data parallelism (e.g. a TP1 x DP16 x EP decode fleet) run one launcher per
46
node with the engine's DP placement flags. Local GPU ``i`` serves DP rank
@@ -168,27 +168,30 @@ def export_entries(
168
169
170
class WeightCacheDaemon:
171
- """Per-GPU process that loads one TP shard and serves CUDA IPC handles."""
171
+ """Per-GPU process that loads one TP/PP shard and serves CUDA IPC handles."""
172
173
def __init__(
174
self,
175
vllm_config: VllmConfig,
176
- tp_rank: int,
176
+ global_rank: int,
177
local_rank: int,
178
distributed_init_method: str,
179
socket_dir: str | None = None,
180
is_draft: bool = False,
181
dp_rank: int = 0,
182
+ pp_rank: int = 0,
183
):
184
parallel_config = vllm_config.parallel_config
184
- if parallel_config.data_parallel_size > 1:
185
+ self.tp_size = parallel_config.tensor_parallel_size
186
+ self.pp_size = parallel_config.pipeline_parallel_size
187
+ self.dp_size = parallel_config.data_parallel_size
188
+ if self.dp_size > 1:
189
# The engine sets these on each DP rank's workers; the MoE/EP
190
# layers read them to place their expert shards.
187
- tp_size = parallel_config.tensor_parallel_size
191
if parallel_config.nnodes > 1:
189
- world_size = parallel_config.data_parallel_size * tp_size
192
+ world_size = self.dp_size * self.tp_size
193
local_world = world_size // parallel_config.nnodes
191
- start_dp_rank = parallel_config.node_rank * local_world // tp_size
194
+ start_dp_rank = parallel_config.node_rank * local_world // self.tp_size
195
else:
196
start_dp_rank = parallel_config.data_parallel_rank
197
parallel_config = replace(
@@ -206,12 +209,12 @@ class WeightCacheDaemon:
209
self.model_config = spec.draft_model_config
210
else:
211
self.model_config = vllm_config.model_config
209
- self.tp_rank = tp_rank
212
+ self.global_rank = global_rank
213
+ rank_in_dp = global_rank - dp_rank * self.pp_size * self.tp_size
214
+ self.tp_rank = rank_in_dp % self.tp_size
215
+ self.pp_rank = rank_in_dp // self.tp_size
216
self.dp_rank = dp_rank
211
- self.tp_size = parallel_config.tensor_parallel_size
212
- self.dp_size = parallel_config.data_parallel_size
213
- self.global_rank = dp_rank * self.tp_size + tp_rank
214
- self.world_size = self.dp_size * self.tp_size
217
+ self.world_size = self.dp_size * self.pp_size * self.tp_size
218
self.local_rank = local_rank
219
self.distributed_init_method = distributed_init_method
220
self.socket_dir = socket_dir
@@ -223,7 +226,9 @@ class WeightCacheDaemon:
226
self.cache_config = WeightCacheKey.from_model_config(
227
self.model_config,
228
tp_size=self.tp_size,
226
- tp_rank=tp_rank,
229
+ tp_rank=self.tp_rank,
230
+ pp_size=self.pp_size,
231
+ pp_rank=self.pp_rank,
232
dp_size=self.dp_size,
233
dp_rank=dp_rank,
234
is_draft=is_draft,
@@ -242,7 +247,7 @@ class WeightCacheDaemon:
247
backend=current_platform.dist_backend,
248
)
249
with set_current_vllm_config(self.vllm_config):
245
- ensure_model_parallel_initialized(self.tp_size, 1)
250
+ ensure_model_parallel_initialized(self.tp_size, self.pp_size)
251
self.model = self.get_model()
252
logger.info(
253
"Weight cache %s daemon rank %d loaded model",
@@ -414,7 +419,7 @@ class WeightCacheDaemon:
419
420
421
def _run_daemon(
417
- tp_rank: int,
422
+ global_rank: int,
423
local_rank: int,
424
vllm_config: VllmConfig,
425
distributed_init_method: str,
@@ -422,15 +427,17 @@ def _run_daemon(
427
ready_queue: "multiprocessing.Queue[tuple[str, int]]",
428
is_draft: bool = False,
429
dp_rank: int = 0,
430
+ pp_rank: int = 0,
431
) -> None:
432
daemon = WeightCacheDaemon(
433
vllm_config,
428
- tp_rank,
434
+ global_rank,
435
local_rank,
436
distributed_init_method,
437
socket_dir,
438
is_draft,
439
dp_rank,
440
+ pp_rank,
441
)
442
daemon.load_model()
443
daemon.serve_forever(
@@ -438,26 +445,34 @@ def _run_daemon(
445
)
446
447
441
-def plan_local_ranks(parallel_config: ParallelConfig) -> list[tuple[int, int, int]]:
442
- """``(local_rank, dp_rank, tp_rank)`` for every GPU this launcher serves.
448
+def plan_local_ranks(
449
+ parallel_config: ParallelConfig,
450
+) -> list[tuple[int, int, int, int]]:
451
+ """``(local_rank, dp_rank, pp_rank, tp_rank)`` for every local daemon.
452
444
- Global ranks enumerate DP then TP: ``global = dp_rank * tp_size + tp_rank``.
453
+ Global ranks enumerate DP, PP, then TP: ``global = dp_rank * pp_size *
454
+ tp_size + pp_rank * tp_size + tp_rank``.
455
With ``--nnodes`` the engine hands each node a contiguous block of global
456
ranks, so node ``r`` serves ``node_rank * local + i``; without it a
457
launcher's block starts at ``--data-parallel-start-rank * tp_size``.
458
"""
459
tp_size = parallel_config.tensor_parallel_size
450
- world_size = parallel_config.data_parallel_size * tp_size
460
+ pp_size = parallel_config.pipeline_parallel_size
461
+ world_size = parallel_config.data_parallel_size * pp_size * tp_size
462
if parallel_config.nnodes > 1:
463
local_world_size = world_size // parallel_config.nnodes
464
base = parallel_config.node_rank * local_world_size
465
else:
455
- local_world_size = parallel_config.data_parallel_size_local * tp_size
456
- base = parallel_config.data_parallel_rank * tp_size
457
- return [
458
- (i, (base + i) // tp_size, (base + i) % tp_size)
459
- for i in range(local_world_size)
460
- ]
466
+ local_world_size = parallel_config.data_parallel_size_local * pp_size * tp_size
467
+ base = parallel_config.data_parallel_rank * pp_size * tp_size
468
+ placements = []
469
+ for local_rank in range(local_world_size):
470
+ global_rank = base + local_rank
471
+ rank_in_dp = global_rank % (pp_size * tp_size)
472
+ dp_rank = global_rank // (pp_size * tp_size)
473
+ pp_rank, tp_rank = divmod(rank_in_dp, tp_size)
474
+ placements.append((local_rank, dp_rank, pp_rank, tp_rank))
475
+ return placements
476
477
478
def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None:
@@ -469,21 +484,18 @@ def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None:
484
485
486
def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:
472
- """Reject pipeline parallelism and placements the daemon cannot map."""
473
- if parallel_config.pipeline_parallel_size > 1:
474
- raise ValueError(
475
- "The weight cache daemon only supports tensor, expert and data "
476
- "parallelism; pipeline parallelism is not supported"
477
- )
487
+ """Reject placements the daemon cannot map."""
488
dp_size = parallel_config.data_parallel_size
489
tp_size = parallel_config.tensor_parallel_size
490
+ pp_size = parallel_config.pipeline_parallel_size
491
if parallel_config.nnodes > 1:
481
- world_size = dp_size * tp_size
492
+ world_size = dp_size * pp_size * tp_size
493
if world_size % parallel_config.nnodes != 0:
494
raise ValueError(
495
f"--nnodes ({parallel_config.nnodes}) must evenly divide the "
496
f"daemon world size ({world_size} = data-parallel-size "
486
- f"{dp_size} x tensor-parallel-size {tp_size})"
497
+ f"{dp_size} x pipeline-parallel-size {pp_size} x "
498
+ f"tensor-parallel-size {tp_size})"
499
)
500
return
501
if dp_size == 1:
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+4
@@ -15,6 +15,7 @@ from vllm.config import ModelConfig, VllmConfig
15
from vllm.config.load import LoadConfig
16
from vllm.distributed import (
17
get_dp_group,
18
+ get_pp_group,
19
get_tensor_model_parallel_rank,
20
get_tensor_model_parallel_world_size,
21
)
@@ -289,10 +290,13 @@ class IpcModelLoader(BaseModelLoader):
290
291
def _fetch_entries(self, model_config: ModelConfig) -> WeightCacheState:
292
dp_group = get_dp_group()
293
+ pp_group = get_pp_group()
294
cache_config = WeightCacheKey.from_model_config(
295
model_config,
296
tp_size=get_tensor_model_parallel_world_size(),
297
tp_rank=get_tensor_model_parallel_rank(),
298
+ pp_size=pp_group.world_size,
299
+ pp_rank=pp_group.rank_in_group,
300
dp_size=dp_group.world_size,
301
dp_rank=dp_group.rank_in_group,
302
is_draft=self.is_draft,
vllm/model_executor/model_loader/weight_cache/protocol.py
+6
@@ -289,6 +289,8 @@ class WeightCacheKey:
289
quant_config_hash: str
290
revision: str | None
291
vllm_version: str
292
+ pp_size: int = 1
293
+ pp_rank: int = 0
294
is_draft: bool = False
295
"""Daemon group the weights come from; False is the target model."""
296
dp_size: int = 1
@@ -301,6 +303,8 @@ class WeightCacheKey:
303
tp_size: int,
304
tp_rank: int,
305
*,
306
+ pp_size: int = 1,
307
+ pp_rank: int = 0,
308
is_draft: bool = False,
309
dp_size: int = 1,
310
dp_rank: int = 0,
@@ -325,6 +329,8 @@ class WeightCacheKey:
329
model_arch=arch,
330
tp_size=tp_size,
331
tp_rank=tp_rank,
332
+ pp_size=pp_size,
333
+ pp_rank=pp_rank,
334
dtype=str(model_config.dtype),
335
quantization=model_config.quantization,
336
quant_config_hash=_hash_quant_config(quant_config),