[Feature][MRv2] pcp/dcp public logic modification for MRv2 (vllm-project#12595)

weiguihua2 · web-flow · commit b2f683ca35a5 · 2026-07-30T09:39:34.000+08:00
### What this PR does / why we need it? Modifications to the public logic of PCP DCP in Model Runner v2. 1. Removed assertion error reporting. 2. Added a new subclass, pcp_manager, to construct the Ascend PCP manager. 3. In Model Runner, the pcp_manager is called to split a request into two separate requests: a header and a tail. ### Does this PR introduce _any_ user-facing change? No ### How was this patch tested? - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@d02df74 --------- Signed-off-by: weiguihua2 <weiguihua2@huawei.com>
diff --git a/tests/ut/worker/test_pcp_manager_v2.py b/tests/ut/worker/test_pcp_manager_v2.py
@@ -0,0 +1,252 @@
+# Adapt from https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/gpu/model_runner.py
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# This file is a part of the vllm-ascend project.
+#
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.worker.gpu.input_batch import InputBatch
+
+from vllm_ascend.utils import vllm_version_is
+
+if not vllm_version_is("0.25.1"):
+    from vllm.v1.worker.gpu.pcp_manager import PCPManager
+
+    import vllm_ascend.worker.v2.pcp_manager as pcp_manager_module
+    from vllm_ascend.worker.v2.input_batch import AscendInputBatch, AscendInputBuffers
+    from vllm_ascend.worker.v2.pcp_manager import (
+        AscendPCPManager,
+        maybe_build_ascend_pcp_manager,
+    )
+
+
+def _mock_async_copy_to_cpu(value, out=None, device=None):
+    """Copy PCP metadata without requiring device hooks in CPU-only UTs."""
+    if isinstance(value, np.ndarray):
+        value = torch.from_numpy(value)
+    elif not isinstance(value, torch.Tensor):
+        value = torch.as_tensor(value)
+
+    if out is not None:
+        out.copy_(value)
+        return out
+
+    return value.to(device="cpu")
+
+
+def _make_local_pcp_batch():
+    """Build a local batch in the shape returned by the community PCP manager."""
+    input_buffers = AscendInputBuffers(
+        max_num_reqs=4,
+        max_num_tokens=16,
+        device=torch.device("cpu"),
+    )
+    base_batch = InputBatch.make_dummy(
+        num_reqs=2,
+        num_tokens=6,
+        input_buffers=input_buffers,
+    )
+
+    # Local PCP rows: one starts at position 6 and contains two tokens; the
+    # other starts at position 13 and contains four tokens.
+    base_batch.req_ids = ["req-head", "req-tail"]
+    base_batch.idx_mapping = torch.tensor([3, 7], dtype=torch.int32)
+    base_batch.idx_mapping_np = np.array([3, 7], dtype=np.int32)
+    base_batch.expanded_idx_mapping = base_batch.idx_mapping
+    base_batch.num_scheduled_tokens = np.array([2, 4], dtype=np.int32)
+    base_batch.query_start_loc_np = np.array([0, 2, 6], dtype=np.int32)
+    base_batch.query_start_loc.copy_(torch.tensor([0, 2, 6], dtype=torch.int32))
+    base_batch.num_computed_tokens_np = np.array([6, 13], dtype=np.int32)
+    base_batch.prefill_len_np = np.array([32, 32], dtype=np.int32)
+    base_batch.num_computed_prefill_tokens_np = np.array([6, 13], dtype=np.int32)
+    base_batch.is_prefilling_np = np.array([True, True])
+    base_batch.seq_lens.copy_(torch.tensor([8, 17], dtype=torch.int32))
+    base_batch.seq_lens_cpu_upper_bound = torch.tensor([500, 600], dtype=torch.int32)
+    base_batch.input_ids.copy_(torch.tensor([10, 11, 20, 21, 22, 23], dtype=torch.int32))
+    base_batch.positions.copy_(torch.tensor([6, 7, 13, 14, 15, 16], dtype=torch.int64))
+    base_batch.is_padding.fill_(False)
+
+    return AscendInputBatch(
+        **base_batch.__dict__,
+        seq_lens_np=np.array([101, 102], dtype=np.int32),
+        attn_state="global-attn-state",
+    )
+
+
+def _make_global_pcp_batch():
+    """Build the global batch that is passed into PCPManager.partition_batch."""
+    input_buffers = AscendInputBuffers(
+        max_num_reqs=4,
+        max_num_tokens=32,
+        device=torch.device("cpu"),
+    )
+    base_batch = InputBatch.make_dummy(
+        num_reqs=1,
+        num_tokens=18,
+        input_buffers=input_buffers,
+    )
+    base_batch.req_ids = ["global-req"]
+    base_batch.idx_mapping = torch.tensor([3], dtype=torch.int32)
+    base_batch.idx_mapping_np = np.array([3], dtype=np.int32)
+    base_batch.expanded_idx_mapping = base_batch.idx_mapping
+    base_batch.num_scheduled_tokens = np.array([18], dtype=np.int32)
+    base_batch.query_start_loc_np = np.array([0, 18], dtype=np.int32)
+    base_batch.query_start_loc.copy_(torch.tensor([0, 18], dtype=torch.int32))
+    base_batch.num_computed_tokens_np = np.array([0], dtype=np.int32)
+    base_batch.prefill_len_np = np.array([18], dtype=np.int32)
+    base_batch.num_computed_prefill_tokens_np = np.array([0], dtype=np.int32)
+    base_batch.is_prefilling_np = np.array([True])
+    base_batch.seq_lens.copy_(torch.tensor([18], dtype=torch.int32))
+    base_batch.seq_lens_cpu_upper_bound = torch.tensor([18], dtype=torch.int32)
+    base_batch.input_ids.copy_(torch.arange(18, dtype=torch.int32))
+    base_batch.positions.copy_(torch.arange(18, dtype=torch.int64))
+    base_batch.is_padding.fill_(False)
+
+    return AscendInputBatch(
+        **base_batch.__dict__,
+        seq_lens_np=np.array([18], dtype=np.int32),
+        attn_state="global-attn-state",
+    )
+
+
+@pytest.mark.skipif(vllm_version_is("0.25.1"), reason="requires vllm main branch")
+def test_partition_batch_refreshes_local_ascend_input_batch_metadata():
+    """Refresh Ascend metadata after the real PCP local-batch rewrite."""
+    vllm_config = object()
+    global_batch = _make_global_pcp_batch()
+    req_states = SimpleNamespace(
+        last_sampled_tokens=torch.zeros(4, dtype=torch.int64),
+        prefill_len=SimpleNamespace(gpu=torch.zeros(4, dtype=torch.int32)),
+        draft_tokens=torch.empty((4, 0), dtype=torch.int64),
+    )
+    manager = AscendPCPManager(
+        pcp_world_size=2,
+        pcp_rank=0,
+        device=torch.device("cpu"),
+        vllm_config=vllm_config,
+        req_states=req_states,
+        max_num_reqs=1,
+        max_num_tokens=18,
+    )
+    attn_state = MagicMock()
+
+    with (
+        # This Triton helper is unrelated to PCP partitioning and has no CPU
+        # implementation. Stub only it; AscendPCPManager.partition_batch and
+        # PCPManager.partition_batch both execute unmocked below.
+        patch(
+            "vllm.v1.worker.gpu.pcp_manager.prepare_pos_seq_lens",
+            return_value=None,
+        ),
+        patch(
+            "vllm.v1.worker.gpu.pcp_manager.combine_sampled_and_draft_tokens",
+            return_value=torch.zeros(2, dtype=torch.int64),
+        ),
+        patch(
+            "vllm.v1.worker.gpu.pcp_manager.async_copy_to_gpu",
+            side_effect=_mock_async_copy_to_cpu,
+        ),
+        patch.object(pcp_manager_module, "build_attn_state", return_value=attn_state) as build_attn_state,
+    ):
+        result = manager.partition_batch(global_batch)
+
+    assert isinstance(result, AscendInputBatch)
+    assert result is not global_batch
+    assert manager._global_batch is global_batch
+    np.testing.assert_array_equal(global_batch.seq_lens_np, np.array([18], dtype=np.int32))
+    assert global_batch.attn_state == "global-attn-state"
+
+    # PCP=2 rank 0 owns the tail chunk then the head chunk; the real base
+    # implementation produces this local row order and pads to rank 1's size.
+    assert result.req_ids == ["global-req", "global-req"]
+    np.testing.assert_array_equal(result.idx_mapping_np, np.array([3, 3], dtype=np.int32))
+    np.testing.assert_array_equal(result.num_scheduled_tokens, np.array([3, 5], dtype=np.int32))
+    np.testing.assert_array_equal(result.query_start_loc_np, np.array([0, 3, 8], dtype=np.int32))
+    assert result.num_tokens == 8
+    assert result.num_tokens_after_padding == 10
+    assert torch.equal(result.input_ids[:8], torch.tensor([15, 16, 17, 0, 1, 2, 3, 4], dtype=torch.int32))
+
+    # dataclasses.replace() retains the global Ascend-only fields by default;
+    # the override must refresh them from real PCP-local CPU rows.
+    expected_seq_lens = np.array([18, 5], dtype=np.int32)
+    np.testing.assert_array_equal(result.seq_lens_np, expected_seq_lens)
+    assert result.attn_state is attn_state
+
+    args = build_attn_state.call_args.args
+    assert args[0] is vllm_config
+    np.testing.assert_array_equal(args[1], expected_seq_lens)
+    assert args[2] == 2
+    np.testing.assert_array_equal(args[3], np.array([3, 5], dtype=np.int32))
+    np.testing.assert_array_equal(args[4], np.array([3, 5], dtype=np.int32))
+
+
+@pytest.mark.skipif(vllm_version_is("0.25.1"), reason="requires vllm main branch")
+def test_maybe_build_ascend_pcp_manager_returns_none_when_pcp_is_disabled():
+    vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(prefill_context_parallel_size=1),
+    )
+
+    assert (
+        maybe_build_ascend_pcp_manager(
+            vllm_config,
+            torch.device("cpu"),
+            supports_mm_inputs=False,
+            req_states=MagicMock(),
+            block_tables=MagicMock(),
+        )
+        is None
+    )
+
+
+@pytest.mark.skipif(vllm_version_is("0.25.1"), reason="requires vllm main branch")
+def test_maybe_build_ascend_pcp_manager_uses_ascend_subclass():
+    vllm_config = SimpleNamespace(
+        parallel_config=SimpleNamespace(
+            prefill_context_parallel_size=2,
+            decode_context_parallel_size=2,
+            cp_kv_cache_interleave_size=4,
+        ),
+        scheduler_config=SimpleNamespace(max_num_seqs=8, max_num_batched_tokens=32),
+    )
+    pcp_group = SimpleNamespace(rank_in_group=1)
+    dcp_group = SimpleNamespace(rank_in_group=0)
+    req_states = MagicMock()
+
+    with (
+        patch.object(PCPManager, "validate_config") as validate_config,
+        patch.object(pcp_manager_module, "get_pcp_group", return_value=pcp_group),
+        patch.object(pcp_manager_module, "get_dcp_group", return_value=dcp_group),
+    ):
+        manager = maybe_build_ascend_pcp_manager(
+            vllm_config,
+            torch.device("cpu"),
+            supports_mm_inputs=False,
+            req_states=req_states,
+            block_tables=None,
+        )
+
+    assert isinstance(manager, AscendPCPManager)
+    assert manager.vllm_config is vllm_config
+    assert manager.pcp_world_size == 2
+    assert manager.pcp_rank == 1
+    assert manager.dcp_world_size == 2
+    assert manager.dcp_rank == 0
+    assert manager.cp_interleave == 4
+    validate_config.assert_called_once_with(vllm_config, False)
diff --git a/vllm_ascend/worker/v2/block_table.py b/vllm_ascend/worker/v2/block_table.py
@@ -68,27 +68,29 @@ def compute_slot_mappings(
         query_start_loc: torch.Tensor,
         positions: torch.Tensor,
         num_tokens_padded: int,
+        out: torch.Tensor | None = None,
     ) -> torch.Tensor:
         num_reqs = idx_mapping.shape[0]
         num_groups = self.num_kv_cache_groups
+        slot_mappings = self.slot_mappings if out is None else out
         _compute_slot_mappings_kernel[(num_groups, num_reqs + 1)](
-            self.max_num_batched_tokens,
+            slot_mappings.shape[1],
             idx_mapping,
             query_start_loc,
             positions,
             self.block_table_ptrs,
             self.block_table_strides,
             self.block_sizes_tensor,
-            self.slot_mappings,
-            self.slot_mappings.stride(0),
+            slot_mappings,
+            slot_mappings.stride(0),
             self.cp_rank,
             CP_SIZE=self.cp_size,
             CP_INTERLEAVE=self.cp_interleave,
             PAD_ID=PAD_SLOT_ID,
             TRITON_BLOCK_SIZE=1024,  # type: ignore
             TOTAL_BLOCK_SIZE=4096,
         )
-        return self.slot_mappings[:, :num_tokens_padded]
+        return slot_mappings[:, :num_tokens_padded]
 
 
 @triton.jit
diff --git a/vllm_ascend/worker/v2/model_runner.py b/vllm_ascend/worker/v2/model_runner.py
@@ -41,6 +41,8 @@
 if not vllm_version_is("0.25.1"):
     from vllm.v1.worker.gpu.model_runner import sort_batch_req_ids
 
+    from vllm_ascend.worker.v2.pcp_manager import maybe_build_ascend_pcp_manager
+
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import (
     MoECommType,
@@ -67,12 +69,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         # Ascend-specific configurations
         self.ascend_config = get_ascend_config()
         # The following features are not yet supported in Ascend NPU model runner v2:
-        # - Context parallelism (prefill or decode)
         # - Dynamic EPLB
-        parallel_config = vllm_config.parallel_config
-        if parallel_config.prefill_context_parallel_size > 1 or parallel_config.decode_context_parallel_size > 1:
-            raise NotImplementedError("Context parallelism is not supported by Ascend NPU model runner v2.")
-
         if self.ascend_config.eplb_config.dynamic_eplb:
             raise NotImplementedError("dynamic_eplb is not supported by Ascend NPU model runner v2.")
 
@@ -136,6 +133,17 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
         with graph_manager_wrapper(self):
             super().initialize_kv_cache(kv_cache_config)
 
+            # GPUModelRunner constructs the community PCP manager while initializing
+            # the KV cache. Replace it with the Ascend subclass.
+            if not vllm_version_is("0.25.1"):
+                self.pcp_manager = maybe_build_ascend_pcp_manager(
+                    self.vllm_config,
+                    self.device,
+                    self.supports_mm_inputs,
+                    self.req_states,
+                    self.block_tables,
+                )
+
     @torch.inference_mode()
     def profile_run(self) -> None:
         """Override GPUModelRunner.profile_run for Ascend NPUs.
@@ -361,6 +369,9 @@ def prepare_inputs(
             attn_state=attn_state,
         )
 
+        if not vllm_version_is("0.25.1"):
+            input_batch = vllm_model_runner.pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch)
+
         # For mla/sfa, update cos/sin. Here is for execute_model.
         update_cos_sin(input_batch.positions)
 
diff --git a/vllm_ascend/worker/v2/pcp_manager.py b/vllm_ascend/worker/v2/pcp_manager.py
