[Refactor][PCP] Make PCPManager construction extensible (#50066)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Qiu committed
Aug 6, 2026 at 07:05 UTC
811622c410c2f1baf2fa7056ca19753264b95815
2 files changed
+9
-2
vllm/v1/worker/gpu/model_runner.py
+6
@@ -510,6 +510,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
510
self.supports_mm_inputs,
511
self.req_states,
512
self.block_tables,
513
+ cls=self.pcp_manager_cls,
514
)
515
initialize_mamba_ssu_backend(
516
self.vllm_config.mamba_config, self.kv_cache_config
@@ -1751,6 +1752,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1752
1753
########### EPLB methods end ###########
1754
1755
+ # Out-of-tree hardware runners can select a PCP manager class.
1756
+ @property
1757
+ def pcp_manager_cls(self) -> type[pcp.PCPManager]:
1758
+ return pcp.PCPManager
1759
+
1760
1761
class ExecuteModelState(NamedTuple):
1762
input_batch: InputBatch
vllm/v1/worker/gpu/pcp_manager.py
+3
-2
@@ -654,19 +654,20 @@ def maybe_build_pcp_manager(
654
supports_mm_inputs: bool,
655
req_states: RequestState,
656
block_tables: BlockTables,
657
+ cls: type[PCPManager] = PCPManager,
658
) -> PCPManager | None:
659
parallel_config = vllm_config.parallel_config
660
pcp_size = parallel_config.prefill_context_parallel_size
661
if pcp_size <= 1:
662
return None
663
663
- PCPManager.validate_config(vllm_config, supports_mm_inputs)
664
+ cls.validate_config(vllm_config, supports_mm_inputs)
665
666
pcp_rank = get_pcp_group().rank_in_group
667
dcp_size = parallel_config.decode_context_parallel_size
668
dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0
669
669
- return PCPManager(
670
+ return cls(
671
pcp_world_size=pcp_size,
672
pcp_rank=pcp_rank,
673
device=device,