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>
Qiu Chunshuo committed
Aug 26, 2026 at 13:58 UTC
b1fbbc2ade51e3826bc92e4733c9c692ee21d42d
3 files changed
+195
-25
tests/v1/worker/test_gpu_pcp_manager.py
new
+109
@@ -0,0 +1,109 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
+import pytest
5
+import torch
6
+
7
+from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module
8
+from vllm.v1.worker.gpu.pcp_manager import PCPManager
9
+
10
+
11
+def _copy_to_cpu(value, out=None, device=None):
12
+ tensor = torch.from_numpy(value) if isinstance(value, np.ndarray) else value
13
+ if out is not None:
14
+ return out.copy_(tensor)
15
+ return tensor
16
+
17
+
18
+def test_replicated_decode_piecewise_graph_padding(monkeypatch):
19
+ manager = PCPManager(
20
+ pcp_world_size=2,
21
+ pcp_rank=0,
22
+ device=torch.device("cpu"),
23
+ dcp_world_size=1,
24
+ )
25
+ monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu)
26
+
27
+ segments_by_rank, per_rank_num_tokens = manager._build_batch_layout(
28
+ num_scheduled_tokens=np.ones(3, dtype=np.int32),
29
+ num_computed_tokens=np.full(3, 16, dtype=np.int32),
30
+ is_prefilling=np.zeros(3, dtype=np.bool_),
31
+ query_start_loc_np=np.arange(4, dtype=np.int32),
32
+ padded_num_tokens=4,
33
+ )
34
+
35
+ assert per_rank_num_tokens == [3, 3]
36
+ request_indices = [
37
+ [segment.global_batch_req_idx for segment in rank] for rank in segments_by_rank
38
+ ]
39
+ assert request_indices == [[0, 1, 2], [0, 1, 2]]
40
+ assert torch.equal(manager._hidden_restore_idx, torch.tensor([0, 1, 2]))
41
+ assert torch.equal(
42
+ manager._padded_gather_idx,
43
+ torch.tensor([0, 1, 2, 0, 0, 1, 2, 0]),
44
+ )
45
+ assert torch.equal(
46
+ manager._gathered_kv_write_mask,
47
+ torch.tensor([True, True, True, False, False, False, False, False]),
48
+ )
49
+
50
+
51
+def test_input_buffers_are_exposed_for_cudagraph_capture():
52
+ manager = PCPManager(
53
+ pcp_world_size=2,
54
+ pcp_rank=0,
55
+ device=torch.device("cpu"),
56
+ max_num_reqs=4,
57
+ max_num_tokens=8,
58
+ )
59
+
60
+ assert manager.input_buffers is manager._input_buffers
61
+ assert manager.input_buffers.input_ids.shape == (8,)
62
+ assert manager.input_buffers.positions.shape == (8,)
63
+ assert manager.input_buffers.is_padding.shape == (8,)
64
+
65
+
66
+@pytest.mark.parametrize(
67
+ ("pcp_world_size", "num_scheduled_tokens", "is_prefilling", "expected"),
68
+ [
69
+ (2, [8], [True], 4),
70
+ (2, [7], [True], 4),
71
+ (2, [3], [False], 3),
72
+ (2, [3, 8], [False, True], 7),
73
+ (4, [2, 9], [False, True], 5),
74
+ ],
75
+)
76
+def test_num_tokens_for_dispatch_uses_largest_pcp_rank(
77
+ pcp_world_size, num_scheduled_tokens, is_prefilling, expected
78
+):
79
+ manager = PCPManager(
80
+ pcp_world_size=pcp_world_size,
81
+ pcp_rank=0,
82
+ device=torch.device("cpu"),
83
+ )
84
+
85
+ actual = manager.get_num_tokens_for_dispatch(
86
+ np.asarray(num_scheduled_tokens, dtype=np.int32),
87
+ np.asarray(is_prefilling, dtype=np.bool_),
88
+ )
89
+
90
+ assert actual == expected
91
+
92
+
93
+def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
94
+ manager = PCPManager(
95
+ pcp_world_size=2,
96
+ pcp_rank=0,
97
+ device=torch.device("cpu"),
98
+ dcp_world_size=1,
99
+ )
100
+ monkeypatch.setattr(pcp_manager_module, "async_copy_to_gpu", _copy_to_cpu)
101
+
102
+ with pytest.raises(ValueError, match="smaller than the largest rank-local batch"):
103
+ manager._build_batch_layout(
104
+ num_scheduled_tokens=np.ones(3, dtype=np.int32),
105
+ num_computed_tokens=np.full(3, 16, dtype=np.int32),
106
+ is_prefilling=np.zeros(3, dtype=np.bool_),
107
+ query_start_loc_np=np.arange(4, dtype=np.int32),
108
+ padded_num_tokens=2,
109
+ )
vllm/v1/worker/gpu/model_runner.py
+15
-3
@@ -907,10 +907,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
907
self.model_state.encoder_runner.capture()
908
909
if capture_decoder:
910
+ input_buffers = self.input_buffers
911
+ if self.pcp_manager is not None:
912
+ input_buffers = self.pcp_manager.input_buffers
913
self.cudagraph_manager.capture(
914
self.model,
915
self.model_state,
913
- self.input_buffers,
916
+ input_buffers,
917
self.intermediate_tensors,
918
self.block_tables,
919
self.attn_groups,
@@ -1141,7 +1144,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1144
batch_desc: BatchExecutionDescriptor,
1145
) -> InputBatch:
1146
num_tokens = batch_req_state.num_tokens
1144
- num_tokens_after_padding = batch_desc.num_tokens
1147
+ num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
1148
assert num_tokens > 0
1149
if envs.VLLM_MOE_SKIP_PADDING:
1150
# Mark trailing cudagraph-padding rows so kernels can skip work for
@@ -1333,7 +1336,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1336
else None
1337
),
1338
)
1336
- return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch)
1339
+ return pcp.maybe_partition_pcp_batch(
1340
+ self.pcp_manager,
1341
+ input_batch,
1342
+ padded_num_tokens=batch_desc.num_tokens,
1343
+ )
1344
1345
def prepare_attn(
1346
self, input_batch: InputBatch
@@ -1517,6 +1524,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1524
)
1525
if batch_req_state is not None:
1526
num_toks = batch_req_state.num_tokens
1527
+ if self.pcp_manager is not None:
1528
+ num_toks = self.pcp_manager.get_num_tokens_for_dispatch(
1529
+ batch_req_state.num_scheduled_tokens,
1530
+ batch_req_state.is_prefilling_np,
1531
+ )
1532
1533
num_active_loras = 0
1534
if self.lora_config:
vllm/v1/worker/gpu/pcp_manager.py
+71
-22
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from collections.abc import Iterator
4
from dataclasses import dataclass, replace
5
6
import numpy as np
@@ -192,15 +193,13 @@ class PCPManager:
193
rank_offset += segment.num_tokens
194
return segments
195
195
- def _get_rank_segments(
196
+ def _iter_rank_chunks(
197
self,
198
rank: int,
199
num_scheduled_tokens: np.ndarray,
199
- num_computed_tokens: np.ndarray,
200
is_prefilling: np.ndarray,
201
- query_start_loc_np: np.ndarray,
202
- ) -> list[RankSegment]:
203
- """Build one rank's attention-compatible DualChunkSwap rows.
201
+ ) -> Iterator[tuple[int, int, int]]:
202
+ """Yield ``(request index, query offset, length)`` for one PCP rank.
203
204
PCP=4 partitions each prefill into eight chunks:
205
@@ -210,14 +209,11 @@ class PCPManager:
209
rank 2: 2 5
210
rank 3: 3 4
211
"""
213
- rank_segments = []
214
- rank_offset = 0
212
num_chunks = 2 * self.pcp_world_size
213
for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens):
214
query_len = int(num_tokens)
215
if query_len == 0:
216
continue
220
- global_batch_start = int(query_start_loc_np[global_batch_req_idx])
217
chunk_indices: tuple[int, ...]
218
if bool(is_prefilling[global_batch_req_idx]):
219
chunk_size = (query_len + num_chunks - 1) // num_chunks
@@ -231,17 +227,31 @@ class PCPManager:
227
chunk_len = min(chunk_size, query_len - chunk_offset)
228
if chunk_len <= 0:
229
continue
234
- chunk_start = global_batch_start + chunk_offset
235
- rank_segments.append(
236
- RankSegment(
237
- global_batch_req_idx=global_batch_req_idx,
238
- global_batch_slice=slice(chunk_start, chunk_start + chunk_len),
239
- rank_local_batch_slice=slice(
240
- rank_offset, rank_offset + chunk_len
241
- ),
242
- )
230
+ yield global_batch_req_idx, chunk_offset, chunk_len
231
+
232
+ def _get_rank_segments(
233
+ self,
234
+ rank: int,
235
+ num_scheduled_tokens: np.ndarray,
236
+ num_computed_tokens: np.ndarray,
237
+ is_prefilling: np.ndarray,
238
+ query_start_loc_np: np.ndarray,
239
+ ) -> list[RankSegment]:
240
+ rank_segments = []
241
+ rank_offset = 0
242
+ for global_batch_req_idx, chunk_offset, chunk_len in self._iter_rank_chunks(
243
+ rank, num_scheduled_tokens, is_prefilling
244
+ ):
245
+ global_batch_start = int(query_start_loc_np[global_batch_req_idx])
246
+ chunk_start = global_batch_start + chunk_offset
247
+ rank_segments.append(
248
+ RankSegment(
249
+ global_batch_req_idx=global_batch_req_idx,
250
+ global_batch_slice=slice(chunk_start, chunk_start + chunk_len),
251
+ rank_local_batch_slice=slice(rank_offset, rank_offset + chunk_len),
252
)
244
- rank_offset += chunk_len
253
+ )
254
+ rank_offset += chunk_len
255
return self._reorder_segments(
256
rank_segments,
257
num_computed_tokens,
@@ -255,6 +265,7 @@ class PCPManager:
265
num_computed_tokens: np.ndarray,
266
is_prefilling: np.ndarray,
267
query_start_loc_np: np.ndarray,
268
+ padded_num_tokens: int | None = None,
269
) -> tuple[list[list[RankSegment]], list[int]]:
270
segments_by_rank = []
271
per_rank_num_tokens = []
@@ -279,7 +290,13 @@ class PCPManager:
290
# Therefore global = gathered[hidden_restore_idx] and
291
# padded_gathered = global[padded_gather_idx].
292
hidden_restore_idx = np.empty(int(query_start_loc_np[-1]), dtype=np.int64)
282
- padded_num_tokens = max(per_rank_num_tokens)
293
+ if padded_num_tokens is None:
294
+ padded_num_tokens = max(per_rank_num_tokens)
295
+ elif padded_num_tokens < max(per_rank_num_tokens):
296
+ raise ValueError(
297
+ "PCP padded token count is smaller than the largest rank-local "
298
+ f"batch: {padded_num_tokens} < {max(per_rank_num_tokens)}."
299
+ )
300
num_expanded_tokens = padded_num_tokens * self.pcp_world_size
301
padded_gather_idx = np.zeros(num_expanded_tokens, dtype=np.int64)
302
gathered_kv_write_mask = np.zeros(num_expanded_tokens, dtype=np.bool_)
@@ -316,7 +333,32 @@ class PCPManager:
333
)
334
return segments_by_rank, per_rank_num_tokens
335
319
- def partition_batch(self, input_batch: InputBatch) -> InputBatch:
336
+ def get_num_tokens_for_dispatch(
337
+ self,
338
+ num_scheduled_tokens: np.ndarray,
339
+ is_prefilling: np.ndarray,
340
+ ) -> int:
341
+ """Return the largest real rank-local batch before graph padding."""
342
+ return max(
343
+ sum(
344
+ chunk_len
345
+ for _, _, chunk_len in self._iter_rank_chunks(
346
+ rank, num_scheduled_tokens, is_prefilling
347
+ )
348
+ )
349
+ for rank in range(self.pcp_world_size)
350
+ )
351
+
352
+ @property
353
+ def input_buffers(self) -> InputBuffers:
354
+ assert self._input_buffers is not None
355
+ return self._input_buffers
356
+
357
+ def partition_batch(
358
+ self,
359
+ input_batch: InputBatch,
360
+ padded_num_tokens: int | None = None,
361
+ ) -> InputBatch:
362
assert self._req_states is not None
363
assert self._input_buffers is not None
364
req_states = self._req_states
@@ -336,6 +378,7 @@ class PCPManager:
378
num_computed_tokens,
379
is_prefilling,
380
global_batch.query_start_loc_np,
381
+ padded_num_tokens=padded_num_tokens,
382
)
383
384
local_segments = segments_by_rank[self.pcp_rank]
@@ -384,7 +427,9 @@ class PCPManager:
427
]
428
429
num_local_tokens = int(local_num_scheduled_tokens.sum())
387
- num_local_tokens_padded = max(per_rank_num_tokens)
430
+ num_local_tokens_padded = (
431
+ max(per_rank_num_tokens) if padded_num_tokens is None else padded_num_tokens
432
+ )
433
fresh_prefills = int(
434
np.count_nonzero(is_prefilling & (num_computed_tokens == 0))
435
)
@@ -622,10 +667,14 @@ class PCPManager:
667
def maybe_partition_pcp_batch(
668
manager: PCPManager | None,
669
input_batch: InputBatch,
670
+ padded_num_tokens: int | None = None,
671
) -> InputBatch:
672
if manager is None:
673
return input_batch
628
- return manager.partition_batch(input_batch)
674
+ return manager.partition_batch(
675
+ input_batch,
676
+ padded_num_tokens=padded_num_tokens,
677
+ )
678
679
680
def maybe_get_pcp_dummy_slot_mappings(