Bugfix: use PCP slot mappings for PIECEWISE capture (#53869)

Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml

@@ -4,7 +4,7 @@ num_questions: 1319
 num_fewshot: 5
 max_concurrency: 100
 server_args: >-
-  --enforce-eager
+  --compilation-config '{"cudagraph_mode":"PIECEWISE"}'
   --max-model-len 4096
   --max-num-batched-tokens 32768
   --safetensors-load-strategy prefetch

tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml

@@ -4,7 +4,7 @@ num_questions: 1319
 num_fewshot: 5
 max_concurrency: 100
 server_args: >-
-  --enforce-eager
+  --compilation-config '{"cudagraph_mode":"PIECEWISE"}'
   --max-model-len 4096
   --max-num-batched-tokens 32768
   --safetensors-load-strategy prefetch

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -17,8 +17,11 @@ from vllm.config import (
     VllmConfig,
 )
 from vllm.distributed.device_communicators import pynccl_allocator
+from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu.input_batch import InputBuffers
+from vllm.v1.worker.gpu.pcp_manager import PCPManager
 
 pytestmark = pytest.mark.cpu_test
 
@@ -112,6 +115,56 @@ def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
     mock_cuda_graph.assert_called_once()
 
 
+def test_piecewise_capture_uses_pcp_dummy_slot_mappings():
+    num_reqs = 32
+    num_tokens = 56
+    pcp_world_size = 2
+    input_buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
+
+    pcp_block_tables = SimpleNamespace(
+        num_kv_cache_groups=1,
+        input_block_tables=(torch.zeros(num_reqs * 2, 1, dtype=torch.int32),),
+    )
+    pcp_manager = PCPManager(
+        pcp_world_size=pcp_world_size,
+        pcp_rank=0,
+        device=torch.device("cpu"),
+        max_num_reqs=num_reqs,
+        max_num_tokens=num_tokens,
+        block_tables=pcp_block_tables,
+    )
+
+    block_tables = MagicMock()
+    block_tables.cp_size = 1
+    block_tables.get_dummy_block_tables.return_value = ()
+    block_tables.get_dummy_slot_mappings.return_value = torch.zeros(
+        1, num_tokens, dtype=torch.int64
+    )
+    model_state = MagicMock()
+    model_state.prepare_attn.return_value = {}
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0,
+        kv_cache_tensors=[],
+        kv_cache_groups=[],
+    )
+
+    gpu_cudagraph_utils.prepare_inputs_to_capture(
+        num_reqs,
+        num_tokens,
+        model_state,
+        input_buffers,
+        block_tables,
+        [],
+        kv_cache_config,
+        full_cudagraph=False,
+        pcp_manager=pcp_manager,
+    )
+
+    slot_mappings = model_state.prepare_attn.call_args.args[3]
+    assert slot_mappings.shape == (1, num_tokens * pcp_world_size)
+    block_tables.get_dummy_slot_mappings.assert_not_called()
+
+
 _DECODE_QUERY_LEN = 3
 
 

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -42,6 +42,7 @@ from vllm.v1.worker.utils import AttentionGroup
 
 if TYPE_CHECKING:
     from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+    from vllm.v1.worker.gpu.pcp_manager import PCPManager
 
 logger = init_logger(__name__)
 
@@ -494,6 +495,7 @@ class ModelCudaGraphManager(CudaGraphManager):
         block_tables: BlockTables,
         attn_groups: list[list[AttentionGroup]],
         kv_cache_config: KVCacheConfig,
+        pcp_manager: "PCPManager | None" = None,
         has_lora: bool = False,
         use_aux_hidden_state_outputs: bool = False,
         lora_capture_hook: Callable[[int, int, int], None] | None = None,
@@ -543,6 +545,7 @@ class ModelCudaGraphManager(CudaGraphManager):
                 kv_cache_config,
                 full_cudagraph=desc.cg_mode == CUDAGraphMode.FULL,
                 max_query_len=desc.max_query_len,
+                pcp_manager=pcp_manager,
             )
 
             # Capture with dummy rows marked as padding.
@@ -635,12 +638,16 @@ def prepare_inputs_to_capture(
     kv_cache_config: KVCacheConfig,
     full_cudagraph: bool,
     max_query_len: int | None = None,
+    pcp_manager: "PCPManager | None" = None,
 ) -> AttentionState:
     input_batch = InputBatch.make_dummy(
         num_reqs, num_tokens, input_buffers, max_query_len=max_query_len
     )
     input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
-    slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)
+    slot_mapping_provider: BlockTables | PCPManager = block_tables
+    if pcp_manager is not None:
+        slot_mapping_provider = pcp_manager
+    slot_mappings = slot_mapping_provider.get_dummy_slot_mappings(num_tokens)
     slot_mappings_by_layer = build_slot_mappings_by_layer(
         slot_mappings, kv_cache_config
     )

vllm/v1/worker/gpu/model_runner.py

@@ -920,6 +920,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     self.block_tables,
                     self.attn_groups,
                     self.kv_cache_config,
+                    pcp_manager=self.pcp_manager,
                     has_lora=self.lora_config is not None,
                     use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
                     lora_capture_hook=create_lora_capture_hook(self.lora_config, self),