Bugfix: use PCP slot mappings for PIECEWISE capture (#53869)
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>
Qiu Chunshuo committed
Aug 27, 2026 at 13:39 UTC
b3af042abd8fe5a297ec3ec72db276fd661a67b3
5 files changed
+64
-3
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+1
-1
@@ -4,7 +4,7 @@ num_questions: 1319
4
num_fewshot: 5
5
max_concurrency: 100
6
server_args: >-
7
- --enforce-eager
7
+ --compilation-config '{"cudagraph_mode":"PIECEWISE"}'
8
--max-model-len 4096
9
--max-num-batched-tokens 32768
10
--safetensors-load-strategy prefetch
tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+1
-1
@@ -4,7 +4,7 @@ num_questions: 1319
4
num_fewshot: 5
5
max_concurrency: 100
6
server_args: >-
7
- --enforce-eager
7
+ --compilation-config '{"cudagraph_mode":"PIECEWISE"}'
8
--max-model-len 4096
9
--max-num-batched-tokens 32768
10
--safetensors-load-strategy prefetch
tests/v1/cudagraph/test_cudagraph_manager.py
+53
@@ -17,8 +17,11 @@ from vllm.config import (
17
VllmConfig,
18
)
19
from vllm.distributed.device_communicators import pynccl_allocator
20
+from vllm.v1.kv_cache_interface import KVCacheConfig
21
from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
22
from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
23
+from vllm.v1.worker.gpu.input_batch import InputBuffers
24
+from vllm.v1.worker.gpu.pcp_manager import PCPManager
25
26
pytestmark = pytest.mark.cpu_test
27
@@ -112,6 +115,56 @@ def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
115
mock_cuda_graph.assert_called_once()
116
117
118
+def test_piecewise_capture_uses_pcp_dummy_slot_mappings():
119
+ num_reqs = 32
120
+ num_tokens = 56
121
+ pcp_world_size = 2
122
+ input_buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
123
+
124
+ pcp_block_tables = SimpleNamespace(
125
+ num_kv_cache_groups=1,
126
+ input_block_tables=(torch.zeros(num_reqs * 2, 1, dtype=torch.int32),),
127
+ )
128
+ pcp_manager = PCPManager(
129
+ pcp_world_size=pcp_world_size,
130
+ pcp_rank=0,
131
+ device=torch.device("cpu"),
132
+ max_num_reqs=num_reqs,
133
+ max_num_tokens=num_tokens,
134
+ block_tables=pcp_block_tables,
135
+ )
136
+
137
+ block_tables = MagicMock()
138
+ block_tables.cp_size = 1
139
+ block_tables.get_dummy_block_tables.return_value = ()
140
+ block_tables.get_dummy_slot_mappings.return_value = torch.zeros(
141
+ 1, num_tokens, dtype=torch.int64
142
+ )
143
+ model_state = MagicMock()
144
+ model_state.prepare_attn.return_value = {}
145
+ kv_cache_config = KVCacheConfig(
146
+ num_blocks=0,
147
+ kv_cache_tensors=[],
148
+ kv_cache_groups=[],
149
+ )
150
+
151
+ gpu_cudagraph_utils.prepare_inputs_to_capture(
152
+ num_reqs,
153
+ num_tokens,
154
+ model_state,
155
+ input_buffers,
156
+ block_tables,
157
+ [],
158
+ kv_cache_config,
159
+ full_cudagraph=False,
160
+ pcp_manager=pcp_manager,
161
+ )
162
+
163
+ slot_mappings = model_state.prepare_attn.call_args.args[3]
164
+ assert slot_mappings.shape == (1, num_tokens * pcp_world_size)
165
+ block_tables.get_dummy_slot_mappings.assert_not_called()
166
+
167
+
168
_DECODE_QUERY_LEN = 3
169
170
vllm/v1/worker/gpu/cudagraph_utils.py
+8
-1
@@ -42,6 +42,7 @@ from vllm.v1.worker.utils import AttentionGroup
42
43
if TYPE_CHECKING:
44
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
45
+ from vllm.v1.worker.gpu.pcp_manager import PCPManager
46
47
logger = init_logger(__name__)
48
@@ -494,6 +495,7 @@ class ModelCudaGraphManager(CudaGraphManager):
495
block_tables: BlockTables,
496
attn_groups: list[list[AttentionGroup]],
497
kv_cache_config: KVCacheConfig,
498
+ pcp_manager: "PCPManager | None" = None,
499
has_lora: bool = False,
500
use_aux_hidden_state_outputs: bool = False,
501
lora_capture_hook: Callable[[int, int, int], None] | None = None,
@@ -543,6 +545,7 @@ class ModelCudaGraphManager(CudaGraphManager):
545
kv_cache_config,
546
full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
547
max_query_len=desc.max_query_len,
548
+ pcp_manager=pcp_manager,
549
)
550
551
# Capture with dummy rows marked as padding.
@@ -635,12 +638,16 @@ def prepare_inputs_to_capture(
638
kv_cache_config: KVCacheConfig,
639
full_cudagraph: bool,
640
max_query_len: int | None = None,
641
+ pcp_manager: "PCPManager | None" = None,
642
) -> AttentionState:
643
input_batch = InputBatch.make_dummy(
644
num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
645
)
646
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
643
- slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)
647
+ slot_mapping_provider: BlockTables | PCPManager = block_tables
648
+ if pcp_manager is not None:
649
+ slot_mapping_provider = pcp_manager
650
+ slot_mappings = slot_mapping_provider.get_dummy_slot_mappings(num_tokens)
651
slot_mappings_by_layer = build_slot_mappings_by_layer(
652
slot_mappings, kv_cache_config
653
)
vllm/v1/worker/gpu/model_runner.py
+1
@@ -920,6 +920,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
920
self.block_tables,
921
self.attn_groups,
922
self.kv_cache_config,
923
+ pcp_manager=self.pcp_manager,
924
has_lora=self.lora_config is not None,
925
use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
926
lora_capture_hook=create_lora_capture_hook(self.lora_config, self),