BugFix(PCP): use persistent input buffers for PIECEWISE CUDA graphs (#53515)

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>
Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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,
+        )

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:

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(