From 676650397ff71cd802e7e0e30c2e46a215ab86ee Mon Sep 17 00:00:00 2001 From: Qiu Chunshuo Date: Tue, 15 Sep 2026 16:27:46 +0800 Subject: [PATCH] [Feature][PCP] Support decode-only FULL CUDA graphs (#53867) Signed-off-by: QiuChunshuo Signed-off-by: QiuChunshuo Co-authored-by: Codex Co-authored-by: Lucas Wilkinson Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/v1/worker/test_gpu_pcp_manager.py | 159 +++++++++++++++++++++++- vllm/v1/worker/gpu/cudagraph_utils.py | 9 +- vllm/v1/worker/gpu/input_batch.py | 8 +- vllm/v1/worker/gpu/model_runner.py | 35 +++--- vllm/v1/worker/gpu/pcp_manager.py | 114 ++++++++++++++--- 5 files changed, 283 insertions(+), 42 deletions(-) diff --git a/tests/v1/worker/test_gpu_pcp_manager.py b/tests/v1/worker/test_gpu_pcp_manager.py index dc935929d4..d3ec562aa7 100644 --- a/tests/v1/worker/test_gpu_pcp_manager.py +++ b/tests/v1/worker/test_gpu_pcp_manager.py @@ -1,15 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock import numpy as np import pytest import torch +from vllm.config import CUDAGraphMode from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens from vllm.v1.worker.gpu import cp_utils as gpu_cp_utils from vllm.v1.worker.gpu import pcp_manager as pcp_manager_module -from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers, set_dummy_context from vllm.v1.worker.gpu.pcp_manager import PCPManager @@ -20,6 +24,58 @@ def _copy_to_cpu(value, out=None, device=None): return tensor +def _make_config(cudagraph_mode: CUDAGraphMode): + return SimpleNamespace( + parallel_config=SimpleNamespace( + prefill_context_parallel_size=2, + decode_context_parallel_size=1, + pipeline_parallel_size=1, + dcp_comm_backend="ag_rs", + ), + model_config=SimpleNamespace( + use_mla=True, + is_encoder_decoder=False, + hf_text_config=SimpleNamespace(), + ), + lora_config=None, + speculative_config=None, + compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode), + ) + + +def _make_capture_manager(block_table: torch.Tensor): + block_tables = SimpleNamespace( + input_block_tables=(block_table,), + num_kv_cache_groups=1, + kernel_block_sizes=(2,), + blocks_per_kv_block=(1,), + ) + manager = PCPManager( + pcp_world_size=2, + pcp_rank=0, + device=torch.device("cpu"), + max_num_reqs=4, + max_num_tokens=8, + block_tables=block_tables, + ) + return manager, block_tables + + +@pytest.mark.parametrize( + "cudagraph_mode", + [CUDAGraphMode.FULL_DECODE_ONLY, CUDAGraphMode.FULL_AND_PIECEWISE], +) +def test_validate_config_accepts_decode_only_full_graphs(cudagraph_mode): + PCPManager.validate_config(_make_config(cudagraph_mode), supports_mm_inputs=False) + + +def test_validate_config_rejects_full_graph_for_prefills(): + with pytest.raises(NotImplementedError, match="decode-only routines"): + PCPManager.validate_config( + _make_config(CUDAGraphMode.FULL), supports_mm_inputs=False + ) + + def test_replicated_decode_piecewise_graph_padding(monkeypatch): manager = PCPManager( pcp_world_size=2, @@ -114,6 +170,89 @@ def test_graph_padding_cannot_be_smaller_than_largest_pcp_rank(monkeypatch): ) +@pytest.mark.parametrize( + ("cg_mode", "num_reqs", "expected_tokens", "expected_reqs"), + [ + (CUDAGraphMode.NONE, 4, None, None), + (CUDAGraphMode.PIECEWISE, None, 8, None), + (CUDAGraphMode.FULL, 4, 8, 4), + ], +) +def test_partition_padding_is_derived_from_batch_descriptor( + cg_mode, num_reqs, expected_tokens, expected_reqs +): + manager = MagicMock() + input_batch = MagicMock() + manager.partition_batch.return_value = input_batch + batch_desc = BatchExecutionDescriptor( + cg_mode=cg_mode, + num_tokens=8, + num_reqs=num_reqs, + ) + + result = pcp_manager_module.maybe_partition_pcp_batch( + manager, + input_batch, + batch_desc, + ) + + assert result is input_batch + manager.partition_batch.assert_called_once_with( + input_batch, + padded_num_tokens=expected_tokens, + padded_num_reqs=expected_reqs, + ) + + +def test_capture_uses_pcp_persistent_inputs(): + manager, _ = _make_capture_manager(torch.ones((4, 2), dtype=torch.int32)) + + dummy_batch = InputBatch.make_dummy( + num_reqs=4, + num_tokens=4, + input_buffers=InputBuffers(4, 8, torch.device("cpu")), + max_query_len=1, + ) + dummy_batch.input_ids.copy_(torch.arange(4)) + dummy_batch.positions.fill_(7) + input_batch = manager.prepare_inputs_to_capture(dummy_batch) + + assert input_batch is not dummy_batch + assert input_batch.req_ids == dummy_batch.req_ids + for name in ("input_ids", "positions", "is_padding", "query_start_loc", "seq_lens"): + actual = getattr(input_batch, name) + torch.testing.assert_close(actual, getattr(dummy_batch, name)) + assert actual.data_ptr() == getattr(manager.input_buffers, name).data_ptr() + + +def test_dummy_context_updates_pcp_local_block_tables(): + global_block_table = torch.full((4, 4), -1, dtype=torch.int32) + manager, block_tables = _make_capture_manager(global_block_table) + input_batch = InputBatch.make_dummy( + num_reqs=2, + num_tokens=2, + input_buffers=manager.input_buffers, + max_query_len=1, + ) + input_batch = manager.prepare_inputs_to_capture(input_batch) + local_block_tables = manager.get_dummy_block_tables(input_batch.num_reqs) + + set_dummy_context( + input_batch, + block_tables, + context_len=3, + num_kv_blocks=16, + max_model_len=16, + input_block_tables=local_block_tables, + ) + + torch.testing.assert_close( + local_block_tables[0][:2, :2], + torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + ) + assert torch.all(global_block_table == -1) + + def _rank_rows( pcp_rank: int, pcp_world_size: int, @@ -331,10 +470,22 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): global_batch.dcp_local_seq_lens = global_buffers.dcp_local_seq_lens[:2] global_batch.dcp_local_seq_lens.fill_(-1) - local_batch = manager.partition_batch(global_batch, padded_num_tokens=2) + local_batch = manager.partition_batch( + global_batch, + padded_num_tokens=4, + padded_num_reqs=4, + ) assert local_batch.dcp_local_seq_lens is None - assert local_batch.seq_lens.tolist() == [17, 25] + assert local_batch.seq_lens.tolist() == [17, 25, 0, 0] + assert local_batch.dcp_local_seq_lens_cpu_upper_bound is not None + expected_dcp_upper_bound = get_dcp_local_seq_lens( + torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + ) + assert torch.equal( + local_batch.dcp_local_seq_lens_cpu_upper_bound, + torch.cat([expected_dcp_upper_bound, torch.zeros(2, dtype=torch.int32)]), + ) # What execute_model does next: derive DCP metadata from the final batch # on the PCP-owned buffers. @@ -348,7 +499,7 @@ def test_partition_defers_dcp_metadata_to_post_partition_batch(): num_reqs_padded=local_batch.num_reqs_after_padding, ) expected = get_dcp_local_seq_lens( - torch.tensor([17, 25], dtype=torch.int32), 2, 0, 1 + torch.tensor([17, 25, 0, 0], dtype=torch.int32), 2, 0, 1 ) assert local_batch.dcp_local_seq_lens is not None assert torch.equal(local_batch.dcp_local_seq_lens.cpu(), expected) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index d197dbbc6d..9657c28fe6 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -777,11 +777,12 @@ def prepare_inputs_to_capture( 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_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) + input_batch = pcp_manager.prepare_inputs_to_capture(input_batch) + + block_table_provider = pcp_manager or block_tables + input_block_tables = block_table_provider.get_dummy_block_tables(num_reqs) + slot_mappings = block_table_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/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 5cbd642d25..d3c65bd06f 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Sequence from dataclasses import dataclass from typing import TYPE_CHECKING @@ -217,9 +218,12 @@ def set_dummy_context( context_len: int, num_kv_blocks: int, max_model_len: int, + input_block_tables: Sequence[torch.Tensor] | None = None, ) -> None: """Give each dummy request context_len of context, used when profiling step cost.""" - if not block_tables.input_block_tables: + if input_block_tables is None: + input_block_tables = block_tables.input_block_tables + if not input_block_tables: # Attention-free models have no KV context to fabricate. return num_reqs = input_batch.num_reqs @@ -241,7 +245,7 @@ def set_dummy_context( seq_len = context_len + query_len for block_table, block_size, bpk in zip( - block_tables.input_block_tables, + input_block_tables, block_tables.kernel_block_sizes, block_tables.blocks_per_kv_block, ): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 761893303a..25395c3330 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1468,7 +1468,7 @@ class GPUModelRunner(LoRAModelRunnerMixin): input_batch = pcp.maybe_partition_pcp_batch( self.pcp_manager, input_batch, - padded_num_tokens=batch_desc.num_tokens, + batch_desc, ) return input_batch @@ -1496,7 +1496,8 @@ class GPUModelRunner(LoRAModelRunnerMixin): def prepare_dummy_attn( self, input_batch: InputBatch, valid_state_slots: bool = False ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: - block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs) + block_table_provider = self.pcp_manager or self.block_tables + block_tables = block_table_provider.get_dummy_block_tables(input_batch.num_reqs) if valid_state_slots: state_slots = torch.arange( 1, @@ -1745,25 +1746,29 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.input_buffers, max_query_len=batch_desc.max_query_len, ) - if not skip_attn_for_dummy_run: - block_tables, slot_mappings = self.prepare_dummy_attn( - input_batch, valid_dummy_state_slots - ) - if context_len: - set_dummy_context( - input_batch, - self.block_tables, - context_len, - self.kv_cache_config.num_blocks, - self.max_model_len, - ) - else: + if self.pcp_manager is not None: + input_batch = self.pcp_manager.prepare_inputs_to_capture(input_batch) + if skip_attn_for_dummy_run: assert batch_desc.cg_mode != CUDAGraphMode.FULL, ( "Attention metadata must be prepared for dummy runs when using " "FULL cudagraph mode." ) block_tables = None slot_mappings = None + else: + block_tables, slot_mappings = self.prepare_dummy_attn( + input_batch, valid_dummy_state_slots + ) + if not skip_attn_for_dummy_run and context_len: + assert block_tables is not None + set_dummy_context( + input_batch, + self.block_tables, + context_len, + self.kv_cache_config.num_blocks, + self.max_model_len, + input_block_tables=block_tables, + ) attn_metadata = None slot_mappings_by_layer = None diff --git a/vllm/v1/worker/gpu/pcp_manager.py b/vllm/v1/worker/gpu/pcp_manager.py index 3a51532f57..adfbf2f9f0 100644 --- a/vllm/v1/worker/gpu/pcp_manager.py +++ b/vllm/v1/worker/gpu/pcp_manager.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Iterator from dataclasses import dataclass, replace +from typing import TYPE_CHECKING import numpy as np import torch @@ -17,6 +18,9 @@ from vllm.v1.worker.gpu.input_batch import ( InputBuffers, ) +if TYPE_CHECKING: + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + logger = init_logger(__name__) @@ -169,8 +173,14 @@ class PCPManager: "MRV2 sparse MLA PCP does not support CUDA graphs yet. " "Set -cc.cudagraph_mode=NONE." ) - if cudagraph_mode.has_full_cudagraphs(): - raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.") + if ( + cudagraph_mode.has_full_cudagraphs() + and not cudagraph_mode.separate_routine() + ): + raise NotImplementedError( + "MRV2 PCP supports full CUDA graphs for decode-only routines. " + "Use FULL_DECODE_ONLY, FULL_AND_PIECEWISE, PIECEWISE, or NONE." + ) if ( parallel_config.decode_context_parallel_size > 1 and parallel_config.dcp_comm_backend != "ag_rs" @@ -375,6 +385,22 @@ class PCPManager: for rank in range(self.pcp_world_size) ) + @staticmethod + def _resolve_num_reqs_after_padding( + input_batch: InputBatch, + padded_num_reqs: int | None, + num_local_reqs: int, + ) -> int: + if padded_num_reqs is None: + return num_local_reqs + if input_batch.has_prefill: + raise RuntimeError("PCP FULL graphs require a decode-only batch.") + assert padded_num_reqs >= num_local_reqs, ( + "PCP graph request capacity must cover the rank-local batch: " + f"{padded_num_reqs} < {num_local_reqs}." + ) + return padded_num_reqs + @property def input_buffers(self) -> InputBuffers: assert self._input_buffers is not None @@ -384,6 +410,7 @@ class PCPManager: self, input_batch: InputBatch, padded_num_tokens: int | None = None, + padded_num_reqs: int | None = None, ) -> InputBatch: assert self._input_buffers is not None input_buffers = self._input_buffers @@ -414,10 +441,15 @@ class PCPManager: ] num_local_reqs = len(local_segments) - if num_local_reqs > input_buffers.max_num_reqs: + num_reqs_after_padding = self._resolve_num_reqs_after_padding( + global_batch, + padded_num_reqs, + num_local_reqs, + ) + if num_reqs_after_padding > input_buffers.max_num_reqs: raise RuntimeError( - "PCP local request count exceeds the MRV2 input buffer size: " - f"{num_local_reqs} > {input_buffers.max_num_reqs}." + "PCP padded local request count exceeds the MRV2 input buffer size: " + f"{num_reqs_after_padding} > {input_buffers.max_num_reqs}." ) local_to_global_batch_req_idx_np = np.fromiter( @@ -505,21 +537,25 @@ class PCPManager: np.cumsum(local_num_scheduled_tokens, out=local_query_start_loc_out) local_query_start_loc_np[num_local_reqs + 1 :] = num_local_tokens async_copy_to_gpu(local_query_start_loc_np, out=input_buffers.query_start_loc) - local_query_start_loc = input_buffers.query_start_loc[: num_local_reqs + 1] + local_query_start_loc = input_buffers.query_start_loc[ + : num_reqs_after_padding + 1 + ] local_to_global_req_idx = async_copy_to_gpu( local_to_global_req_idx_np, device=self.device ) - seq_lens = input_buffers.seq_lens[:num_local_reqs] + seq_lens = input_buffers.seq_lens[:num_reqs_after_padding] + real_seq_lens = seq_lens[:num_local_reqs] if num_local_tokens > 0: local_end_positions = torch.index_select( input_buffers.positions, 0, - local_query_start_loc[1:] - 1, + local_query_start_loc[1 : num_local_reqs + 1] - 1, ) - seq_lens.copy_(local_end_positions + 1) + real_seq_lens.copy_(local_end_positions + 1) else: - seq_lens.zero_() + real_seq_lens.zero_() + seq_lens[num_local_reqs:].zero_() is_padding = input_buffers.is_padding[:num_local_tokens_padded] is_padding[:num_local_tokens].fill_(False) is_padding[num_local_tokens:].fill_(True) @@ -552,11 +588,15 @@ class PCPManager: local_num_computed_prefill_tokens_np = np.minimum( local_start_pos_np, local_prefill_len_np ) - local_is_prefilling_np = ( + real_local_is_prefilling_np = ( local_num_computed_prefill_tokens_np < local_prefill_len_np ) - seq_lens_cpu_upper_bound_np = np.zeros(num_local_reqs, dtype=np.int32) - seq_lens_cpu_upper_bound_np[:] = local_start_pos_np + local_num_scheduled_tokens + local_is_prefilling_np = np.zeros(num_reqs_after_padding, dtype=np.bool_) + local_is_prefilling_np[:num_local_reqs] = real_local_is_prefilling_np + seq_lens_cpu_upper_bound_np = np.zeros(num_reqs_after_padding, dtype=np.int32) + seq_lens_cpu_upper_bound_np[:num_local_reqs] = ( + local_start_pos_np + local_num_scheduled_tokens + ) dcp_local_seq_lens_cpu_upper_bound = None if self.dcp_world_size > 1: # The largest DCP shard of each row's whole request, identical on @@ -564,18 +604,24 @@ class PCPManager: request_seq_lens = (num_computed_tokens + num_scheduled_tokens)[ local_to_global_batch_req_idx_np ] - dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( + real_dcp_local_seq_lens_cpu_upper_bound = get_dcp_local_seq_lens( torch.from_numpy(request_seq_lens.astype(np.int32)), self.dcp_world_size, 0, self.cp_interleave, ) + dcp_local_seq_lens_cpu_upper_bound = torch.zeros( + num_reqs_after_padding, dtype=torch.int32 + ) + dcp_local_seq_lens_cpu_upper_bound[:num_local_reqs].copy_( + real_dcp_local_seq_lens_cpu_upper_bound + ) self._local_batch = replace( input_batch, req_ids=local_req_ids, num_reqs=num_local_reqs, - num_reqs_after_padding=num_local_reqs, + num_reqs_after_padding=num_reqs_after_padding, idx_mapping=local_to_global_req_idx, idx_mapping_np=local_to_global_req_idx_np, expanded_idx_mapping=local_to_global_req_idx, @@ -588,7 +634,7 @@ class PCPManager: num_draft_tokens=0, num_draft_tokens_per_req=None, query_start_loc=local_query_start_loc, - query_start_loc_np=local_query_start_loc_np[: num_local_reqs + 1], + query_start_loc_np=local_query_start_loc_np[: num_reqs_after_padding + 1], seq_lens=seq_lens, seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np), dcp_local_seq_lens=None, @@ -608,6 +654,31 @@ class PCPManager: ) return self._local_batch + def prepare_inputs_to_capture(self, input_batch: InputBatch) -> InputBatch: + """Stage a capture or dummy batch in persistent PCP input buffers.""" + input_buffers = self.input_buffers + num_reqs = input_batch.num_reqs_after_padding + num_tokens = input_batch.num_tokens_after_padding + input_batch = replace( + input_batch, + input_ids=input_buffers.input_ids[:num_tokens].copy_(input_batch.input_ids), + positions=input_buffers.positions[:num_tokens].copy_(input_batch.positions), + is_padding=input_buffers.is_padding[:num_tokens].copy_( + input_batch.is_padding + ), + query_start_loc=input_buffers.query_start_loc[: num_reqs + 1].copy_( + input_batch.query_start_loc + ), + seq_lens=input_buffers.seq_lens[:num_reqs].copy_(input_batch.seq_lens), + ) + return input_batch + + def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]: + assert self._local_block_tables is not None + return tuple( + block_table[:num_reqs].zero_() for block_table in self._local_block_tables + ) + def prepare_attn( self, input_batch: InputBatch ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]: @@ -728,13 +799,22 @@ class PCPManager: def maybe_partition_pcp_batch( manager: PCPManager | None, input_batch: InputBatch, - padded_num_tokens: int | None = None, + batch_desc: "BatchExecutionDescriptor", ) -> InputBatch: if manager is None: return input_batch + + padded_num_tokens = None + padded_num_reqs = None + if batch_desc.cg_mode != CUDAGraphMode.NONE: + padded_num_tokens = batch_desc.num_tokens + if batch_desc.cg_mode == CUDAGraphMode.FULL: + padded_num_reqs = batch_desc.num_reqs + return manager.partition_batch( input_batch, padded_num_tokens=padded_num_tokens, + padded_num_reqs=padded_num_reqs, )