[Feat] Support PP with PCP in GPU Model Runner V2 (#59139)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Qiu Chunshuo committed
Sep 30, 2026 at 07:21 UTC
98b45c1a94db509e75fa9ef268cf5d2ae71875b7
4 files changed
+46
-3
tests/v1/worker/test_gpu_model_runner_v2.py
+38
-1
@@ -3,6 +3,7 @@
3
4
import contextlib
5
from types import SimpleNamespace
6
+from unittest.mock import Mock
7
8
import pytest
9
import torch
@@ -18,7 +19,43 @@ from vllm.v1.kv_cache_interface import (
19
UniformTypeKVCacheSpecs,
20
)
21
from vllm.v1.worker.gpu.block_table import BlockTables
21
-from vllm.v1.worker.gpu.model_runner import GPUModelRunner
22
+from vllm.v1.worker.gpu.model_runner import ExecuteModelState, GPUModelRunner
23
+
24
+
25
+def test_non_last_pp_rank_uses_global_batch_for_sample_feedback():
26
+ runner = GPUModelRunner.__new__(GPUModelRunner)
27
+ runner.is_last_pp_rank = False
28
+ local_batch = object()
29
+ global_batch = SimpleNamespace(idx_mapping=object())
30
+ runner.pcp_manager = SimpleNamespace(
31
+ global_batch=global_batch,
32
+ restore_for_sampling=Mock(),
33
+ )
34
+ runner.pp_handler = SimpleNamespace(receive=Mock(return_value=False))
35
+ runner.postprocess_num_computed_tokens = Mock()
36
+ runner.model_state = SimpleNamespace(postprocess_state=Mock())
37
+ runner.kv_connector = SimpleNamespace(post_forward=Mock(return_value=None))
38
+ runner.eplb = SimpleNamespace(step=Mock())
39
+ runner.execute_model_state = ExecuteModelState(
40
+ input_batch=local_batch,
41
+ attn_metadata=None,
42
+ slot_mappings_by_layer=None,
43
+ hidden_states=None,
44
+ aux_hidden_states=None,
45
+ dp_sync=None,
46
+ finished_req_ids=set(),
47
+ ec_connector_output=None,
48
+ cudagraph_stats=None,
49
+ )
50
+
51
+ runner.sample_tokens(None)
52
+
53
+ runner.pp_handler.receive.assert_called_once_with(global_batch)
54
+ runner.postprocess_num_computed_tokens.assert_called_once_with(global_batch)
55
+ runner.model_state.postprocess_state.assert_called_once_with(
56
+ global_batch.idx_mapping, 0
57
+ )
58
+ runner.pcp_manager.restore_for_sampling.assert_not_called()
59
60
61
def test_qsa_circular_group_uses_custom_slot_mapping(monkeypatch):
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+1
@@ -79,6 +79,7 @@ def _make_runner(**overrides: Any) -> Any:
79
runner.use_aux_hidden_state_outputs = False
80
runner.speculative_config = None
81
runner.speculator = None
82
+ runner.pcp_manager = None
83
runner.num_speculative_steps = 0
84
runner.encoder_cache = None
85
runner.is_pooling_model = False
vllm/v1/worker/gpu/model_runner.py
+2
@@ -2064,6 +2064,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
2064
# Non-last PP rank: hidden_states is None because this rank produced
2065
# IntermediateTensors instead of final hidden states. Receive the
2066
# sampled tokens broadcast from the last rank and update local state.
2067
+ if self.pcp_manager is not None:
2068
+ input_batch = self.pcp_manager.global_batch
2069
assert self.pp_handler is not None
2070
all_decode_next = self.pp_handler.receive(input_batch)
2071
# Optimistically update num_computed_tokens for entire batch here.
vllm/v1/worker/gpu/pcp_manager.py
+5
-2
@@ -131,8 +131,6 @@ class PCPManager:
131
132
if not model_config.use_mla:
133
raise NotImplementedError("MRV2 PCP currently supports MLA models only.")
134
- if parallel_config.pipeline_parallel_size > 1:
135
- raise NotImplementedError("MRV2 PCP does not support PP yet.")
134
if model_config.is_encoder_decoder:
135
raise NotImplementedError(
136
"MRV2 PCP does not support encoder-decoder models yet."
@@ -406,6 +404,11 @@ class PCPManager:
404
assert self._input_buffers is not None
405
return self._input_buffers
406
407
+ @property
408
+ def global_batch(self) -> InputBatch:
409
+ assert self._global_batch is not None
410
+ return self._global_batch
411
+
412
def partition_batch(
413
self, input_batch: InputBatch, batch_desc: "BatchExecutionDescriptor"
414
) -> InputBatch: