From b1fbbc2ade51e3826bc92e4733c9c692ee21d42d Mon Sep 17 00:00:00 2001 From: Qiu Chunshuo Date: Wed, 26 Aug 2026 13:58:13 +0800 Subject: [PATCH] BugFix(PCP): use persistent input buffers for PIECEWISE CUDA graphs (#53515) Signed-off-by: QiuChunshuo Signed-off-by: Lucas Wilkinson Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_gpu_pcp_manager.py | 109 ++++++++++++++++++++++++ vllm/v1/worker/gpu/model_runner.py | 18 +++- vllm/v1/worker/gpu/pcp_manager.py | 93 +++++++++++++++----- 3 files changed, 195 insertions(+), 25 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 0000000000..efb69ab120 --- /dev/null +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import numpy as np +import pytest +import torch + +from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module +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_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_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, + ) + + 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,) + + +@pytest.mark.parametrize( + ("pcp_world_size", "num_scheduled_tokens", "is_prefilling", "expected"), + [ + (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_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 32fbd63fe8..a58c86166d 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -907,10 +907,13 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.model_state.encoder_runner.capture() if capture_decoder: + 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, - self.input_buffers, + input_buffers, self.intermediate_tensors, self.block_tables, self.attn_groups, @@ -1141,7 +1144,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): batch_desc: BatchExecutionDescriptor, ) -> InputBatch: num_tokens = batch_req_state.num_tokens - num_tokens_after_padding = batch_desc.num_tokens + 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 @@ -1333,7 +1336,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): else None ), ) - return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch) + return pcp.maybe_partition_pcp_batch( + self.pcp_manager, + input_batch, + padded_num_tokens=batch_desc.num_tokens, + ) def prepare_attn( self, input_batch: InputBatch @@ -1517,6 +1524,11 @@ class GPUModelRunner(LoRAModelRunnerMixin): ) 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 d80e246f17..f8c5a70550 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 @@ class PCPManager: 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 @@ class PCPManager: 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 @@ class PCPManager: 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, @@ -255,6 +265,7 @@ class PCPManager: 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 +290,13 @@ class PCPManager: # 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) + 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_) @@ -316,7 +333,32 @@ class PCPManager: ) return segments_by_rank, per_rank_num_tokens - def partition_batch(self, input_batch: InputBatch) -> InputBatch: + 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: + assert self._input_buffers is not None + return self._input_buffers + + def partition_batch( + self, + input_batch: InputBatch, + padded_num_tokens: int | None = None, + ) -> InputBatch: assert self._req_states is not None assert self._input_buffers is not None req_states = self._req_states @@ -336,6 +378,7 @@ class PCPManager: num_computed_tokens, is_prefilling, global_batch.query_start_loc_np, + padded_num_tokens=padded_num_tokens, ) local_segments = segments_by_rank[self.pcp_rank] @@ -384,7 +427,9 @@ class PCPManager: ] 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 padded_num_tokens is None else padded_num_tokens + ) fresh_prefills = int( np.count_nonzero(is_prefilling & (num_computed_tokens == 0)) ) @@ -622,10 +667,14 @@ class PCPManager: def maybe_partition_pcp_batch( manager: PCPManager | None, input_batch: InputBatch, + padded_num_tokens: int | None = None, ) -> InputBatch: if manager is None: return input_batch - return manager.partition_batch(input_batch) + return manager.partition_batch( + input_batch, + padded_num_tokens=padded_num_tokens, + ) def maybe_get_pcp_dummy_slot_mappings(