[Feature][PCP] Support decode-only FULL CUDA graphs (#53867)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com> Signed-off-by: QiuChunshuo <chunshuoq@gmail.com> Co-authored-by: Codex <noreply@openai.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
Sep 15, 2026 at 16:27 UTC
676650397ff71cd802e7e0e30c2e46a215ab86ee
5 files changed
+283
-42
tests/v1/worker/test_gpu_pcp_manager.py
+155
-4
@@ -1,15 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from dataclasses import replace
4
+from types import SimpleNamespace
5
+from unittest.mock import MagicMock
6
7
import numpy as np
8
import pytest
9
import torch
10
11
+from vllm.config import CUDAGraphMode
12
from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens
13
from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils
14
from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module
12
-from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
15
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
16
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_context
17
from vllm.v1.worker.gpu.pcp_manager import PCPManager
18
19
@@ -20,6 +24,58 @@ def _copy_to_cpu(value, out=None, device=None):
24
return tensor
25
26
27
+def _make_config(cudagraph_mode: CUDAGraphMode):
28
+ return SimpleNamespace(
29
+ parallel_config=SimpleNamespace(
30
+ prefill_context_parallel_size=2,
31
+ decode_context_parallel_size=1,
32
+ pipeline_parallel_size=1,
33
+ dcp_comm_backend="ag_rs",
34
+ ),
35
+ model_config=SimpleNamespace(
36
+ use_mla=True,
37
+ is_encoder_decoder=False,
38
+ hf_text_config=SimpleNamespace(),
39
+ ),
40
+ lora_config=None,
41
+ speculative_config=None,
42
+ compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
43
+ )
44
+
45
+
46
+def _make_capture_manager(block_table: torch.Tensor):
47
+ block_tables = SimpleNamespace(
48
+ input_block_tables=(block_table,),
49
+ num_kv_cache_groups=1,
50
+ kernel_block_sizes=(2,),
51
+ blocks_per_kv_block=(1,),
52
+ )
53
+ manager = PCPManager(
54
+ pcp_world_size=2,
55
+ pcp_rank=0,
56
+ device=torch.device("cpu"),
57
+ max_num_reqs=4,
58
+ max_num_tokens=8,
59
+ block_tables=block_tables,
60
+ )
61
+ return manager, block_tables
62
+
63
+
64
+@pytest.mark.parametrize(
65
+ "cudagraph_mode",
66
+ [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE],
67
+)
68
+def test_validate_config_accepts_decode_only_full_graphs(cudagraph_mode):
69
+ PCPManager.validate_config(_make_config(cudagraph_mode), supports_mm_inputs=False)
70
+
71
+
72
+def test_validate_config_rejects_full_graph_for_prefills():
73
+ with pytest.raises(NotImplementedError, match="decode-only routines"):
74
+ PCPManager.validate_config(
75
+ _make_config(CUDAGraphMode.FULL), supports_mm_inputs=False
76
+ )
77
+
78
+
79
def test_replicated_decode_piecewise_graph_padding(monkeypatch):
80
manager = PCPManager(
81
pcp_world_size=2,
@@ -114,6 +170,89 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch):
170
)
171
172
173
+@pytest.mark.parametrize(
174
+ ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"),
175
+ [
176
+ (CUDAGraphMode.NONE, 4, None, None),
177
+ (CUDAGraphMode.PIECEWISE, None, 8, None),
178
+ (CUDAGraphMode.FULL, 4, 8, 4),
179
+ ],
180
+)
181
+def test_partition_padding_is_derived_from_batch_descriptor(
182
+ cg_mode, num_reqs, expected_tokens, expected_reqs
183
+):
184
+ manager = MagicMock()
185
+ input_batch = MagicMock()
186
+ manager.partition_batch.return_value = input_batch
187
+ batch_desc = BatchExecutionDescriptor(
188
+ cg_mode=cg_mode,
189
+ num_tokens=8,
190
+ num_reqs=num_reqs,
191
+ )
192
+
193
+ result = pcp_manager_module.maybe_partition_pcp_batch(
194
+ manager,
195
+ input_batch,
196
+ batch_desc,
197
+ )
198
+
199
+ assert result is input_batch
200
+ manager.partition_batch.assert_called_once_with(
201
+ input_batch,
202
+ padded_num_tokens=expected_tokens,
203
+ padded_num_reqs=expected_reqs,
204
+ )
205
+
206
+
207
+def test_capture_uses_pcp_persistent_inputs():
208
+ manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32))
209
+
210
+ dummy_batch = InputBatch.make_dummy(
211
+ num_reqs=4,
212
+ num_tokens=4,
213
+ input_buffers=InputBuffers(4, 8, torch.device("cpu")),
214
+ max_query_len=1,
215
+ )
216
+ dummy_batch.input_ids.copy_(torch.arange(4))
217
+ dummy_batch.positions.fill_(7)
218
+ input_batch = manager.prepare_inputs_to_capture(dummy_batch)
219
+
220
+ assert input_batch is not dummy_batch
221
+ assert input_batch.req_ids == dummy_batch.req_ids
222
+ for name in ("input_ids", "positions", "is_padding", "query_start_loc", "seq_lens"):
223
+ actual = getattr(input_batch, name)
224
+ torch.testing.assert_close(actual, getattr(dummy_batch, name))
225
+ assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr()
226
+
227
+
228
+def test_dummy_context_updates_pcp_local_block_tables():
229
+ global_block_table = torch.full((4, 4), -1, dtype=torch.int32)
230
+ manager, block_tables = _make_capture_manager(global_block_table)
231
+ input_batch = InputBatch.make_dummy(
232
+ num_reqs=2,
233
+ num_tokens=2,
234
+ input_buffers=manager.input_buffers,
235
+ max_query_len=1,
236
+ )
237
+ input_batch = manager.prepare_inputs_to_capture(input_batch)
238
+ local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs)
239
+
240
+ set_dummy_context(
241
+ input_batch,
242
+ block_tables,
243
+ context_len=3,
244
+ num_kv_blocks=16,
245
+ max_model_len=16,
246
+ input_block_tables=local_block_tables,
247
+ )
248
+
249
+ torch.testing.assert_close(
250
+ local_block_tables[0][:2, :2],
251
+ torch.tensor([[0, 1], [2, 3]], dtype=torch.int32),
252
+ )
253
+ assert torch.all(global_block_table == -1)
254
+
255
+
256
def _rank_rows(
257
pcp_rank: int,
258
pcp_world_size: int,
@@ -331,10 +470,22 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
470
global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2]
471
global_batch.dcp_local_seq_lens.fill_(-1)
472
334
- local_batch = manager.partition_batch(global_batch, padded_num_tokens=2)
473
+ local_batch = manager.partition_batch(
474
+ global_batch,
475
+ padded_num_tokens=4,
476
+ padded_num_reqs=4,
477
+ )
478
479
assert local_batch.dcp_local_seq_lens is None
337
- assert local_batch.seq_lens.tolist() == [17, 25]
480
+ assert local_batch.seq_lens.tolist() == [17, 25, 0, 0]
481
+ assert local_batch.dcp_local_seq_lens_cpu_upper_bound is not None
482
+ expected_dcp_upper_bound = get_dcp_local_seq_lens(
483
+ torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1
484
+ )
485
+ assert torch.equal(
486
+ local_batch.dcp_local_seq_lens_cpu_upper_bound,
487
+ torch.cat([expected_dcp_upper_bound, torch.zeros(2, dtype=torch.int32)]),
488
+ )
489
490
# What execute_model does next: derive DCP metadata from the final batch
491
# on the PCP-owned buffers.
@@ -348,7 +499,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch():
499
num_reqs_padded=local_batch.num_reqs_after_padding,
500
)
501
expected = get_dcp_local_seq_lens(
351
- torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1
502
+ torch.tensor([17, 25, 0, 0], dtype=torch.int32), 2, 0, 1
503
)
504
assert local_batch.dcp_local_seq_lens is not None
505
assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected)
vllm/v1/worker/gpu/cudagraph_utils.py
+5
-4
@@ -777,11 +777,12 @@ def prepare_inputs_to_capture(
777
input_batch = InputBatch.make_dummy(
778
num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
779
)
780
- input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
781
- slot_mapping_provider: BlockTables | PCPManager = block_tables
780
if pcp_manager is not None:
783
- slot_mapping_provider = pcp_manager
784
- slot_mappings = slot_mapping_provider.get_dummy_slot_mappings(num_tokens)
781
+ input_batch = pcp_manager.prepare_inputs_to_capture(input_batch)
782
+
783
+ block_table_provider = pcp_manager or block_tables
784
+ input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs)
785
+ slot_mappings = block_table_provider.get_dummy_slot_mappings(num_tokens)
786
slot_mappings_by_layer = build_slot_mappings_by_layer(
787
slot_mappings, kv_cache_config
788
)
vllm/v1/worker/gpu/input_batch.py
+6
-2
@@ -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 Sequence
4
from dataclasses import dataclass
5
from typing import TYPE_CHECKING
6
@@ -217,9 +218,12 @@ def set_dummy_context(
218
context_len: int,
219
num_kv_blocks: int,
220
max_model_len: int,
221
+ input_block_tables: Sequence[torch.Tensor] | None = None,
222
) -> None:
223
"""Give each dummy request context_len of context, used when profiling step cost."""
222
- if not block_tables.input_block_tables:
224
+ if input_block_tables is None:
225
+ input_block_tables = block_tables.input_block_tables
226
+ if not input_block_tables:
227
# Attention-free models have no KV context to fabricate.
228
return
229
num_reqs = input_batch.num_reqs
@@ -241,7 +245,7 @@ def set_dummy_context(
245
246
seq_len = context_len + query_len
247
for block_table, block_size, bpk in zip(
244
- block_tables.input_block_tables,
248
+ input_block_tables,
249
block_tables.kernel_block_sizes,
250
block_tables.blocks_per_kv_block,
251
):
vllm/v1/worker/gpu/model_runner.py
+20
-15
@@ -1468,7 +1468,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1468
input_batch = pcp.maybe_partition_pcp_batch(
1469
self.pcp_manager,
1470
input_batch,
1471
- padded_num_tokens=batch_desc.num_tokens,
1471
+ batch_desc,
1472
)
1473
return input_batch
1474
@@ -1496,7 +1496,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1496
def prepare_dummy_attn(
1497
self, input_batch: InputBatch, valid_state_slots: bool = False
1498
) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
1499
- block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
1499
+ block_table_provider = self.pcp_manager or self.block_tables
1500
+ block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs)
1501
if valid_state_slots:
1502
state_slots = torch.arange(
1503
1,
@@ -1745,25 +1746,29 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1746
self.input_buffers,
1747
max_query_len=batch_desc.max_query_len,
1748
)
1748
- if not skip_attn_for_dummy_run:
1749
- block_tables, slot_mappings = self.prepare_dummy_attn(
1750
- input_batch, valid_dummy_state_slots
1751
- )
1752
- if context_len:
1753
- set_dummy_context(
1754
- input_batch,
1755
- self.block_tables,
1756
- context_len,
1757
- self.kv_cache_config.num_blocks,
1758
- self.max_model_len,
1759
- )
1760
- else:
1749
+ if self.pcp_manager is not None:
1750
+ input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch)
1751
+ if skip_attn_for_dummy_run:
1752
assert batch_desc.cg_mode != CUDAGraphMode.FULL, (
1753
"Attention metadata must be prepared for dummy runs when using "
1754
"FULL cudagraph mode."
1755
)
1756
block_tables = None
1757
slot_mappings = None
1758
+ else:
1759
+ block_tables, slot_mappings = self.prepare_dummy_attn(
1760
+ input_batch, valid_dummy_state_slots
1761
+ )
1762
+ if not skip_attn_for_dummy_run and context_len:
1763
+ assert block_tables is not None
1764
+ set_dummy_context(
1765
+ input_batch,
1766
+ self.block_tables,
1767
+ context_len,
1768
+ self.kv_cache_config.num_blocks,
1769
+ self.max_model_len,
1770
+ input_block_tables=block_tables,
1771
+ )
1772
1773
attn_metadata = None
1774
slot_mappings_by_layer = None
vllm/v1/worker/gpu/pcp_manager.py
+97
-17
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Iterator
4
from dataclasses import dataclass, replace
5
+from typing import TYPE_CHECKING
6
7
import numpy as np
8
import torch
@@ -17,6 +18,9 @@ from vllm.v1.worker.gpu.input_batch import (
18
InputBuffers,
19
)
20
21
+if TYPE_CHECKING:
22
+ from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
23
+
24
logger = init_logger(__name__)
25
26
@@ -169,8 +173,14 @@ class PCPManager:
173
"MRV2 sparse MLA PCP does not support CUDA graphs yet. "
174
"Set -cc.cudagraph_mode=NONE."
175
)
172
- if cudagraph_mode.has_full_cudagraphs():
173
- raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.")
176
+ if (
177
+ cudagraph_mode.has_full_cudagraphs()
178
+ and not cudagraph_mode.separate_routine()
179
+ ):
180
+ raise NotImplementedError(
181
+ "MRV2 PCP supports full CUDA graphs for decode-only routines. "
182
+ "Use FULL_DECODE_ONLY, FULL_AND_PIECEWISE, PIECEWISE, or NONE."
183
+ )
184
if (
185
parallel_config.decode_context_parallel_size > 1
186
and parallel_config.dcp_comm_backend != "ag_rs"
@@ -375,6 +385,22 @@ class PCPManager:
385
for rank in range(self.pcp_world_size)
386
)
387
388
+ @staticmethod
389
+ def _resolve_num_reqs_after_padding(
390
+ input_batch: InputBatch,
391
+ padded_num_reqs: int | None,
392
+ num_local_reqs: int,
393
+ ) -> int:
394
+ if padded_num_reqs is None:
395
+ return num_local_reqs
396
+ if input_batch.has_prefill:
397
+ raise RuntimeError("PCP FULL graphs require a decode-only batch.")
398
+ assert padded_num_reqs >= num_local_reqs, (
399
+ "PCP graph request capacity must cover the rank-local batch: "
400
+ f"{padded_num_reqs} < {num_local_reqs}."
401
+ )
402
+ return padded_num_reqs
403
+
404
@property
405
def input_buffers(self) -> InputBuffers:
406
assert self._input_buffers is not None
@@ -384,6 +410,7 @@ class PCPManager:
410
self,
411
input_batch: InputBatch,
412
padded_num_tokens: int | None = None,
413
+ padded_num_reqs: int | None = None,
414
) -> InputBatch:
415
assert self._input_buffers is not None
416
input_buffers = self._input_buffers
@@ -414,10 +441,15 @@ class PCPManager:
441
]
442
443
num_local_reqs = len(local_segments)
417
- if num_local_reqs > input_buffers.max_num_reqs:
444
+ num_reqs_after_padding = self._resolve_num_reqs_after_padding(
445
+ global_batch,
446
+ padded_num_reqs,
447
+ num_local_reqs,
448
+ )
449
+ if num_reqs_after_padding > input_buffers.max_num_reqs:
450
raise RuntimeError(
419
- "PCP local request count exceeds the MRV2 input buffer size: "
420
- f"{num_local_reqs} > {input_buffers.max_num_reqs}."
451
+ "PCP padded local request count exceeds the MRV2 input buffer size: "
452
+ f"{num_reqs_after_padding} > {input_buffers.max_num_reqs}."
453
)
454
455
local_to_global_batch_req_idx_np = np.fromiter(
@@ -505,21 +537,25 @@ class PCPManager:
537
np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out)
538
local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens
539
async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc)
508
- local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1]
540
+ local_query_start_loc = input_buffers.query_start_loc[
541
+ : num_reqs_after_padding + 1
542
+ ]
543
544
local_to_global_req_idx = async_copy_to_gpu(
545
local_to_global_req_idx_np, device=self.device
546
)
513
- seq_lens = input_buffers.seq_lens[:num_local_reqs]
547
+ seq_lens = input_buffers.seq_lens[:num_reqs_after_padding]
548
+ real_seq_lens = seq_lens[:num_local_reqs]
549
if num_local_tokens > 0:
550
local_end_positions = torch.index_select(
551
input_buffers.positions,
552
0,
518
- local_query_start_loc[1:] - 1,
553
+ local_query_start_loc[1 : num_local_reqs + 1] - 1,
554
)
520
- seq_lens.copy_(local_end_positions + 1)
555
+ real_seq_lens.copy_(local_end_positions + 1)
556
else:
522
- seq_lens.zero_()
557
+ real_seq_lens.zero_()
558
+ seq_lens[num_local_reqs:].zero_()
559
is_padding = input_buffers.is_padding[:num_local_tokens_padded]
560
is_padding[:num_local_tokens].fill_(False)
561
is_padding[num_local_tokens:].fill_(True)
@@ -552,11 +588,15 @@ class PCPManager:
588
local_num_computed_prefill_tokens_np = np.minimum(
589
local_start_pos_np, local_prefill_len_np
590
)
555
- local_is_prefilling_np = (
591
+ real_local_is_prefilling_np = (
592
local_num_computed_prefill_tokens_np < local_prefill_len_np
593
)
558
- seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32)
559
- seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens
594
+ local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_)
595
+ local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np
596
+ seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32)
597
+ seq_lens_cpu_upper_bound_np[:num_local_reqs] = (
598
+ local_start_pos_np + local_num_scheduled_tokens
599
+ )
600
dcp_local_seq_lens_cpu_upper_bound = None
601
if self.dcp_world_size > 1:
602
# The largest DCP shard of each row's whole request, identical on
@@ -564,18 +604,24 @@ class PCPManager:
604
request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[
605
local_to_global_batch_req_idx_np
606
]
567
- dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens(
607
+ real_dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens(
608
torch.from_numpy(request_seq_lens.astype(np.int32)),
609
self.dcp_world_size,
610
0,
611
self.cp_interleave,
612
)
613
+ dcp_local_seq_lens_cpu_upper_bound = torch.zeros(
614
+ num_reqs_after_padding, dtype=torch.int32
615
+ )
616
+ dcp_local_seq_lens_cpu_upper_bound[:num_local_reqs].copy_(
617
+ real_dcp_local_seq_lens_cpu_upper_bound
618
+ )
619
620
self._local_batch = replace(
621
input_batch,
622
req_ids=local_req_ids,
623
num_reqs=num_local_reqs,
578
- num_reqs_after_padding=num_local_reqs,
624
+ num_reqs_after_padding=num_reqs_after_padding,
625
idx_mapping=local_to_global_req_idx,
626
idx_mapping_np=local_to_global_req_idx_np,
627
expanded_idx_mapping=local_to_global_req_idx,
@@ -588,7 +634,7 @@ class PCPManager:
634
num_draft_tokens=0,
635
num_draft_tokens_per_req=None,
636
query_start_loc=local_query_start_loc,
591
- query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1],
637
+ query_start_loc_np=local_query_start_loc_np[: num_reqs_after_padding + 1],
638
seq_lens=seq_lens,
639
seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np),
640
dcp_local_seq_lens=None,
@@ -608,6 +654,31 @@ class PCPManager:
654
)
655
return self._local_batch
656
657
+ def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch:
658
+ """Stage a capture or dummy batch in persistent PCP input buffers."""
659
+ input_buffers = self.input_buffers
660
+ num_reqs = input_batch.num_reqs_after_padding
661
+ num_tokens = input_batch.num_tokens_after_padding
662
+ input_batch = replace(
663
+ input_batch,
664
+ input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids),
665
+ positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions),
666
+ is_padding=input_buffers.is_padding[:num_tokens].copy_(
667
+ input_batch.is_padding
668
+ ),
669
+ query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_(
670
+ input_batch.query_start_loc
671
+ ),
672
+ seq_lens=input_buffers.seq_lens[:num_reqs].copy_(input_batch.seq_lens),
673
+ )
674
+ return input_batch
675
+
676
+ def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
677
+ assert self._local_block_tables is not None
678
+ return tuple(
679
+ block_table[:num_reqs].zero_() for block_table in self._local_block_tables
680
+ )
681
+
682
def prepare_attn(
683
self, input_batch: InputBatch
684
) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
@@ -728,13 +799,22 @@ class PCPManager:
799
def maybe_partition_pcp_batch(
800
manager: PCPManager | None,
801
input_batch: InputBatch,
731
- padded_num_tokens: int | None = None,
802
+ batch_desc: "BatchExecutionDescriptor",
803
) -> InputBatch:
804
if manager is None:
805
return input_batch
806
+
807
+ padded_num_tokens = None
808
+ padded_num_reqs = None
809
+ if batch_desc.cg_mode != CUDAGraphMode.NONE:
810
+ padded_num_tokens = batch_desc.num_tokens
811
+ if batch_desc.cg_mode == CUDAGraphMode.FULL:
812
+ padded_num_reqs = batch_desc.num_reqs
813
+
814
return manager.partition_batch(
815
input_batch,
816
padded_num_tokens=padded_num_tokens,
817
+ padded_num_reqs=padded_num_reqs,
818
)
819
820