From c82aa7d441cd5209c65018b3b82468fdad11166e Mon Sep 17 00:00:00 2001 From: QiuChunshuo Date: Sun, 23 Aug 2026 12:27:30 +0800 Subject: [PATCH 1/9] fix(pcp): stabilize piecewise cudagraph replay Signed-off-by: QiuChunshuo --- tests/v1/worker/test_gpu_pcp_manager.py | 93 +++++++++++++++++++++++++ vllm/v1/worker/gpu/cudagraph_utils.py | 22 ++++-- vllm/v1/worker/gpu/model_runner.py | 5 ++ vllm/v1/worker/gpu/pcp_manager.py | 28 +++++++- 4 files changed, 142 insertions(+), 6 deletions(-) create mode 100644 tests/v1/worker/test_gpu_pcp_manager.py diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py new file mode 100644 index 000000000000..b77eed8261e7 --- /dev/null +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np +import pytest +import torch + +from vllm.config import CUDAGraphMode +from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module +from vllm.v1.worker.gpu.cudagraph_utils import ( + BatchExecutionDescriptor, + CudaGraphManager, +) +from vllm.v1.worker.gpu.pcp_manager import PCPManager + + +def _copy_to_cpu(value, out=None, device=None): + tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value + if out is not None: + return out.copy_(tensor) + return tensor + + +def test_pcp_piecewise_dispatch_accepts_only_uniform_decode(): + manager = CudaGraphManager.__new__(CudaGraphManager) + manager.decode_query_len = 1 + manager.pcp_size = 2 + manager._graphs_captured = True + manager._lora_dispatch_map = {} + manager._max_lora_case = 0 + piecewise = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.PIECEWISE, + num_tokens=40, + num_reqs=None, + ) + manager._candidates = {(40, 0): [piecewise]} + + mixed = manager.dispatch(20, 40, None, num_active_loras=0) + decode = manager.dispatch(40, 40, 1, num_active_loras=0) + + assert mixed.cg_mode == CUDAGraphMode.NONE + assert decode == piecewise + + +def test_replicated_decode_piecewise_graph_padding(monkeypatch): + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + dcp_world_size=1, + ) + monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + + segments_by_rank, per_rank_num_tokens = manager._build_batch_layout( + num_scheduled_tokens=np.ones(3, dtype=np.int32), + num_computed_tokens=np.full(3, 16, dtype=np.int32), + is_prefilling=np.zeros(3, dtype=np.bool_), + query_start_loc_np=np.arange(4, dtype=np.int32), + padded_num_tokens=4, + ) + + assert per_rank_num_tokens == [3, 3] + request_indices = [ + [segment.global_batch_req_idx for segment in rank] for rank in segments_by_rank + ] + assert request_indices == [[0, 1, 2], [0, 1, 2]] + assert torch.equal(manager._hidden_restore_idx, torch.tensor([0, 1, 2])) + assert torch.equal( + manager._padded_gather_idx, + torch.tensor([0, 1, 2, 0, 0, 1, 2, 0]), + ) + assert torch.equal( + manager._gathered_kv_write_mask, + torch.tensor([True, True, True, False, False, False, False, False]), + ) + + +def test_replicated_decode_rejects_undersized_graph_padding(): + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + dcp_world_size=1, + ) + + with pytest.raises(RuntimeError, match="graph token capacity"): + manager._build_batch_layout( + num_scheduled_tokens=np.ones(3, dtype=np.int32), + num_computed_tokens=np.full(3, 16, dtype=np.int32), + is_prefilling=np.zeros(3, dtype=np.bool_), + query_start_loc_np=np.arange(4, dtype=np.int32), + padded_num_tokens=2, + ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 8e24a0ee0259..b9090bda6710 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -120,6 +120,7 @@ def __init__( self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size + self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size self.is_first_pp_rank = get_pp_group().is_first_rank self.is_last_pp_rank = get_pp_group().is_last_rank self.lora_capture_cases = lora_capture_cases or [0] @@ -393,6 +394,15 @@ def dispatch( key = (num_tokens, effective_loras) if self._graphs_captured and num_tokens > 0 and key in self._candidates: for desc in self._candidates[key]: + # PCP mixed batches have rank-local layouts that can change with + # the prefill split. Uniform decode keeps a stable replicated + # layout and is safe for PIECEWISE replay. + if ( + self.pcp_size > 1 + and desc.cg_mode == CUDAGraphMode.PIECEWISE + and uniform_token_count != self.decode_query_len + ): + continue if _is_compatible( desc, num_reqs, @@ -472,6 +482,7 @@ def capture( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, + pcp_input_buffers: InputBuffers | None = None, has_lora: bool = False, use_aux_hidden_state_outputs: bool = False, lora_capture_hook: Callable[[int, int, int], None] | None = None, @@ -488,6 +499,9 @@ def create_forward_fn( ) -> Callable[[CUDAGraphMode], None]: num_tokens = desc.num_tokens num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs) + capture_input_buffers = ( + pcp_input_buffers if pcp_input_buffers is not None else input_buffers + ) # Set LoRA state before capture so kernels see correct adapters. if lora_capture_hook is not None: @@ -500,8 +514,8 @@ def create_forward_fn( ) model_inputs = { - "input_ids": input_buffers.input_ids[:num_tokens], - "positions": input_buffers.positions[:num_tokens], + "input_ids": capture_input_buffers.input_ids[:num_tokens], + "positions": capture_input_buffers.positions[:num_tokens], **model_state.prepare_dummy_inputs(num_reqs, num_tokens), } if not self.is_first_pp_rank: @@ -524,7 +538,7 @@ def create_forward_fn( ) # Capture with dummy rows marked as padding. - input_buffers.is_padding.fill_(True) + capture_input_buffers.is_padding.fill_(True) def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None @@ -542,7 +556,7 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None: num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, - is_padding=input_buffers.is_padding[:num_tokens], + is_padding=capture_input_buffers.is_padding[:num_tokens], ): if cg_mode == CUDAGraphMode.PIECEWISE: # PIECEWISE graph (compiled PW or breakable, chosen inside diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3861bb52649d..5e32624bb769 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -885,6 +885,11 @@ def capture_model(self) -> int: self.block_tables, self.attn_groups, self.kv_cache_config, + pcp_input_buffers=( + self.pcp_manager.input_buffers + if self.pcp_manager is not None + else None + ), has_lora=self.lora_config is not None, use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, lora_capture_hook=create_lora_capture_hook(self.lora_config, self), diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index d80e246f17c9..9b221aad4467 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -160,6 +160,11 @@ def validate_config( if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(): raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + @property + def input_buffers(self) -> InputBuffers: + assert self._input_buffers is not None + return self._input_buffers + @staticmethod def _reorder_segments( segments: list[RankSegment], @@ -255,6 +260,7 @@ def _build_batch_layout( num_computed_tokens: np.ndarray, is_prefilling: np.ndarray, query_start_loc_np: np.ndarray, + padded_num_tokens: int | None = None, ) -> tuple[list[list[RankSegment]], list[int]]: segments_by_rank = [] per_rank_num_tokens = [] @@ -279,7 +285,14 @@ def _build_batch_layout( # Therefore global = gathered[hidden_restore_idx] and # padded_gathered = global[padded_gather_idx]. hidden_restore_idx = np.empty(int(query_start_loc_np[-1]), dtype=np.int64) - padded_num_tokens = max(per_rank_num_tokens) + required_num_tokens = max(per_rank_num_tokens) + if padded_num_tokens is None: + padded_num_tokens = required_num_tokens + elif padded_num_tokens < required_num_tokens: + raise RuntimeError( + "PCP graph token capacity is smaller than the rank-local batch: " + f"{padded_num_tokens} < {required_num_tokens}." + ) num_expanded_tokens = padded_num_tokens * self.pcp_world_size padded_gather_idx = np.zeros(num_expanded_tokens, dtype=np.int64) gathered_kv_write_mask = np.zeros(num_expanded_tokens, dtype=np.bool_) @@ -327,6 +340,12 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: global_batch = input_batch self._global_batch = global_batch + num_tokens_after_padding = ( + global_batch.num_tokens_after_padding + if global_batch.num_tokens_after_padding > global_batch.num_tokens + else None + ) + num_scheduled_tokens = global_batch.num_scheduled_tokens num_computed_tokens = global_batch.num_computed_tokens_np is_prefilling = global_batch.is_prefilling_np @@ -336,6 +355,7 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: num_computed_tokens, is_prefilling, global_batch.query_start_loc_np, + padded_num_tokens=num_tokens_after_padding, ) local_segments = segments_by_rank[self.pcp_rank] @@ -384,7 +404,11 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: ] num_local_tokens = int(local_num_scheduled_tokens.sum()) - num_local_tokens_padded = max(per_rank_num_tokens) + num_local_tokens_padded = ( + max(per_rank_num_tokens) + if num_tokens_after_padding is None + else num_tokens_after_padding + ) fresh_prefills = int( np.count_nonzero(is_prefilling & (num_computed_tokens == 0)) ) From ce6229207d34a4be7b4dd3f9662e979338c4e61c Mon Sep 17 00:00:00 2001 From: QiuChunshuo Date: Sun, 23 Aug 2026 15:26:04 +0800 Subject: [PATCH 2/9] refactor(pcp): minimize piecewise graph fix Signed-off-by: QiuChunshuo --- tests/v1/worker/test_gpu_pcp_manager.py | 45 ------------------------- vllm/v1/worker/gpu/cudagraph_utils.py | 22 +++--------- vllm/v1/worker/gpu/model_runner.py | 5 --- vllm/v1/worker/gpu/pcp_manager.py | 13 +------ 4 files changed, 5 insertions(+), 80 deletions(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index b77eed8261e7..4f9fc7232971 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -2,15 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import numpy as np -import pytest import torch -from vllm.config import CUDAGraphMode from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module -from vllm.v1.worker.gpu.cudagraph_utils import ( - BatchExecutionDescriptor, - CudaGraphManager, -) from vllm.v1.worker.gpu.pcp_manager import PCPManager @@ -21,27 +15,6 @@ def _copy_to_cpu(value, out=None, device=None): return tensor -def test_pcp_piecewise_dispatch_accepts_only_uniform_decode(): - manager = CudaGraphManager.__new__(CudaGraphManager) - manager.decode_query_len = 1 - manager.pcp_size = 2 - manager._graphs_captured = True - manager._lora_dispatch_map = {} - manager._max_lora_case = 0 - piecewise = BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.PIECEWISE, - num_tokens=40, - num_reqs=None, - ) - manager._candidates = {(40, 0): [piecewise]} - - mixed = manager.dispatch(20, 40, None, num_active_loras=0) - decode = manager.dispatch(40, 40, 1, num_active_loras=0) - - assert mixed.cg_mode == CUDAGraphMode.NONE - assert decode == piecewise - - def test_replicated_decode_piecewise_graph_padding(monkeypatch): manager = PCPManager( pcp_world_size=2, @@ -73,21 +46,3 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch): manager._gathered_kv_write_mask, torch.tensor([True, True, True, False, False, False, False, False]), ) - - -def test_replicated_decode_rejects_undersized_graph_padding(): - manager = PCPManager( - pcp_world_size=2, - pcp_rank=0, - device=torch.device("cpu"), - dcp_world_size=1, - ) - - with pytest.raises(RuntimeError, match="graph token capacity"): - manager._build_batch_layout( - num_scheduled_tokens=np.ones(3, dtype=np.int32), - num_computed_tokens=np.full(3, 16, dtype=np.int32), - is_prefilling=np.zeros(3, dtype=np.bool_), - query_start_loc_np=np.arange(4, dtype=np.int32), - padded_num_tokens=2, - ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index b9090bda6710..8e24a0ee0259 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -120,7 +120,6 @@ def __init__( self.dp_size = vllm_config.parallel_config.data_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size - self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size self.is_first_pp_rank = get_pp_group().is_first_rank self.is_last_pp_rank = get_pp_group().is_last_rank self.lora_capture_cases = lora_capture_cases or [0] @@ -394,15 +393,6 @@ def dispatch( key = (num_tokens, effective_loras) if self._graphs_captured and num_tokens > 0 and key in self._candidates: for desc in self._candidates[key]: - # PCP mixed batches have rank-local layouts that can change with - # the prefill split. Uniform decode keeps a stable replicated - # layout and is safe for PIECEWISE replay. - if ( - self.pcp_size > 1 - and desc.cg_mode == CUDAGraphMode.PIECEWISE - and uniform_token_count != self.decode_query_len - ): - continue if _is_compatible( desc, num_reqs, @@ -482,7 +472,6 @@ def capture( block_tables: BlockTables, attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, - pcp_input_buffers: InputBuffers | None = None, has_lora: bool = False, use_aux_hidden_state_outputs: bool = False, lora_capture_hook: Callable[[int, int, int], None] | None = None, @@ -499,9 +488,6 @@ def create_forward_fn( ) -> Callable[[CUDAGraphMode], None]: num_tokens = desc.num_tokens num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs) - capture_input_buffers = ( - pcp_input_buffers if pcp_input_buffers is not None else input_buffers - ) # Set LoRA state before capture so kernels see correct adapters. if lora_capture_hook is not None: @@ -514,8 +500,8 @@ def create_forward_fn( ) model_inputs = { - "input_ids": capture_input_buffers.input_ids[:num_tokens], - "positions": capture_input_buffers.positions[:num_tokens], + "input_ids": input_buffers.input_ids[:num_tokens], + "positions": input_buffers.positions[:num_tokens], **model_state.prepare_dummy_inputs(num_reqs, num_tokens), } if not self.is_first_pp_rank: @@ -538,7 +524,7 @@ def create_forward_fn( ) # Capture with dummy rows marked as padding. - capture_input_buffers.is_padding.fill_(True) + input_buffers.is_padding.fill_(True) def forward_fn(cg_mode: CUDAGraphMode) -> None: batch_descriptor = None @@ -556,7 +542,7 @@ def forward_fn(cg_mode: CUDAGraphMode) -> None: num_tokens_across_dp=num_tokens_across_dp, slot_mapping=slot_mappings, batch_descriptor=batch_descriptor, - is_padding=capture_input_buffers.is_padding[:num_tokens], + is_padding=input_buffers.is_padding[:num_tokens], ): if cg_mode == CUDAGraphMode.PIECEWISE: # PIECEWISE graph (compiled PW or breakable, chosen inside diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 5e32624bb769..3861bb52649d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -885,11 +885,6 @@ def capture_model(self) -> int: self.block_tables, self.attn_groups, self.kv_cache_config, - pcp_input_buffers=( - self.pcp_manager.input_buffers - if self.pcp_manager is not None - else None - ), has_lora=self.lora_config is not None, use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs, lora_capture_hook=create_lora_capture_hook(self.lora_config, self), diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 9b221aad4467..25bf7d382fa9 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -160,11 +160,6 @@ def validate_config( if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs(): raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") - @property - def input_buffers(self) -> InputBuffers: - assert self._input_buffers is not None - return self._input_buffers - @staticmethod def _reorder_segments( segments: list[RankSegment], @@ -285,14 +280,8 @@ def _build_batch_layout( # Therefore global = gathered[hidden_restore_idx] and # padded_gathered = global[padded_gather_idx]. hidden_restore_idx = np.empty(int(query_start_loc_np[-1]), dtype=np.int64) - required_num_tokens = max(per_rank_num_tokens) if padded_num_tokens is None: - padded_num_tokens = required_num_tokens - elif padded_num_tokens < required_num_tokens: - raise RuntimeError( - "PCP graph token capacity is smaller than the rank-local batch: " - f"{padded_num_tokens} < {required_num_tokens}." - ) + padded_num_tokens = max(per_rank_num_tokens) num_expanded_tokens = padded_num_tokens * self.pcp_world_size padded_gather_idx = np.zeros(num_expanded_tokens, dtype=np.int64) gathered_kv_write_mask = np.zeros(num_expanded_tokens, dtype=np.bool_) From a8aa573e53c7a83a82d8604dbd34346d00c63e3a Mon Sep 17 00:00:00 2001 From: QiuChunshuo Date: Sun, 23 Aug 2026 17:46:09 +0800 Subject: [PATCH 3/9] fix(pcp): preserve cudagraph token padding Signed-off-by: QiuChunshuo --- tests/v1/worker/test_gpu_pcp_manager.py | 32 +++++++++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 6 ++++- vllm/v1/worker/gpu/pcp_manager.py | 28 +++++++++++++++++----- 3 files changed, 59 insertions(+), 7 deletions(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index 4f9fc7232971..a2923dd3bf7c 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -1,9 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace import numpy as np +import pytest import torch +from vllm.config import CUDAGraphMode from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module from vllm.v1.worker.gpu.pcp_manager import PCPManager @@ -46,3 +49,32 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch): manager._gathered_kv_write_mask, torch.tensor([True, True, True, False, False, False, False, False]), ) + + +def test_global_cudagraph_padding_is_disabled_for_none(): + input_batch = SimpleNamespace(num_tokens=3, num_tokens_after_padding=4) + + assert ( + PCPManager._get_cudagraph_padded_num_tokens(input_batch, CUDAGraphMode.NONE) + is None + ) + + +@pytest.mark.parametrize( + "cudagraph_mode", + [ + CUDAGraphMode.PIECEWISE, + CUDAGraphMode.FULL, + CUDAGraphMode.FULL_DECODE_ONLY, + CUDAGraphMode.FULL_AND_PIECEWISE, + ], +) +def test_graph_modes_use_global_cudagraph_padding(cudagraph_mode): + input_batch = SimpleNamespace(num_tokens=3, num_tokens_after_padding=4) + assert PCPManager._get_cudagraph_padded_num_tokens(input_batch, cudagraph_mode) == 4 + + input_batch.num_tokens_after_padding = input_batch.num_tokens + assert ( + PCPManager._get_cudagraph_padded_num_tokens(input_batch, cudagraph_mode) + == input_batch.num_tokens + ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 3861bb52649d..390ce1116067 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1300,7 +1300,11 @@ def prepare_inputs( else None ), ) - return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch) + return pcp.maybe_partition_pcp_batch( + self.pcp_manager, + input_batch, + cudagraph_mode=batch_desc.cg_mode, + ) def prepare_attn( self, input_batch: InputBatch diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 25bf7d382fa9..7f7327a45598 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -318,7 +318,20 @@ def _build_batch_layout( ) return segments_by_rank, per_rank_num_tokens - def partition_batch(self, input_batch: InputBatch) -> InputBatch: + @staticmethod + def _get_cudagraph_padded_num_tokens( + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + ) -> int | None: + if cudagraph_mode == CUDAGraphMode.NONE: + return None + return input_batch.num_tokens_after_padding + + def partition_batch( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE, + ) -> InputBatch: assert self._req_states is not None assert self._input_buffers is not None req_states = self._req_states @@ -329,10 +342,9 @@ def partition_batch(self, input_batch: InputBatch) -> InputBatch: global_batch = input_batch self._global_batch = global_batch - num_tokens_after_padding = ( - global_batch.num_tokens_after_padding - if global_batch.num_tokens_after_padding > global_batch.num_tokens - else None + num_tokens_after_padding = self._get_cudagraph_padded_num_tokens( + global_batch, + cudagraph_mode, ) num_scheduled_tokens = global_batch.num_scheduled_tokens @@ -635,10 +647,14 @@ def restore_for_sampling( def maybe_partition_pcp_batch( manager: PCPManager | None, input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE, ) -> InputBatch: if manager is None: return input_batch - return manager.partition_batch(input_batch) + return manager.partition_batch( + input_batch, + cudagraph_mode=cudagraph_mode, + ) def maybe_get_pcp_dummy_slot_mappings( From 48ac3cc03720ffff801343a80d1695b8ed79e4d5 Mon Sep 17 00:00:00 2001 From: QiuChunshuo Date: Mon, 24 Aug 2026 09:51:23 +0800 Subject: [PATCH 4/9] fix(pcp): reuse capture-bound padding mask Signed-off-by: QiuChunshuo --- tests/v1/worker/test_gpu_pcp_manager.py | 27 +++++++++++++++++++++++++ vllm/v1/worker/gpu/pcp_manager.py | 14 ++++++++++++- 2 files changed, 40 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index a2923dd3bf7c..f84a38eebe6b 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -51,6 +51,33 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch): ) +def test_local_is_padding_reuses_global_batch_buffer(): + global_is_padding = torch.zeros(8, dtype=torch.bool) + global_batch = SimpleNamespace(is_padding=global_is_padding) + + local_is_padding = PCPManager._get_local_is_padding(global_batch, 4) + + assert local_is_padding.data_ptr() == global_is_padding.data_ptr() + local_is_padding[3] = True + assert global_is_padding.tolist() == [ + False, + False, + False, + True, + False, + False, + False, + False, + ] + + +def test_local_is_padding_rejects_short_global_batch_buffer(): + global_batch = SimpleNamespace(is_padding=torch.zeros(3, dtype=torch.bool)) + + with pytest.raises(RuntimeError, match="3 < 4"): + PCPManager._get_local_is_padding(global_batch, 4) + + def test_global_cudagraph_padding_is_disabled_for_none(): input_batch = SimpleNamespace(num_tokens=3, num_tokens_after_padding=4) diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 7f7327a45598..142b1d5a7781 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -327,6 +327,18 @@ def _get_cudagraph_padded_num_tokens( return None return input_batch.num_tokens_after_padding + @staticmethod + def _get_local_is_padding( + global_batch: InputBatch, num_local_tokens_padded: int + ) -> torch.Tensor: + if global_batch.is_padding.numel() < num_local_tokens_padded: + raise RuntimeError( + "Global is_padding buffer is smaller than the PCP-local padded " + f"token count: {global_batch.is_padding.numel()} < " + f"{num_local_tokens_padded}." + ) + return global_batch.is_padding[:num_local_tokens_padded] + def partition_batch( self, input_batch: InputBatch, @@ -470,7 +482,7 @@ def partition_batch( input_buffers.seq_lens[:num_local_reqs], ) seq_lens = input_buffers.seq_lens[:num_local_reqs] - is_padding = input_buffers.is_padding[:num_local_tokens_padded] + is_padding = self._get_local_is_padding(global_batch, num_local_tokens_padded) is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) if num_local_tokens_padded > num_local_tokens: From 3cfe87edbed16966b4b690262653221db32a0f40 Mon Sep 17 00:00:00 2001 From: QiuChunshuo Date: Mon, 24 Aug 2026 11:25:11 +0800 Subject: [PATCH 5/9] fix(pcp): capture graphs with PCP input buffers Signed-off-by: QiuChunshuo --- tests/v1/worker/test_gpu_pcp_manager.py | 36 +++++++++---------------- vllm/v1/worker/gpu/model_runner.py | 7 ++++- vllm/v1/worker/gpu/pcp_manager.py | 17 ++++-------- 3 files changed, 23 insertions(+), 37 deletions(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index f84a38eebe6b..65f676b5ff66 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -51,31 +51,19 @@ def test_replicated_decode_piecewise_graph_padding(monkeypatch): ) -def test_local_is_padding_reuses_global_batch_buffer(): - global_is_padding = torch.zeros(8, dtype=torch.bool) - global_batch = SimpleNamespace(is_padding=global_is_padding) - - local_is_padding = PCPManager._get_local_is_padding(global_batch, 4) - - assert local_is_padding.data_ptr() == global_is_padding.data_ptr() - local_is_padding[3] = True - assert global_is_padding.tolist() == [ - False, - False, - False, - True, - False, - False, - False, - False, - ] - - -def test_local_is_padding_rejects_short_global_batch_buffer(): - global_batch = SimpleNamespace(is_padding=torch.zeros(3, dtype=torch.bool)) +def test_input_buffers_are_exposed_for_cudagraph_capture(): + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + max_num_reqs=4, + max_num_tokens=8, + ) - with pytest.raises(RuntimeError, match="3 < 4"): - PCPManager._get_local_is_padding(global_batch, 4) + assert manager.input_buffers is manager._input_buffers + assert manager.input_buffers.input_ids.shape == (8,) + assert manager.input_buffers.positions.shape == (8,) + assert manager.input_buffers.is_padding.shape == (8,) def test_global_cudagraph_padding_is_disabled_for_none(): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 390ce1116067..4504a2fc5b2e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -877,10 +877,15 @@ def capture_model(self) -> int: self.model_state.encoder_runner.capture() if capture_decoder: + input_buffers = ( + self.pcp_manager.input_buffers + if self.pcp_manager is not None + else self.input_buffers + ) self.cudagraph_manager.capture( self.model, self.model_state, - self.input_buffers, + input_buffers, self.intermediate_tensors, self.block_tables, self.attn_groups, diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 142b1d5a7781..e3da56188431 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -327,17 +327,10 @@ def _get_cudagraph_padded_num_tokens( return None return input_batch.num_tokens_after_padding - @staticmethod - def _get_local_is_padding( - global_batch: InputBatch, num_local_tokens_padded: int - ) -> torch.Tensor: - if global_batch.is_padding.numel() < num_local_tokens_padded: - raise RuntimeError( - "Global is_padding buffer is smaller than the PCP-local padded " - f"token count: {global_batch.is_padding.numel()} < " - f"{num_local_tokens_padded}." - ) - return global_batch.is_padding[:num_local_tokens_padded] + @property + def input_buffers(self) -> InputBuffers: + assert self._input_buffers is not None + return self._input_buffers def partition_batch( self, @@ -482,7 +475,7 @@ def partition_batch( input_buffers.seq_lens[:num_local_reqs], ) seq_lens = input_buffers.seq_lens[:num_local_reqs] - is_padding = self._get_local_is_padding(global_batch, num_local_tokens_padded) + is_padding = input_buffers.is_padding[:num_local_tokens_padded] is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) if num_local_tokens_padded > num_local_tokens: From 1dfcbc289db596b38052c33606a6f32c09d65ee9 Mon Sep 17 00:00:00 2001 From: QiuChunshuo Date: Tue, 25 Aug 2026 15:04:25 +0800 Subject: [PATCH 6/9] perf(pcp): dispatch graphs by rank-local token count Co-authored-by: OpenAI Codex Co-authored-by: Lucas Wilkinson Signed-off-by: QiuChunshuo --- tests/v1/worker/test_gpu_pcp_manager.py | 64 ++++++++++------- vllm/v1/worker/gpu/model_runner.py | 11 ++- vllm/v1/worker/gpu/pcp_manager.py | 91 ++++++++++++++----------- 3 files changed, 101 insertions(+), 65 deletions(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index 65f676b5ff66..efb69ab120ad 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -1,12 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from types import SimpleNamespace - import numpy as np import pytest import torch -from vllm.config import CUDAGraphMode from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module from vllm.v1.worker.gpu.pcp_manager import PCPManager @@ -66,30 +63,47 @@ def test_input_buffers_are_exposed_for_cudagraph_capture(): assert manager.input_buffers.is_padding.shape == (8,) -def test_global_cudagraph_padding_is_disabled_for_none(): - input_batch = SimpleNamespace(num_tokens=3, num_tokens_after_padding=4) - - assert ( - PCPManager._get_cudagraph_padded_num_tokens(input_batch, CUDAGraphMode.NONE) - is None - ) - - @pytest.mark.parametrize( - "cudagraph_mode", + ("pcp_world_size", "num_scheduled_tokens", "is_prefilling", "expected"), [ - CUDAGraphMode.PIECEWISE, - CUDAGraphMode.FULL, - CUDAGraphMode.FULL_DECODE_ONLY, - CUDAGraphMode.FULL_AND_PIECEWISE, + (2, [8], [True], 4), + (2, [7], [True], 4), + (2, [3], [False], 3), + (2, [3, 8], [False, True], 7), + (4, [2, 9], [False, True], 5), ], ) -def test_graph_modes_use_global_cudagraph_padding(cudagraph_mode): - input_batch = SimpleNamespace(num_tokens=3, num_tokens_after_padding=4) - assert PCPManager._get_cudagraph_padded_num_tokens(input_batch, cudagraph_mode) == 4 - - input_batch.num_tokens_after_padding = input_batch.num_tokens - assert ( - PCPManager._get_cudagraph_padded_num_tokens(input_batch, cudagraph_mode) - == input_batch.num_tokens +def test_num_tokens_for_dispatch_uses_largest_pcp_rank( + pcp_world_size, num_scheduled_tokens, is_prefilling, expected +): + manager = PCPManager( + pcp_world_size=pcp_world_size, + pcp_rank=0, + device=torch.device("cpu"), ) + + actual = manager.get_num_tokens_for_dispatch( + np.asarray(num_scheduled_tokens, dtype=np.int32), + np.asarray(is_prefilling, dtype=np.bool_), + ) + + assert actual == expected + + +def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + dcp_world_size=1, + ) + monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu) + + with pytest.raises(ValueError, match="smaller than the largest rank-local batch"): + manager._build_batch_layout( + num_scheduled_tokens=np.ones(3, dtype=np.int32), + num_computed_tokens=np.full(3, 16, dtype=np.int32), + is_prefilling=np.zeros(3, dtype=np.bool_), + query_start_loc_np=np.arange(4, dtype=np.int32), + padded_num_tokens=2, + ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4504a2fc5b2e..d4461dd32f88 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1113,7 +1113,9 @@ def prepare_inputs( batch_desc: BatchExecutionDescriptor, ) -> InputBatch: num_tokens = batch_req_state.num_tokens - num_tokens_after_padding = batch_desc.num_tokens + num_tokens_after_padding = ( + num_tokens if self.pcp_manager is not None else batch_desc.num_tokens + ) assert num_tokens > 0 if envs.VLLM_MOE_SKIP_PADDING: # Mark trailing cudagraph-padding rows so kernels can skip work for @@ -1308,7 +1310,7 @@ def prepare_inputs( return pcp.maybe_partition_pcp_batch( self.pcp_manager, input_batch, - cudagraph_mode=batch_desc.cg_mode, + padded_num_tokens=batch_desc.num_tokens, ) def prepare_attn( @@ -1449,6 +1451,11 @@ def execute_model( ) if batch_req_state is not None: num_toks = batch_req_state.num_tokens + if self.pcp_manager is not None: + num_toks = self.pcp_manager.get_num_tokens_for_dispatch( + batch_req_state.num_scheduled_tokens, + batch_req_state.is_prefilling_np, + ) num_active_loras = 0 if self.lora_config: diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index e3da56188431..f8c5a705507d 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator from dataclasses import dataclass, replace import numpy as np @@ -192,15 +193,13 @@ def is_pure_prefill(segment: RankSegment) -> bool: rank_offset += segment.num_tokens return segments - def _get_rank_segments( + def _iter_rank_chunks( self, rank: int, num_scheduled_tokens: np.ndarray, - num_computed_tokens: np.ndarray, is_prefilling: np.ndarray, - query_start_loc_np: np.ndarray, - ) -> list[RankSegment]: - """Build one rank's attention-compatible DualChunkSwap rows. + ) -> Iterator[tuple[int, int, int]]: + """Yield ``(request index, query offset, length)`` for one PCP rank. PCP=4 partitions each prefill into eight chunks: @@ -210,14 +209,11 @@ def _get_rank_segments( rank 2: 2 5 rank 3: 3 4 """ - rank_segments = [] - rank_offset = 0 num_chunks = 2 * self.pcp_world_size for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens): query_len = int(num_tokens) if query_len == 0: continue - global_batch_start = int(query_start_loc_np[global_batch_req_idx]) chunk_indices: tuple[int, ...] if bool(is_prefilling[global_batch_req_idx]): chunk_size = (query_len + num_chunks - 1) // num_chunks @@ -231,17 +227,31 @@ def _get_rank_segments( chunk_len = min(chunk_size, query_len - chunk_offset) if chunk_len <= 0: continue - chunk_start = global_batch_start + chunk_offset - rank_segments.append( - RankSegment( - global_batch_req_idx=global_batch_req_idx, - global_batch_slice=slice(chunk_start, chunk_start + chunk_len), - rank_local_batch_slice=slice( - rank_offset, rank_offset + chunk_len - ), - ) + yield global_batch_req_idx, chunk_offset, chunk_len + + def _get_rank_segments( + self, + rank: int, + num_scheduled_tokens: np.ndarray, + num_computed_tokens: np.ndarray, + is_prefilling: np.ndarray, + query_start_loc_np: np.ndarray, + ) -> list[RankSegment]: + rank_segments = [] + rank_offset = 0 + for global_batch_req_idx, chunk_offset, chunk_len in self._iter_rank_chunks( + rank, num_scheduled_tokens, is_prefilling + ): + global_batch_start = int(query_start_loc_np[global_batch_req_idx]) + chunk_start = global_batch_start + chunk_offset + rank_segments.append( + RankSegment( + global_batch_req_idx=global_batch_req_idx, + global_batch_slice=slice(chunk_start, chunk_start + chunk_len), + rank_local_batch_slice=slice(rank_offset, rank_offset + chunk_len), ) - rank_offset += chunk_len + ) + rank_offset += chunk_len return self._reorder_segments( rank_segments, num_computed_tokens, @@ -282,6 +292,11 @@ def _build_batch_layout( hidden_restore_idx = np.empty(int(query_start_loc_np[-1]), dtype=np.int64) if padded_num_tokens is None: padded_num_tokens = max(per_rank_num_tokens) + elif padded_num_tokens < max(per_rank_num_tokens): + raise ValueError( + "PCP padded token count is smaller than the largest rank-local " + f"batch: {padded_num_tokens} < {max(per_rank_num_tokens)}." + ) num_expanded_tokens = padded_num_tokens * self.pcp_world_size padded_gather_idx = np.zeros(num_expanded_tokens, dtype=np.int64) gathered_kv_write_mask = np.zeros(num_expanded_tokens, dtype=np.bool_) @@ -318,14 +333,21 @@ def _build_batch_layout( ) return segments_by_rank, per_rank_num_tokens - @staticmethod - def _get_cudagraph_padded_num_tokens( - input_batch: InputBatch, - cudagraph_mode: CUDAGraphMode, - ) -> int | None: - if cudagraph_mode == CUDAGraphMode.NONE: - return None - return input_batch.num_tokens_after_padding + def get_num_tokens_for_dispatch( + self, + num_scheduled_tokens: np.ndarray, + is_prefilling: np.ndarray, + ) -> int: + """Return the largest real rank-local batch before graph padding.""" + return max( + sum( + chunk_len + for _, _, chunk_len in self._iter_rank_chunks( + rank, num_scheduled_tokens, is_prefilling + ) + ) + for rank in range(self.pcp_world_size) + ) @property def input_buffers(self) -> InputBuffers: @@ -335,7 +357,7 @@ def input_buffers(self) -> InputBuffers: def partition_batch( self, input_batch: InputBatch, - cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE, + padded_num_tokens: int | None = None, ) -> InputBatch: assert self._req_states is not None assert self._input_buffers is not None @@ -347,11 +369,6 @@ def partition_batch( global_batch = input_batch self._global_batch = global_batch - num_tokens_after_padding = self._get_cudagraph_padded_num_tokens( - global_batch, - cudagraph_mode, - ) - num_scheduled_tokens = global_batch.num_scheduled_tokens num_computed_tokens = global_batch.num_computed_tokens_np is_prefilling = global_batch.is_prefilling_np @@ -361,7 +378,7 @@ def partition_batch( num_computed_tokens, is_prefilling, global_batch.query_start_loc_np, - padded_num_tokens=num_tokens_after_padding, + padded_num_tokens=padded_num_tokens, ) local_segments = segments_by_rank[self.pcp_rank] @@ -411,9 +428,7 @@ def partition_batch( num_local_tokens = int(local_num_scheduled_tokens.sum()) num_local_tokens_padded = ( - max(per_rank_num_tokens) - if num_tokens_after_padding is None - else num_tokens_after_padding + max(per_rank_num_tokens) if padded_num_tokens is None else padded_num_tokens ) fresh_prefills = int( np.count_nonzero(is_prefilling & (num_computed_tokens == 0)) @@ -652,13 +667,13 @@ def restore_for_sampling( def maybe_partition_pcp_batch( manager: PCPManager | None, input_batch: InputBatch, - cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE, + padded_num_tokens: int | None = None, ) -> InputBatch: if manager is None: return input_batch return manager.partition_batch( input_batch, - cudagraph_mode=cudagraph_mode, + padded_num_tokens=padded_num_tokens, ) From 8e5fa50a6e47b2a62abe6056b2da2944785fd793 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 25 Aug 2026 16:45:05 -0400 Subject: [PATCH 7/9] Update vllm/v1/worker/gpu/model_runner.py Signed-off-by: Lucas Wilkinson --- vllm/v1/worker/gpu/model_runner.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index d4461dd32f88..be03e1e51584 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1114,10 +1114,7 @@ def prepare_inputs( ) -> InputBatch: num_tokens = batch_req_state.num_tokens num_tokens_after_padding = ( - num_tokens if self.pcp_manager is not None else batch_desc.num_tokens - ) - assert num_tokens > 0 - if envs.VLLM_MOE_SKIP_PADDING: + num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens) # Mark trailing cudagraph-padding rows so kernels can skip work for # them when supported. is_padding = self.input_buffers.is_padding From 8bdd185c6e7b16b8751ca96bfa294cdff292b61b Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 25 Aug 2026 16:48:19 -0400 Subject: [PATCH 8/9] Update vllm/v1/worker/gpu/model_runner.py Signed-off-by: Lucas Wilkinson --- vllm/v1/worker/gpu/model_runner.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index be03e1e51584..61955c5d23fb 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -877,11 +877,9 @@ def capture_model(self) -> int: self.model_state.encoder_runner.capture() if capture_decoder: - input_buffers = ( - self.pcp_manager.input_buffers - if self.pcp_manager is not None - else self.input_buffers - ) + input_buffers = self.input_buffers + if self.pcp_manager is not None: + input_buffers = self.pcp_manager.input_buffers self.cudagraph_manager.capture( self.model, self.model_state, From 5149d19bea24d7a53d35a45303aa5bda8213bdad Mon Sep 17 00:00:00 2001 From: QiuChunshuo Date: Wed, 26 Aug 2026 09:40:34 +0800 Subject: [PATCH 9/9] fix wrong update Signed-off-by: QiuChunshuo --- vllm/v1/worker/gpu/model_runner.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 46f5e0bc27d0..a58c86166df2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1144,8 +1144,9 @@ def prepare_inputs( batch_desc: BatchExecutionDescriptor, ) -> InputBatch: num_tokens = batch_req_state.num_tokens - num_tokens_after_padding = ( num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens) + assert num_tokens > 0 + if envs.VLLM_MOE_SKIP_PADDING: # Mark trailing cudagraph-padding rows so kernels can skip work for # them when supported. is_padding = self.input_buffers.is_padding