From b3af042abd8fe5a297ec3ec72db276fd661a67b3 Mon Sep 17 00:00:00 2001 From: Qiu Chunshuo Date: Thu, 27 Aug 2026 13:39:36 +0800 Subject: [PATCH] Bugfix: use PCP slot mappings for PIECEWISE capture (#53869) Signed-off-by: QiuChunshuo --- .../configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml | 2 +- .../configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml | 2 +- tests/v1/cudagraph/test_cudagraph_manager.py | 53 +++++++++++++++++++ vllm/v1/worker/gpu/cudagraph_utils.py | 9 +++- vllm/v1/worker/gpu/model_runner.py | 1 + 5 files changed, 64 insertions(+), 3 deletions(-) diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml index d4ddfcc99f..9679fbebb7 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml +++ b/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 diff --git a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml b/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml index 1a96afe46d..040e076edb 100644 --- a/tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml +++ b/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 diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index b8cda3c11a..9b4f27a424 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/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 diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 9311b966ee..3c334ca31f 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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 ) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index cde043d1cb..f92de74ed9 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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),