[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>

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)

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
     )

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,
     ):

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

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,
     )