[Core] Add MRV2 virtual-batch PCP for MLA (#46570)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

.buildkite/test_areas/lm_eval.yaml

@@ -78,6 +78,28 @@ steps:
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small-tp.txt
 
+- label: LM Eval PCP (4xB200)
+  key: lm-eval-pcp-4xb200
+  timeout_in_minutes: 360
+  device: b200-k8s
+  num_devices: 4
+  optional: true
+  source_file_dependencies:
+  - csrc/
+  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+  - tests/evals/gsm8k/configs/GLM-5.2-NVFP4-TP1-PCP4-EP.yaml
+  - tests/evals/gsm8k/configs/models-pcp.txt
+  - vllm/model_executor/layers/quantization
+  - vllm/config/parallel.py
+  - vllm/distributed/parallel_state.py
+  - vllm/model_executor/layers/attention/mla_attention.py
+  - vllm/model_executor/layers/attention/pcp.py
+  - vllm/v1/worker/gpu/model_runner.py
+  - vllm/v1/worker/gpu/pcp_manager.py
+  autorun_on_main: true
+  commands:
+  - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-pcp.txt
+
 - label: LM Eval Large Models EP (2xB200)
   key: lm-eval-large-models-ep-2xb200
   timeout_in_minutes: 60

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

@@ -0,0 +1,16 @@
+model_name: "nvidia/GLM-5.2-NVFP4"
+accuracy_threshold: 0.90
+num_questions: 1319
+num_fewshot: 5
+max_concurrency: 100
+server_args: >-
+  --enforce-eager
+  --max-model-len 4096
+  --safetensors-load-strategy prefetch
+  --moe-backend flashinfer_cutlass
+  --prefill-context-parallel-size 4
+  --enable-expert-parallel
+  --kv-cache-dtype fp8
+env:
+  VLLM_LOGGING_LEVEL: "DEBUG"
+  VLLM_USE_V2_MODEL_RUNNER: "1"

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

@@ -0,0 +1,17 @@
+model_name: "nvidia/GLM-5.2-NVFP4"
+accuracy_threshold: 0.90
+num_questions: 1319
+num_fewshot: 5
+max_concurrency: 100
+server_args: >-
+  --enforce-eager
+  --max-model-len 4096
+  --safetensors-load-strategy prefetch
+  --moe-backend flashinfer_cutlass
+  --tensor-parallel-size 2
+  --prefill-context-parallel-size 2
+  --enable-expert-parallel
+  --kv-cache-dtype fp8
+env:
+  VLLM_LOGGING_LEVEL: "DEBUG"
+  VLLM_USE_V2_MODEL_RUNNER: "1"

tests/evals/gsm8k/configs/models-pcp.txt

@@ -0,0 +1,2 @@
+GLM-5.2-NVFP4-TP2-PCP2-EP.yaml
+GLM-5.2-NVFP4-TP1-PCP4-EP.yaml

tests/evals/gsm8k/gsm8k_eval.py

@@ -217,6 +217,7 @@ def evaluate_gsm8k(
     seed: int | None = 42,
     request_timeout_seconds: float = 600,
     gen_prefix: str = "",
+    max_concurrency: int | None = None,
 ) -> dict[str, float | int]:
     """
     Evaluate GSM8K accuracy using vLLM serve endpoint.
@@ -261,7 +262,14 @@ def evaluate_gsm8k(
             return answer, tokens
 
         timeout = aiohttp.ClientTimeout(total=request_timeout_seconds)
-        async with aiohttp.ClientSession(timeout=timeout) as session:
+        connector = (
+            aiohttp.TCPConnector(limit=max_concurrency)
+            if max_concurrency is not None
+            else None
+        )
+        async with aiohttp.ClientSession(
+            timeout=timeout, connector=connector
+        ) as session:
             tasks = [get_answer(session, i) for i in range(num_questions)]
             await tqdm.gather(*tasks, desc="Evaluating")
 
@@ -343,6 +351,11 @@ def main() -> None:
     parser.add_argument(
         "--seed", type=int, default=42, help="Random seed for reproducibility"
     )
+    parser.add_argument(
+        "--max-concurrency",
+        type=int,
+        help="Maximum number of concurrent requests",
+    )
     parser.add_argument("--save-results", type=str, help="Save results to JSON file")
 
     args = parser.parse_args()
@@ -355,6 +368,7 @@ def main() -> None:
         port=args.port,
         temperature=args.temperature,
         seed=args.seed,
+        max_concurrency=args.max_concurrency,
     )
 
     # Print results to terminal

tests/evals/gsm8k/test_gsm8k_correctness.py

@@ -59,6 +59,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
         seed=eval_config.get("seed", 42),
         request_timeout_seconds=request_timeout_seconds,
         gen_prefix=eval_config.get("gen_prefix", ""),
+        max_concurrency=eval_config.get("max_concurrency"),
     )
 
     return results

tests/v1/attention/test_attention_splitting.py

@@ -160,6 +160,8 @@ def apply_split_decodes_and_prefills(
     decode_threshold: int,
     require_uniform: bool,
     padded_num_tokens: int | None = None,
+    is_prefilling: list[bool] | None = None,
+    treat_short_extends_as_decodes: bool = True,
 ):
     """Helper function to apply split_decodes_and_prefills and return
     the results."""
@@ -173,11 +175,14 @@ def apply_split_decodes_and_prefills(
 
     if padded_num_tokens is not None:
         common_metadata.num_actual_tokens = padded_num_tokens
+    if is_prefilling is not None:
+        common_metadata.is_prefilling = torch.tensor(is_prefilling)
 
     return split_decodes_and_prefills(
         common_metadata,
         decode_threshold=decode_threshold,
         require_uniform=require_uniform,
+        treat_short_extends_as_decodes=treat_short_extends_as_decodes,
     )
 
 
@@ -236,6 +241,17 @@ def test_split_decodes_and_prefills_uniform_all_ones():
     assert num_prefill_tokens == 0
 
 
+def test_split_decodes_and_prefills_uniform_short_extend():
+    result = apply_split_decodes_and_prefills(
+        [1, 1],
+        decode_threshold=1,
+        require_uniform=True,
+        is_prefilling=[False, True],
+        treat_short_extends_as_decodes=False,
+    )
+    assert result == (1, 1, 1, 1)
+
+
 def test_split_decodes_and_prefills_uniform_all_short_decodes():
     query_lens = [2, 2, 1, 3, 2, 1, 2]
     num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

tests/v1/attention/test_sparse_mla_backends.py

@@ -1114,6 +1114,13 @@ def test_sparse_backend_prefill_correctness(
 @pytest.mark.parametrize(
     "seq_lens,query_lens,workspace_size,max_logits_bytes,expected",
     [
+        (
+            torch.tensor([0]),
+            torch.tensor([0]),
+            100,
+            1000,
+            [],
+        ),
         # Logits constraint triggers split (M*N exceeds budget)
         # req0: M=10, N=100 -> 1000 elems (4000 bytes) - fits in 5000
         # req1: adding M=10, N=100 -> new_M=20, new_N=200 -> 4000 elems > 1250

tests/v1/kv_offload/test_factory.py

@@ -385,7 +385,7 @@ def test_tiering_spec_aligns_row_size():
     assert spec.num_blocks == cpu_bytes_to_use // alignment
 
 
-def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():
+def test_offloading_spec_kv_sharding_ignores_prefill_context_parallel():
     config = _make_layout_vllm_config(
         cpu_bytes_to_use=65536,
         extra_config={"block_size": 64},
@@ -394,9 +394,9 @@ def test_offloading_spec_resolves_prefill_context_parallel_block_sizes():
 
     spec = _create_spec(config, _make_kv_cache_config())
 
-    assert spec.tokens_per_block == (32,)
-    assert spec.tokens_per_hash == 32
-    assert spec.blocks_per_chunk == 2
+    assert spec.tokens_per_block == (16,)
+    assert spec.tokens_per_hash == 16
+    assert spec.blocks_per_chunk == 4
 
 
 def test_offloading_config_preserves_data_parallel_index():

tests/v1/simple_kv_offload/test_scheduler.py

@@ -1551,9 +1551,8 @@ def _make_cp_scheduler(
     num_gpu_blocks: int = 16,
     lazy: bool = False,
 ) -> SchedulerFixture:
-    """Build a SimpleCPUOffloadScheduler with CP-scaled virtual block size."""
-    cp_world_size = dcp_world_size * pcp_world_size
-    virtual_block_size = BLOCK_SIZE * cp_world_size
+    """Build a SimpleCPUOffloadScheduler with DCP-scaled block size."""
+    virtual_block_size = BLOCK_SIZE * dcp_world_size
 
     kv_cache_config = _make_kv_cache_config(num_gpu_blocks)
     vllm_config = _make_cp_vllm_config(dcp_world_size, pcp_world_size)
@@ -1645,17 +1644,14 @@ def _allocate_cp_gpu_blocks(
     ],
 )
 def test_cp_block_size_scaling(dcp_world_size: int, pcp_world_size: int) -> None:
-    """Verify that the scheduler's block_size and cp_world_size are correctly
-    scaled when context parallelism is enabled."""
+    """Verify block size scaling follows DCP ownership."""
     fix = _make_cp_scheduler(
         dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
     )
     sched = fix.scheduler
 
-    expected_cp = dcp_world_size * pcp_world_size
-    assert sched.cp_world_size == expected_cp
-    assert sched.block_size == BLOCK_SIZE * expected_cp
-    assert sched.fa_block_size == BLOCK_SIZE * expected_cp
+    assert sched.cp_world_size == dcp_world_size
+    assert sched.block_size == BLOCK_SIZE * dcp_world_size
 
 
 # ---------------------------------------------------------------------------
@@ -1682,8 +1678,7 @@ def test_cp_eager_store_and_load_roundtrip(
         lazy=False,
     )
     sched = fix.scheduler
-    cp = dcp_world_size * pcp_world_size
-    vbs = BLOCK_SIZE * cp
+    vbs = BLOCK_SIZE * dcp_world_size
 
     num_blocks = 2
     req = _make_cp_request(num_blocks, vbs)
@@ -1749,18 +1744,16 @@ def test_cp_eager_store_and_load_roundtrip(
 def test_cp_effective_block_size_store_and_load(
     dcp_world_size: int, pcp_world_size: int
 ) -> None:
-    """Verify ready_blocks_g (store) and n_take_g (load) use the effective
-    (physical * cp) block size, not the per-rank physical size."""
+    """Verify store/load physical block counts scale with DCP, not PCP."""
     fix = _make_cp_scheduler(
         dcp_world_size=dcp_world_size, pcp_world_size=pcp_world_size
     )
     sched = fix.scheduler
     gpu_pool = fix.gpu_block_pool
-    cp = dcp_world_size * pcp_world_size
-    vbs = BLOCK_SIZE * cp
+    vbs = BLOCK_SIZE * dcp_world_size
+    expected_blocks = 1
 
-    # Store: allocate 2 blocks, confirm only 1. Without the fix,
-    # ready_blocks_g = vbs / BLOCK_SIZE = 2, storing both blocks.
+    # Store one DCP-scaled logical block worth of tokens.
     req = _make_cp_request(num_blocks=2, virtual_block_size=vbs)
     gpu_blocks = _allocate_cp_gpu_blocks(gpu_pool, req, 2, vbs)
     kv = KVCacheBlocks(blocks=(gpu_blocks,))
@@ -1772,12 +1765,11 @@ def test_cp_effective_block_size_store_and_load(
             new_reqs={req.request_id: kv.get_block_ids()},
         )
     )
-    assert len(m1.store_gpu_blocks) == 1
-    assert len(m1.store_cpu_blocks) == 1
+    assert len(m1.store_gpu_blocks) == expected_blocks
+    assert len(m1.store_cpu_blocks) == expected_blocks
     simulate_store_completion(sched, m1.store_event)
 
-    # Load: store 2 blocks from a second request, accept only 1 as external.
-    # Without the fix, n_take_g = vbs / BLOCK_SIZE = 2, loading both.
+    # Load one DCP-scaled logical block from a two-block CPU hit.
     req2 = _make_cp_request(num_blocks=2, virtual_block_size=vbs)
     kv2 = KVCacheBlocks(blocks=(_allocate_cp_gpu_blocks(gpu_pool, req2, 2, vbs),))
     req2.num_computed_tokens = 2 * vbs
@@ -1801,7 +1793,7 @@ def test_cp_effective_block_size_store_and_load(
     hit, _ = sched.get_num_new_matched_tokens(req3, num_computed_tokens=0)
     assert hit == 2 * vbs
 
-    kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(2),))
+    kv3 = KVCacheBlocks(blocks=(gpu_pool.get_new_blocks(expected_blocks),))
     sched.update_state_after_alloc(req3, kv3, num_external_tokens=vbs)
     m3 = sched.build_connector_meta(
         make_scheduler_output(
@@ -1810,8 +1802,8 @@ def test_cp_effective_block_size_store_and_load(
         )
     )
     assert m3.load_event >= 0
-    assert len(m3.load_gpu_blocks) == 1
-    assert len(m3.load_cpu_blocks) == 1
+    assert len(m3.load_gpu_blocks) == expected_blocks
+    assert len(m3.load_cpu_blocks) == expected_blocks
     assert m3.load_gpu_blocks == [kv3.get_block_ids()[0][0]]
 
 

vllm/config/parallel.py

@@ -122,10 +122,11 @@ class ParallelConfig:
     tensor_parallel_size: int = Field(default=1, ge=1)
     """Number of tensor parallel groups."""
     prefill_context_parallel_size: int = Field(default=1, ge=1)
-    """Number of prefill context parallel groups."""
+    """Number of ranks that split prefill sequence computation. PCP expands
+    the process world size but does not increase the KV-cache shard count."""
     data_parallel_size: int = Field(default=1, ge=1)
     """Number of data parallel groups. MoE layers will be sharded according to
-    the product of the tensor parallel size and data parallel size."""
+    the product of the tensor, prefill-context, and data parallel sizes."""
     data_parallel_size_local: int = Field(default=1, ge=0)
     """Number of local data parallel groups. A value of 0 is a sentinel used by
     the engine-args layer to signal that data parallelism was specified
@@ -337,9 +338,9 @@ class ParallelConfig:
     connect as clients to exchange self-picked group ports at runtime."""
 
     decode_context_parallel_size: int = Field(default=1, ge=1)
-    """Number of decode context parallel groups, because the world size does
-    not change by dcp, it simply reuse the GPUs of TP group, and tp_size
-    needs to be divisible by dcp_size."""
+    """Number of ranks that shard the decode KV cache. DCP does not expand
+    the process world size. Without PCP, DCP reuses TP ranks. With PCP, DCP
+    either spans the PCP axis or the full TP x PCP block."""
 
     dcp_kv_cache_interleave_size: int = 1
     """
@@ -357,13 +358,11 @@ class ParallelConfig:
     """
 
     cp_kv_cache_interleave_size: int = 1
-    """Interleave size of kv_cache storage while using DCP or PCP.
-    For `total_cp_rank = pcp_rank * dcp_world_size + dcp_rank`,
-        and `total_cp_world_size = pcp_world_size * dcp_world_size`.
-    store interleave_size tokens on total_cp_rank i,
-    then store next interleave_size tokens on total_cp_rank i+1.
+    """Interleave size of kv_cache storage while using DCP.
+    Store interleave_size tokens on dcp_rank i, then store next
+    interleave_size tokens on dcp_rank i+1.
     Interleave_size=1: token-level alignment, where token `i` is stored on
-        total_cp_rank `i % total_cp_world_size`.
+        dcp_rank `i % dcp_world_size`.
     Interleave_size=block_size: block-level alignment, where tokens are
         first populated to the preceding ranks. Tokens are then stored
         in (rank i+1, block j) only after (rank i, block j) is fully occupied.
@@ -480,11 +479,19 @@ class ParallelConfig:
                 )
             if not self.enable_expert_parallel:
                 raise ValueError("enable_expert_parallel must be True to use EPLB.")
-            if self.tensor_parallel_size * self.data_parallel_size <= 1:
+            # The EP group spans the TP x PCP x DP ranks. EPLB therefore needs
+            # TP, PCP, or DP > 1.
+            if (
+                self.tensor_parallel_size
+                * self.prefill_context_parallel_size
+                * self.data_parallel_size
+                <= 1
+            ):
                 raise ValueError(
-                    "EPLB requires tensor_parallel_size or data_parallel_size "
-                    f"to be greater than 1, but got "
-                    f"TP={self.tensor_parallel_size},DP={self.data_parallel_size}."
+                    "EPLB requires tensor, prefill-context, or data parallelism, "
+                    f"but got TP={self.tensor_parallel_size}, "
+                    f"PCP={self.prefill_context_parallel_size}, "
+                    f"DP={self.data_parallel_size}."
                 )
         else:
             if self.eplb_config.num_redundant_experts != 0:
@@ -495,15 +502,21 @@ class ParallelConfig:
                     "num_redundant_experts."
                 )
 
-        # Note(hc): In the current implementation of decode context
-        # parallel(DCP), tp_size needs to be divisible by dcp_size,
-        # because the world size does not change by dcp, it simply
-        # reuses the GPUs of TP group, and split one TP group into
-        # tp_size//dcp_size DCP groups.
-        if self.tensor_parallel_size % self.decode_context_parallel_size != 0:
+        tp = self.tensor_parallel_size
+        pcp = self.prefill_context_parallel_size
+        dcp = self.decode_context_parallel_size
+        if pcp > 1 and self.data_parallel_size > 1:
+            raise ValueError("PCP does not support data parallelism yet.")
+        if pcp == 1:
+            # DCP reuses the TP ranks when PCP is disabled.
+            if tp % dcp != 0:
+                raise ValueError(f"tp_size={tp} must be divisible by dcp_size={dcp}.")
+        elif dcp not in (1, pcp, tp * pcp):
             raise ValueError(
-                f"tp_size={self.tensor_parallel_size} must be divisible by"
-                f"dcp_size={self.decode_context_parallel_size}."
+                "When PCP is enabled, DCP must be disabled, span the PCP "
+                "axis, or span the full TP x PCP axis. "
+                f"Got TP={tp}, PCP={pcp}, DCP={dcp}; valid DCP sizes are "
+                f"{sorted({1, pcp, tp * pcp})}."
             )
 
         if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1:
@@ -515,8 +528,7 @@ class ParallelConfig:
 
     @property
     def world_size_across_dp(self) -> int:
-        """world_size_across_dp is TPxPPxDP, it is the size of the world
-        including data parallelism."""
+        """Process world size across TP, PCP, PP, and DP."""
         return self.world_size * self.data_parallel_size
 
     @property

vllm/config/vllm.py

@@ -2130,9 +2130,10 @@ class VllmConfig:
         model_config = self.model_config
         speculative_config = self.speculative_config
 
-        if self.parallel_config.prefill_context_parallel_size > 1:
+        if self.parallel_config.prefill_context_parallel_size > 1 and not (
+            model_config is not None and model_config.use_mla
+        ):
             unsupported.append("prefill context parallelism")
-
         if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
             unsupported.append("stock torch.compile")
 

vllm/distributed/device_communicators/all2all.py

@@ -8,7 +8,7 @@ import torch
 import torch.distributed as dist
 
 import vllm.envs as envs
-from vllm.distributed import get_dp_group, get_ep_group
+from vllm.distributed import get_dp_group, get_ep_group, get_pcp_group
 from vllm.distributed.utils import StatelessProcessGroup
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
@@ -49,6 +49,23 @@ class AgRsAll2AllManager(All2AllManagerBase):
     def __init__(self, cpu_group, tcp_store_group=None):
         super().__init__(cpu_group, tcp_store_group)
 
+    def _get_comm_group(self, is_sequence_parallel: bool) -> Any:
+        if is_sequence_parallel:
+            return get_ep_group()
+        if self.dp_world_size > 1:
+            return get_dp_group()
+        return get_pcp_group()
+
+    def _get_sizes(self, num_local_tokens: int, comm_group: Any) -> list[int]:
+        if self.dp_world_size == 1:
+            return [num_local_tokens] * comm_group.world_size
+
+        dp_metadata = get_forward_context().dp_metadata
+        assert dp_metadata is not None
+        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
+        assert sizes is not None
+        return sizes
+
     def dispatch_router_logits(
         self,
         hidden_states: torch.Tensor,
@@ -62,11 +79,8 @@ class AgRsAll2AllManager(All2AllManagerBase):
         """
         Gather hidden_states and router_logits from all dp ranks.
         """
-        dp_metadata = get_forward_context().dp_metadata
-        assert dp_metadata is not None
-        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
-        assert sizes is not None
-        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
+        dist_group = self._get_comm_group(is_sequence_parallel)
+        sizes = self._get_sizes(hidden_states.shape[0], dist_group)
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
 
         tensors_to_gather = [hidden_states, router_logits]
@@ -97,11 +111,8 @@ class AgRsAll2AllManager(All2AllManagerBase):
         """
         Gather hidden_states and router_logits from all dp ranks.
         """
-        dp_metadata = get_forward_context().dp_metadata
-        assert dp_metadata is not None
-        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
-        assert sizes is not None
-        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
+        dist_group = self._get_comm_group(is_sequence_parallel)
+        sizes = self._get_sizes(hidden_states.shape[0], dist_group)
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
 
         tensors_to_gather = [hidden_states, topk_weights, topk_ids]
@@ -129,12 +140,11 @@ class AgRsAll2AllManager(All2AllManagerBase):
         """
         Reduce-scatter hidden_states across all dp ranks.
         """
-        dp_metadata = get_forward_context().dp_metadata
-        assert dp_metadata is not None
-        sizes = dp_metadata.get_chunk_sizes_across_dp_rank()
-        assert sizes is not None
-
-        dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
+        dist_group = self._get_comm_group(is_sequence_parallel)
+        sizes = self._get_sizes(
+            hidden_states.shape[0] // dist_group.world_size,
+            dist_group,
+        )
         hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
         return hidden_states
 

vllm/distributed/device_communicators/base_device_communicator.py

@@ -176,11 +176,16 @@ class DeviceCommunicatorBase:
         config = get_current_vllm_config_or_none()
         if config is not None:
             # initialize the all2all manager for DP or sequence-parallel EP.
+            parallel_config = config.parallel_config
             use_ep = (
-                config.parallel_config.data_parallel_size > 1
-                or config.parallel_config.use_sequence_parallel_moe
+                parallel_config.data_parallel_size > 1
+                or parallel_config.use_sequence_parallel_moe
+                or (
+                    parallel_config.enable_expert_parallel
+                    and parallel_config.prefill_context_parallel_size > 1
+                )
             )
-            all2all_backend = config.parallel_config.all2all_backend
+            all2all_backend = parallel_config.all2all_backend
 
         self.is_ep_communicator = unique_name.split(":")[0] == "ep"
         self.use_all2all = self.is_ep_communicator and use_ep

vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

@@ -36,13 +36,12 @@ def build_offloading_config(
     engine_id = kv_transfer_config.engine_id
 
     parallel_config = vllm_config.parallel_config
-    context_parallel_factor = (
-        parallel_config.decode_context_parallel_size
-        * parallel_config.prefill_context_parallel_size
-    )
     groups = tuple(
         OffloadingGroupConfig(
-            tokens_per_block=(group.kv_cache_spec.block_size * context_parallel_factor),
+            tokens_per_block=(
+                group.kv_cache_spec.block_size
+                * parallel_config.decode_context_parallel_size
+            ),
             layer_names=tuple(group.layer_names),
         )
         for group in kv_cache_config.kv_cache_groups

vllm/distributed/parallel_state.py

@@ -1784,7 +1784,7 @@ def initialize_model_parallel(
             get_world_group().device_group
         )
 
-    # the layout order is: ExternalDP x DP x PP x TP
+    # the layout order is: ExternalDP x DP x PP x PCP x TP
     # ExternalDP is the data parallel group that is not part of the model,
     # every dp rank can generate independently (in verl integration).
     # DP is the data parallel group that is part of the model,
@@ -1821,17 +1821,13 @@ def initialize_model_parallel(
     # Build the DCP model-parallel groups.
     global _DCP
     assert _DCP is None, "decode context model parallel group is already initialized"
-    # Note(hc): In the current implementation of decode context parallel,
-    # dcp_size must not exceed tp_size, because the world size does not
-    # change by DCP, it simply reuses the GPUs of TP group, and split one
-    # TP group into tp_size//dcp_size DCP groups.
-    group_ranks = all_ranks.reshape(-1, decode_context_model_parallel_size).unbind(0)
+    dcp_size = decode_context_model_parallel_size or 1
+    dcp_ranks = local_all_ranks if enable_elastic_ep else all_ranks
+    if dcp_size > 1:
+        # DCP spans PCP first, then TP for full TP x PCP groups.
+        dcp_ranks = dcp_ranks.transpose(-1, -2)
+    group_ranks = dcp_ranks.reshape(-1, dcp_size).unbind(0)
     group_ranks = [x.tolist() for x in group_ranks]
-    if enable_elastic_ep:
-        group_ranks = local_all_ranks.reshape(
-            -1, decode_context_model_parallel_size
-        ).unbind(0)
-        group_ranks = [x.tolist() for x in group_ranks]
     _DCP = init_model_parallel_group(
         group_ranks,
         get_world_group().local_rank,
@@ -2005,6 +2001,13 @@ def ensure_model_parallel_initialized(
         f"{pcp_world_size=} vs. "
         f"{prefill_context_model_parallel_size=}"
     )
+    dcp_world_size = get_dcp_group().world_size
+    dcp_model_parallel_size = decode_context_model_parallel_size or 1
+    assert dcp_world_size == dcp_model_parallel_size, (
+        "decode context parallel group already initialized, but of unexpected size: "
+        f"{dcp_world_size=} vs. "
+        f"{dcp_model_parallel_size=}"
+    )
 
 
 def prepare_communication_buffer_for_model(model: torch.nn.Module):

vllm/model_executor/layers/attention/mla_attention.py

@@ -211,6 +211,7 @@ from vllm.config import (
 from vllm.config.cache import CacheDType
 from vllm.distributed.parallel_state import (
     get_dcp_group,
+    get_tp_group,
     is_global_first_rank,
 )
 from vllm.forward_context import ForwardContext, get_forward_context
@@ -225,6 +226,10 @@ from vllm.model_executor.layers.attention.attention import (
 from vllm.model_executor.layers.attention.kv_transfer_utils import (
     maybe_transfer_kv_layer,
 )
+from vllm.model_executor.layers.attention.pcp import (
+    finalize_mla_pcp_decode,
+    maybe_gather_mla_latent_cache_inputs,
+)
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -269,7 +274,7 @@ from vllm.v1.attention.backends.utils import (
     get_dcp_local_seq_lens,
     split_decodes_and_prefills,
 )
-from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
+from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs
 from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 from vllm.v1.attention.selector import get_attn_backend
@@ -499,6 +504,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         self.use_direct_call = not current_platform.opaque_attention_op()
 
         vllm_config = get_current_vllm_config()
+        parallel_config = vllm_config.parallel_config
+        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
         compilation_config = vllm_config.compilation_config
         if prefix in compilation_config.static_forward_context:
             raise ValueError(f"Duplicate layer name: {prefix}")
@@ -611,11 +618,23 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             assert isinstance(slot_mapping, dict), (
                 f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
             )
+            layer_slot_mapping = slot_mapping.get(self.layer_name)
+            kv_for_cache, kpe_for_cache, layer_slot_mapping = (
+                maybe_gather_mla_latent_cache_inputs(
+                    kv_c_normed,
+                    k_pe,
+                    layer_slot_mapping,
+                    attn_metadata.num_decode_tokens
+                    if attn_metadata is not None
+                    else None,
+                    self.use_pcp,
+                )
+            )
             self.impl.do_kv_cache_update(  # type: ignore[attr-defined]
-                kv_c_normed,
-                k_pe,
+                kv_for_cache,
+                kpe_for_cache,
                 self_kv_cache,
-                slot_mapping.get(self.layer_name),
+                layer_slot_mapping,
                 self.kv_cache_dtype,
                 self._k_scale,
             )
@@ -710,6 +729,10 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
 
         num_actual_toks = attn_metadata.num_actual_tokens
+        if self.use_pcp and self.impl.dcp_world_size > 1 and quant_key is not None:
+            raise NotImplementedError(
+                "MRV2 MLA PCP+DCP does not support fused output quantization yet."
+            )
 
         # Inputs and outputs may be padded for CUDA graphs
         output_padded = output
@@ -833,11 +856,15 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             else:
                 mqa_q = (mqa_ql_nope, mqa_q_pe)
             if self.impl.dcp_world_size > 1:
-                if isinstance(mqa_q, tuple):
-                    # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P)
-                    mqa_q = torch.cat(mqa_q, dim=-1)
-                # mqa_q do allgather in head dim.
-                mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
+                if self.use_pcp:
+                    if self.impl.dcp_world_size > self.impl.pcp_world_size:
+                        if isinstance(mqa_q, tuple):
+                            mqa_q = torch.cat(mqa_q, dim=-1)
+                        mqa_q = get_tp_group().all_gather(mqa_q, dim=1)
+                else:
+                    if isinstance(mqa_q, tuple):
+                        mqa_q = torch.cat(mqa_q, dim=-1)
+                    mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
 
             # call decode attn
             if not self.impl.is_sparse:
@@ -846,6 +873,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
 
             # correct dcp attn_out with lse.
             if self.impl.dcp_world_size > 1:
+                assert lse is not None
                 if self.dcp_a2a:
                     attn_out = dcp_a2a_lse_reduce(
                         attn_out,
@@ -853,6 +881,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                         get_dcp_group(),
                         is_lse_base_on_e=self.impl.lse_base_on_e,
                     )
+                elif self.use_pcp:
+                    attn_out = cp_lse_ag_out_ar(
+                        attn_out,
+                        lse,
+                        get_dcp_group(),
+                        is_lse_base_on_e=self.impl.lse_base_on_e,
+                    )
                 else:
                     attn_out = cp_lse_ag_out_rs(
                         attn_out,
@@ -860,6 +895,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                         get_dcp_group(),
                         is_lse_base_on_e=self.impl.lse_base_on_e,
                     )
+                if self.use_pcp:
+                    attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads)
 
             # v_up projection
             self._v_up_proj(attn_out, out=mqa_output_slice)
@@ -903,6 +940,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 raise ValueError(f"Unsupported quant_key: {quant_key}")
             return quant_output
 
+        if self.use_pcp and output_padded.shape[0] > num_actual_toks:
+            output_padded[num_actual_toks:].zero_()
         return output_padded
 
     def process_weights_after_loading(self, act_dtype: torch.dtype):
@@ -1082,8 +1121,17 @@ def unified_mla_kv_cache_update(
     the data dependency between them to ensure torch.compile preserves ordering.
     """
     layer_name = _resolve_layer_name(layer_name)
-    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
+    attn_metadata, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(
+        layer_name
+    )
     if layer_slot_mapping is not None:
+        kv_c_normed, k_pe, layer_slot_mapping = maybe_gather_mla_latent_cache_inputs(
+            kv_c_normed,
+            k_pe,
+            layer_slot_mapping,
+            attn_metadata.num_decode_tokens if attn_metadata is not None else None,
+            attn_layer.use_pcp,
+        )
         attn_layer.impl.do_kv_cache_update(  # type: ignore[attr-defined]
             kv_c_normed,
             k_pe,
@@ -1723,6 +1771,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
         self.compilation_config = vllm_config.compilation_config
         self.vllm_config = vllm_config
         self.device = device
+        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
 
         self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
         self.mla_dims = get_mla_dims(self.model_config)
@@ -1735,11 +1784,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
 
         try:
             self.dcp_world_size = get_dcp_group().world_size
-            self.dcp_rank = get_dcp_group().rank_in_group
         except AssertionError:
             # DCP might not be initialized in testing
             self.dcp_world_size = 1
-            self.dcp_rank = 0
         self.dcp_local_block_size = parallel_config.cp_kv_cache_interleave_size
         self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
         self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
@@ -1857,6 +1904,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                 common_attn_metadata,
                 decode_threshold=self.reorder_batch_threshold,
                 require_uniform=(self.query_len_support != QueryLenSupport.VARLEN),
+                treat_short_extends_as_decodes=not self.use_pcp,
             )
         )
 

vllm/model_executor/layers/attention/pcp.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.distributed.parallel_state import (
+    get_pcp_group,
+    get_tp_group,
+)
+
+
+def _gather_prefill_cache_inputs(
+    tensors: tuple[torch.Tensor, ...],
+    slot_mapping: torch.Tensor,
+    num_decode_tokens: int,
+) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+    """Keep replicated decode writes local and gather partitioned prefills."""
+    local_num_tokens = tensors[0].shape[0]
+    assert all(tensor.shape[0] == local_num_tokens for tensor in tensors)
+    assert 0 <= num_decode_tokens <= local_num_tokens
+
+    if num_decode_tokens == local_num_tokens:
+        return tensors, slot_mapping[:num_decode_tokens]
+
+    pcp_group = get_pcp_group()
+    gathered_prefills = tuple(
+        pcp_group.all_gather(tensor[num_decode_tokens:].contiguous(), dim=0)
+        for tensor in tensors
+    )
+    pcp_size = pcp_group.world_size
+    gathered_slot_mapping = slot_mapping[: pcp_size * local_num_tokens]
+    if num_decode_tokens == 0:
+        return gathered_prefills, gathered_slot_mapping
+
+    cache_inputs = tuple(
+        torch.cat((tensor[:num_decode_tokens], gathered_prefill), dim=0)
+        for tensor, gathered_prefill in zip(tensors, gathered_prefills)
+    )
+    rank_slot_mappings = gathered_slot_mapping.view(pcp_size, local_num_tokens)
+    cache_slot_mapping = torch.cat(
+        (
+            rank_slot_mappings[0, :num_decode_tokens],
+            rank_slot_mappings[:, num_decode_tokens:].flatten(),
+        )
+    )
+    return cache_inputs, cache_slot_mapping
+
+
+def maybe_gather_mla_latent_cache_inputs(
+    kv_c_normed: torch.Tensor,
+    k_pe: torch.Tensor,
+    slot_mapping: torch.Tensor | None,
+    num_decode_tokens: int | None,
+    use_pcp: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+    if not use_pcp or num_decode_tokens is None:
+        return kv_c_normed, k_pe, slot_mapping
+    assert slot_mapping is not None
+    num_tokens = kv_c_normed.shape[0]
+    k_pe_flat = k_pe.reshape(num_tokens, -1)
+    (cache_kv_c, cache_k_pe_flat), cache_slot_mapping = _gather_prefill_cache_inputs(
+        (kv_c_normed, k_pe_flat),
+        slot_mapping,
+        num_decode_tokens,
+    )
+    cache_k_pe = cache_k_pe_flat.view(-1, *k_pe.shape[1:])
+    return cache_kv_c, cache_k_pe, cache_slot_mapping
+
+
+def maybe_gather_indexer_k(
+    k: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    num_decode_tokens: int,
+    use_pcp: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not use_pcp:
+        return k, slot_mapping
+    (cache_k,), cache_slot_mapping = _gather_prefill_cache_inputs(
+        (k,), slot_mapping, num_decode_tokens
+    )
+    return cache_k, cache_slot_mapping
+
+
+def finalize_mla_pcp_decode(
+    output: torch.Tensor,
+    num_heads: int,
+) -> torch.Tensor:
+    if output.shape[1] < num_heads:
+        output = get_pcp_group().all_gather(output, dim=1)
+    elif output.shape[1] > num_heads:
+        head_start = get_tp_group().rank_in_group * num_heads
+        output = output[:, head_start : head_start + num_heads]
+    return output

vllm/model_executor/layers/attention/sparse_mla_attention.py

@@ -59,13 +59,12 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
             device=device,
         )
         parallel_config = vllm_config.parallel_config
+        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
         try:
             self.dcp_world_size = get_dcp_group().world_size
-            self.dcp_rank = get_dcp_group().rank_in_group
         except AssertionError:
             # DCP might not be initialized in testing
             self.dcp_world_size = 1
-            self.dcp_rank = 0
         self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
         self.dcp_local_block_size = self.cp_kv_cache_interleave_size
         self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
@@ -174,6 +173,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
         num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
             common_attn_metadata,
             decode_threshold=self.reorder_batch_threshold or 1,
+            treat_short_extends_as_decodes=not self.use_pcp,
             require_uniform=self.require_uniform_decodes,
         )
         (

vllm/model_executor/layers/fused_moe/config.py

@@ -1050,7 +1050,9 @@ class FusedMoEParallelConfig:
 
     @property
     def use_all2all_kernels(self):
-        return self.use_ep and (self.dp_size > 1 or self.is_sequence_parallel)
+        return self.use_ep and (
+            self.dp_size > 1 or self.pcp_size > 1 or self.is_sequence_parallel
+        )
 
     @property
     def use_deepep_ht_kernels(self):

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -752,18 +752,12 @@ class MoERunner(MoERunnerInterface):
             assert len(result) == 2
             hidden_states, router_logits = result
 
-        # NOTE: Similar with DP, PCP also needs dispatch and combine. For
-        # simplicity, AgRsAll2All was added separately for PCP here. Maybe
-        # we should modify All2AllManager abstraction to better support PCP.
-        if self.moe_config.pcp_size > 1:
-            hidden_states = get_pcp_group().all_gather(
-                hidden_states,
-                dim=0,
-            )
-            router_logits = get_pcp_group().all_gather(
-                router_logits,
-                dim=0,
-            )
+        if (
+            self.moe_config.pcp_size > 1
+            and not self.moe_config.moe_parallel_config.use_all2all_kernels
+        ):
+            hidden_states = get_pcp_group().all_gather(hidden_states, dim=0)
+            router_logits = get_pcp_group().all_gather(router_logits, dim=0)
 
         return hidden_states, router_logits
 
@@ -777,11 +771,11 @@ class MoERunner(MoERunnerInterface):
                 hidden_states, self.moe_config.is_sequence_parallel
             )
 
-        if self.moe_config.pcp_size > 1:
-            hidden_states = get_pcp_group().reduce_scatter(
-                hidden_states,
-                dim=0,
-            )
+        if (
+            self.moe_config.pcp_size > 1
+            and not self.moe_config.moe_parallel_config.use_all2all_kernels
+        ):
+            hidden_states = get_pcp_group().reduce_scatter(hidden_states, dim=0)
 
         if self.shared_experts is not None:
             assert shared_output is not None

vllm/model_executor/layers/sparse_attn_indexer.py

@@ -9,10 +9,11 @@ from vllm import _custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.compilation.breakable_cudagraph import eager_break_during_capture
 from vllm.config import get_current_vllm_config
-from vllm.distributed import get_dcp_group
+from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.attention.pcp import maybe_gather_indexer_k
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     get_fp8_min_max,
 )
@@ -308,6 +309,7 @@ def sparse_attn_indexer(
     total_seq_lens: int,
     topk_indices_buffer: torch.Tensor,
     skip_k_cache_insert: bool,
+    use_pcp: bool,
     use_fp4_cache: bool = False,
     dcp_rank: int = 0,
     dcp_world_size: int = 1,
@@ -354,6 +356,7 @@ def sparse_attn_indexer(
             total_seq_lens,
             topk_indices_buffer,
             skip_k_cache_insert,
+            use_pcp,
             use_fp4_cache,
         )
     attn_metadata_narrowed = attn_metadata[k_cache_prefix]
@@ -373,18 +376,28 @@ def sparse_attn_indexer(
     # During speculative decoding, k may be padded to the CUDA graph batch
     # size while slot_mapping only covers actual tokens. Truncate k to avoid
     # out-of-bounds reads in the kernel.
+    # Keep PCP padding so every rank contributes the same all-gather shape.
     num_tokens = slot_mapping.shape[0]
+    if use_pcp:
+        num_tokens //= get_pcp_group().world_size
     if k is not None:
         k = k[:num_tokens]
 
     if not skip_k_cache_insert:
+        assert k is not None
+        k, slot_mapping_for_cache = maybe_gather_indexer_k(
+            k,
+            slot_mapping,
+            num_decode_tokens,
+            use_pcp,
+        )
         # scale_fmt can be None, but the function expects str
         assert scale_fmt is not None
         assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet"
         ops.indexer_k_quant_and_cache(
             k,
             kv_cache,
-            slot_mapping,
+            slot_mapping_for_cache,
             quant_block_size,
             scale_fmt,
         )
@@ -498,7 +511,14 @@ def sparse_attn_indexer(
         assert decode_metadata is not None
         kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache)
         decode_lens = decode_metadata.decode_lens
-        if decode_metadata.requires_padding:
+        if num_decode_tokens == 0:
+            padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:])
+            padded_q_scale = (
+                q_scale[:1].reshape(1, 1, *q_scale.shape[1:])
+                if q_scale is not None
+                else None
+            )
+        elif decode_metadata.requires_padding:
             # pad in edge case where we have short chunked prefill length <
             # decode_threshold since we unstrictly split
             # prefill and decode by decode_threshold
@@ -663,6 +683,7 @@ def sparse_attn_indexer_fake(
     total_seq_lens: int,
     topk_indices_buffer: torch.Tensor | None,
     skip_k_cache_insert: bool,
+    use_pcp: bool,
     use_fp4_cache: bool = False,
     dcp_rank: int = 0,
     dcp_world_size: int = 1,
@@ -725,6 +746,7 @@ class SparseAttnIndexer(CustomOp):
         self.dcp_world_size = parallel_config.decode_context_parallel_size
         self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
         self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
+        self.use_pcp = parallel_config.prefill_context_parallel_size > 1
         if current_platform.is_cuda() and not has_deep_gemm():
             raise RuntimeError(
                 "Sparse Attention Indexer CUDA op requires DeepGEMM support in "
@@ -777,6 +799,7 @@ class SparseAttnIndexer(CustomOp):
             self.max_total_seq_len,
             self.topk_indices_buffer,
             self.skip_k_cache_insert,
+            self.use_pcp,
             self.use_fp4_cache,
             self.dcp_rank,
             self.dcp_world_size,

vllm/v1/attention/backend.py

@@ -281,6 +281,13 @@ class AttentionBackend(ABC):
     def supports_kv_connector(cls) -> bool:
         return True
 
+    @classmethod
+    def supports_pcp(cls) -> bool:
+        try:
+            return cls.get_impl_cls().supports_pcp
+        except NotImplementedError:
+            return False
+
     @classmethod
     def supports_attn_type(cls, attn_type: str) -> bool:
         """Check if backend supports a given attention type.
@@ -327,6 +334,7 @@ class AttentionBackend(ABC):
         use_non_causal: bool = False,
         use_batch_invariant: bool = False,
         use_kv_connector: bool = False,
+        use_pcp: bool = False,
     ) -> list[str]:
         invalid_reasons = []
         if not cls.supports_head_size(head_size):
@@ -367,6 +375,8 @@ class AttentionBackend(ABC):
             invalid_reasons.append("batch invariance not supported")
         if use_kv_connector and not cls.supports_kv_connector():
             invalid_reasons.append("KV connector not supported")
+        if use_pcp and not cls.supports_pcp():
+            invalid_reasons.append("PCP not supported")
         combination_reason = cls.supports_combination(
             head_size,
             dtype,
@@ -858,8 +868,8 @@ class AttentionImplBase(ABC, Generic[T]):
         except AssertionError:
             self.pcp_world_size = 1
             self.pcp_rank = 0
-        self.total_cp_world_size = self.pcp_world_size * self.dcp_world_size
-        self.total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
+        self.total_cp_world_size = self.dcp_world_size
+        self.total_cp_rank = self.dcp_rank
 
         self.need_to_return_lse_for_decode = (
             self.dcp_world_size > 1 and self.can_return_lse_for_decode
@@ -987,6 +997,8 @@ class AttentionImpl(AttentionImplBase[T], Generic[T]):
 class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
     """MLA attention implementation with forward_mqa and forward_mha methods."""
 
+    supports_pcp: bool = True
+
     @abstractmethod
     def __init__(
         self,

vllm/v1/attention/backends/mla/indexer.py

@@ -6,7 +6,7 @@ import torch
 
 import vllm.envs as envs
 from vllm.config import VllmConfig
-from vllm.distributed import get_dcp_group
+from vllm.distributed import get_dcp_group, get_pcp_group
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
@@ -29,7 +29,7 @@ from vllm.v1.attention.backends.utils import (
     split_decodes_and_prefills,
 )
 from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
-from vllm.v1.worker.cp_utils import get_total_cp_world_size
+from vllm.v1.worker.cp_utils import get_kv_cache_shard_count
 
 logger = init_logger(__name__)
 
@@ -109,7 +109,7 @@ def split_indexer_prefill_chunks(
             end += 1
 
         req_slice = slice(start + request_offset, end + request_offset)
-        max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m
+        max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m)
         for q_off in range(0, chunk_m, max_q):
             sub_m = min(max_q, chunk_m - q_off)
             chunks.append((req_slice, slice(q_off, q_off + sub_m)))
@@ -118,6 +118,10 @@ def split_indexer_prefill_chunks(
 
 
 class DeepseekV32IndexerBackend(AttentionBackend):
+    @classmethod
+    def supports_pcp(cls) -> bool:
+        return True
+
     @staticmethod
     def get_name() -> str:
         return "DEEPSEEK_V32_INDEXER"
@@ -256,6 +260,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         parallel_config = self.vllm_config.parallel_config
         self.dcp_world_size = parallel_config.decode_context_parallel_size
         self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0
+        self.pcp_world_size = parallel_config.prefill_context_parallel_size
+        self.use_pcp = self.pcp_world_size > 1
         self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
         # The DCP sparse-indexer code is parameterized by interleave size, but
         # interleave > 1 is not yet validated end-to-end (gsm8k parity fails),
@@ -339,7 +345,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         )
         max_num_blocks_per_req = cdiv(
             self.vllm_config.model_config.max_model_len,
-            self.kv_cache_spec.block_size * get_total_cp_world_size(),
+            self.kv_cache_spec.block_size * get_kv_cache_shard_count(),
         )
         self.expanded_block_table_buffer = torch.zeros(
             (
@@ -570,6 +576,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                 common_attn_metadata,
                 decode_threshold=self.decode_threshold,
                 require_uniform=not self.use_flattening,
+                treat_short_extends_as_decodes=not self.use_pcp,
             )
         )
 
@@ -579,6 +586,9 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         compressed_slot_mapping = slot_mapping
         compressed_seq_lens = seq_lens
         if self.compress_ratio > 1:
+            padded_num_tokens = num_tokens
+            if self.pcp_world_size > 1:
+                padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size
             compressed_slot_mapping = get_compressed_slot_mapping(
                 num_tokens,
                 query_start_loc,
@@ -588,6 +598,11 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
                 self.compress_ratio,
                 out=self.compressed_slot_mapping_buffer,
             )
+            if self.pcp_world_size > 1:
+                compressed_slot_mapping = get_pcp_group().all_gather(
+                    self.compressed_slot_mapping_buffer[:padded_num_tokens],
+                    dim=0,
+                )
             compressed_seq_lens = seq_lens // self.compress_ratio
 
         prefill_metadata = None

vllm/v1/attention/backends/utils.py

@@ -612,7 +612,9 @@ def split_decodes_and_prefills(
         # check if we are in a padded uniform batch; this is used for full-CGs, some
         # requests may have a query length of 0 but since they are padding its fine
         # to treat them as decodes (ensures num_decodes matches the captured size)
-        if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):
+        if treat_short_extends_as_decodes and torch.all(
+            (query_lens == query_lens[0]) | (query_lens == 0)
+        ):
             return num_reqs, 0, num_tokens, 0  # all decodes
         is_prefill = query_lens != query_lens[0]
     else:

vllm/v1/attention/selector.py

@@ -36,6 +36,7 @@ class AttentionSelectorConfig(NamedTuple):
     use_non_causal: bool = False
     use_batch_invariant: bool = False
     use_kv_connector: bool = False
+    use_pcp: bool = False
 
     def __repr__(self):
         return (
@@ -52,7 +53,8 @@ class AttentionSelectorConfig(NamedTuple):
             f"has_sliding_window={self.has_sliding_window}, "
             f"use_non_causal={self.use_non_causal}, "
             f"use_batch_invariant={self.use_batch_invariant}, "
-            f"use_kv_connector={self.use_kv_connector})"
+            f"use_kv_connector={self.use_kv_connector}, "
+            f"use_pcp={self.use_pcp})"
         )
 
 
@@ -150,6 +152,7 @@ def get_attn_backend(
         use_non_causal=vllm_config.attention_config.use_non_causal,
         use_batch_invariant=envs.VLLM_BATCH_INVARIANT,
         use_kv_connector=use_kv_connector,
+        use_pcp=vllm_config.parallel_config.prefill_context_parallel_size > 1,
     )
 
     # A per-KV-group override (keyed by KVCacheSpecKind) takes precedence over

vllm/v1/core/kv_cache_coordinator.py

@@ -472,8 +472,6 @@ class UnitaryKVCacheCoordinator(KVCacheCoordinator):
         self.pcp_world_size = pcp_world_size
         if dcp_world_size > 1:
             self.block_size *= dcp_world_size
-        if pcp_world_size > 1:
-            self.block_size *= pcp_world_size
         # For models using only Mamba, block_size is set to max_model_len when
         # prefix caching is disabled, and hash_block_size validation is skipped.
         assert not enable_caching or (hash_block_size == self.block_size), (

vllm/v1/core/kv_cache_utils.py

@@ -631,8 +631,8 @@ def resolve_kv_cache_block_sizes(
 
     - ``scheduler_block_size`` is the token-alignment invariant used by the
       scheduler (e.g. for ``num_computed_tokens`` rounding). Single group:
-      ``cache_config.block_size * dcp * pcp``. Multiple groups: LCM of every
-      group's effective block size. Attention groups are scaled by DCP/PCP;
+      ``cache_config.block_size * dcp``. Multiple groups: LCM of every
+      group's effective block size. Attention groups are scaled by DCP;
       Mamba groups keep their full per-rank state and are not scaled.
     - ``hash_block_size`` is the granularity at which ``Request.block_hashes``
       is computed. Single group: equals scheduler block size. Multiple groups:
@@ -644,15 +644,14 @@ def resolve_kv_cache_block_sizes(
     """
     cache_config = vllm_config.cache_config
     dcp = vllm_config.parallel_config.decode_context_parallel_size
-    pcp = vllm_config.parallel_config.prefill_context_parallel_size
     groups = kv_cache_config.kv_cache_groups
 
-    if len(groups) <= 1:  # Single group: block_size * dcp * pcp
-        bs = cache_config.block_size * dcp * pcp
+    if len(groups) <= 1:
+        bs = cache_config.block_size * dcp
         return bs, bs
 
     group_block_sizes = [
-        g.kv_cache_spec.block_size * dcp * pcp
+        g.kv_cache_spec.block_size * dcp
         if isinstance(g.kv_cache_spec, AttentionSpec)
         else g.kv_cache_spec.block_size
         for g in groups

vllm/v1/core/sched/scheduler.py

@@ -271,7 +271,7 @@ class Scheduler(SchedulerInterface):
             log_stats=self.log_stats,
             enable_kv_cache_events=self.enable_kv_cache_events,
             dcp_world_size=self.dcp_world_size,
-            pcp_world_size=self.pcp_world_size,
+            pcp_world_size=1,
             scheduler_block_size=self.block_size,
             hash_block_size=hash_block_size,
             metrics_collector=self.kv_metrics_collector,

vllm/v1/core/single_type_kv_cache_manager.py

@@ -75,8 +75,8 @@ class SingleTypeKVCacheManager(ABC):
         self.block_size = kv_cache_spec.block_size
         self.dcp_world_size = dcp_world_size
         self.pcp_world_size = pcp_world_size
-        if dcp_world_size * pcp_world_size > 1:
-            self.block_size *= dcp_world_size * pcp_world_size
+        if dcp_world_size > 1:
+            self.block_size *= dcp_world_size
         self.kv_cache_spec = kv_cache_spec
         self.block_pool = block_pool
         self.enable_caching = enable_caching
@@ -674,10 +674,10 @@ class FullAttentionManager(SingleTypeKVCacheManager):
             "and chunked local attention groups"
         )
         block_size = kv_cache_spec.block_size
-        if dcp_world_size * pcp_world_size > 1:
-            # DCP/PCP shard each block's KV across ranks; hashes must be
-            # viewed at the sharded (scaled) block size.
-            block_size *= dcp_world_size * pcp_world_size
+        if dcp_world_size > 1:
+            # DCP shards each block's KV across ranks; hashes must be viewed at
+            # the sharded block size.
+            block_size *= dcp_world_size
         block_hashes = resolve_block_hashes(
             block_hashes,
             block_pool.hash_block_size,

vllm/v1/kv_cache_interface.py

@@ -218,14 +218,9 @@ class AttentionSpec(KVCacheSpec):
         )
 
     def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int:
-        # Attention KV is token-interleaved across DCP/PCP ranks, so each rank
-        # only stores max_len // (dcp * pcp) tokens per request.
         parallel_config = vllm_config.parallel_config
-        total_cp_size = (
-            parallel_config.decode_context_parallel_size
-            * parallel_config.prefill_context_parallel_size
-        )
-        return cdiv(max_len, self.block_size * total_cp_size)
+        kv_shard_count = parallel_config.decode_context_parallel_size
+        return cdiv(max_len, self.block_size * kv_shard_count)
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -263,11 +258,8 @@ class FullAttentionSpec(AttentionSpec):
     def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
         max_model_len = vllm_config.model_config.max_model_len
         dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
-        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
-        # Note(hc): each dcp rank only need save
-        # (max_model_len//dcp_world_size) tokens locally.
-        if dcp_world_size * pcp_world_size > 1:
-            max_model_len = cdiv(max_model_len, dcp_world_size * pcp_world_size)
+        if dcp_world_size > 1:
+            max_model_len = cdiv(max_model_len, dcp_world_size)
         return cdiv(max_model_len, self.block_size) * self.page_size_bytes
 
     @classmethod

vllm/v1/simple_kv_offload/manager.py

@@ -83,8 +83,7 @@ class SimpleCPUOffloadScheduler:
             and vllm_config.kv_events_config.enable_kv_cache_events
         )
         dcp_world_size = vllm_config.parallel_config.decode_context_parallel_size
-        pcp_world_size = vllm_config.parallel_config.prefill_context_parallel_size
-        self.cp_world_size = dcp_world_size * pcp_world_size
+        self.cp_world_size = dcp_world_size
         self.block_size = scheduler_block_size
         self.hash_block_size = hash_block_size
         assert self.block_size % self.hash_block_size == 0
@@ -127,7 +126,7 @@ class SimpleCPUOffloadScheduler:
             enable_caching=True,
             enable_kv_cache_events=self.enable_kv_cache_events,
             dcp_world_size=dcp_world_size,
-            pcp_world_size=pcp_world_size,
+            pcp_world_size=1,
             scheduler_block_size=self.block_size,
             hash_block_size=self.hash_block_size,
         )
@@ -354,7 +353,7 @@ class SimpleCPUOffloadScheduler:
             )
             assert num_external_tokens % g_block_size == 0, (
                 f"num_external_tokens={num_external_tokens} not aligned to "
-                f"group {g} effective block_size={g_block_size}"
+                f"group {g} block_size={g_block_size}"
             )
             n_take_g = num_external_tokens // g_block_size
             cpu_hit_blocks.append(cpu_hit_blocks_full[g][:n_take_g])

vllm/v1/worker/block_table.py

@@ -163,8 +163,6 @@ class BlockTable:
             return
         assert self.slot_mapping_mode == SlotMappingMode.TOKEN_TO_KV_SLOT
 
-        total_cp_world_size = self.pcp_world_size * self.dcp_world_size
-        total_cp_rank = self.pcp_rank * self.dcp_world_size + self.dcp_rank
         _compute_slot_mapping_kernel[(num_reqs + 1,)](
             num_tokens,
             self.max_num_batched_tokens,
@@ -176,8 +174,8 @@ class BlockTable:
             self.slot_mapping.gpu,
             KV_CACHE_BLOCK_SIZE=self.kv_cache_block_size,
             BLOCKS_PER_KV_BLOCK=self.blocks_per_kv_block,
-            TOTAL_CP_WORLD_SIZE=total_cp_world_size,
-            TOTAL_CP_RANK=total_cp_rank,
+            TOTAL_CP_WORLD_SIZE=self.dcp_world_size,
+            TOTAL_CP_RANK=self.dcp_rank,
             CP_KV_CACHE_INTERLEAVE_SIZE=self.cp_kv_cache_interleave_size,
             PAD_ID=PAD_SLOT_ID,
             BLOCK_SIZE=1024,

vllm/v1/worker/cp_utils.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, cast
 import torch
 
 from vllm.config import VllmConfig, get_layers_from_vllm_config
-from vllm.distributed import get_dcp_group, get_pcp_group
+from vllm.distributed import get_dcp_group
 from vllm.logger import init_logger
 from vllm.v1.attention.backend import CommonAttentionMetadata
 from vllm.v1.attention.backends.utils import split_decodes_prefills_and_extends
@@ -27,6 +27,13 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
         layer_type = cast(type[Any], AttentionLayerBase)
         layers = get_layers_from_vllm_config(vllm_config, layer_type)
         for layer in layers.values():
+            get_attn_backend = getattr(layer, "get_attn_backend", None)
+            if pcp_size > 1 and get_attn_backend is not None:
+                backend = get_attn_backend()
+                assert backend.supports_pcp(), (
+                    "PCP requires attention backend support, "
+                    f"but {backend.get_name()} does not support PCP."
+                )
             layer_impl = getattr(layer, "impl", None)
             if layer_impl is None:
                 continue
@@ -44,26 +51,14 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
                     "--attention-backend or disable DCP."
                 )
 
-            if pcp_size > 1:
-                assert layer_impl.supports_pcp, (
-                    "PCP requires attention impls' support, "
-                    f"but the impl {layer_impl.__class__.__name__} "
-                    "does not support PCP."
-                )
 
-
-def get_total_cp_world_size():
-    try:
-        pcp_world_size = get_pcp_group().world_size
-    except AssertionError:
-        # PCP might not be initialized in testing
-        pcp_world_size = 1
+def get_kv_cache_shard_count() -> int:
     try:
         dcp_world_size = get_dcp_group().world_size
     except AssertionError:
         # DCP might not be initialized in testing
         dcp_world_size = 1
-    return dcp_world_size * pcp_world_size
+    return dcp_world_size
 
 
 def get_dcp_dummy_context_len(

vllm/v1/worker/gpu/attn_utils.py

@@ -578,6 +578,7 @@ def build_attn_metadata(
     seq_lens_cpu_upper_bound: torch.Tensor | None = None,
     dcp_local_seq_lens: torch.Tensor | None = None,
     positions: torch.Tensor | None = None,
+    is_prefilling: torch.Tensor | None = None,
     mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,
     model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
     for_cudagraph_capture: bool = False,
@@ -605,6 +606,11 @@ def build_attn_metadata(
             if model_specific_attn_metadata is not None
             else {}
         )
+        # Model-specific metadata (e.g. Mamba hybrid) may supply its own
+        # padding-aware is_prefilling, which takes precedence over the default.
+        group_is_prefilling = common_attn_metadata_extra_kwargs.pop(
+            "is_prefilling", is_prefilling
+        )
         common_attn_metadata = CommonAttentionMetadata(
             query_start_loc=query_start_loc_gpu,
             query_start_loc_cpu=query_start_loc_cpu,
@@ -619,6 +625,7 @@ def build_attn_metadata(
             causal=group_causal,
             dcp_local_seq_lens=dcp_local_seq_lens,
             positions=positions,
+            is_prefilling=group_is_prefilling,
             mm_req_doc_ranges=mm_req_doc_ranges,
             rswa_prefix_lens=rswa_prefix_lens,
             **common_attn_metadata_extra_kwargs,

vllm/v1/worker/gpu/block_table.py

@@ -135,20 +135,28 @@ class BlockTables:
         self,
         idx_mapping: torch.Tensor,
         num_reqs_padded: int,
+        out: tuple[torch.Tensor, ...] | None = None,
+        out_ptrs: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, ...]:
+        if out is None:
+            out = tuple(self.input_block_tables)
+            out_ptrs = self.input_block_table_ptrs
+        else:
+            assert out_ptrs is not None
+            assert len(out) == self.num_kv_cache_groups
         num_reqs = idx_mapping.shape[0]
         # Launch kernel with num_reqs_padded to fuse zeroing of padded rows.
         _gather_block_tables_kernel[(self.num_kv_cache_groups, num_reqs_padded)](
             idx_mapping,
             self.block_table_ptrs,
-            self.input_block_table_ptrs,
+            out_ptrs,
             self.block_table_strides,
             self.num_blocks.gpu,
             self.num_blocks.gpu.stride(0),
             num_reqs,
             BLOCK_SIZE=1024,  # type: ignore
         )
-        return tuple(bt[:num_reqs_padded] for bt in self.input_block_tables)
+        return tuple(bt[:num_reqs_padded] for bt in out)
 
     def get_dummy_block_tables(self, num_reqs: int) -> tuple[torch.Tensor, ...]:
         # NOTE(woosuk): The output may be used for CUDA graph capture.
@@ -163,26 +171,28 @@ class BlockTables:
         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
         )
-        return self.slot_mappings[:, :num_tokens_padded]
+        return slot_mappings[:, :num_tokens_padded]
 
     def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
         # Fill the entire slot_mappings tensor, not just the first `num_tokens` entries.

vllm/v1/worker/gpu/model_runner.py

@@ -52,6 +52,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
 from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
 from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
+from vllm.v1.worker.gpu import pcp_manager as pcp
 from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
 from vllm.v1.worker.gpu.attn_utils import (
     build_slot_mappings_by_layer,
@@ -205,6 +206,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         # Draft tokens propagation - for spec-dec + struct outputs.
         self.draft_tokens_handler = DraftTokensHandler(self.device)
 
+        self.pcp_manager: pcp.PCPManager | None = None
+
         # Pooling models.
         self.is_pooling_model = self.model_config.runner_type == "pooling"
         self.pooling_runner: PoolingRunner | None = None
@@ -456,6 +459,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             cp_rank=self.dcp_rank,
             cp_interleave=self.cp_interleave,
         )
+        self.pcp_manager = pcp.maybe_build_pcp_manager(
+            self.vllm_config,
+            self.device,
+            self.supports_mm_inputs,
+            self.req_states,
+            self.block_tables,
+        )
         initialize_mamba_ssu_backend(
             self.vllm_config.mamba_config, self.kv_cache_config
         )
@@ -1008,7 +1018,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             # prompt_lens is only used in R-SWA case.
             prompt_lens = self.req_states.prompt_len.gpu[idx_mapping]
 
-        return InputBatch(
+        input_batch = InputBatch(
             req_ids=req_ids,
             num_reqs=num_reqs,
             num_reqs_after_padding=num_reqs_padded,
@@ -1040,10 +1050,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             has_structured_output_reqs=scheduler_output.has_structured_output_requests,
             prompt_lens=prompt_lens,
         )
+        return pcp.maybe_partition_pcp_batch(self.pcp_manager, input_batch)
 
     def prepare_attn(
         self, input_batch: InputBatch
     ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        if self.pcp_manager is not None:
+            return self.pcp_manager.prepare_attn(input_batch)
+
         # Block tables: num_kv_cache_groups x [num_reqs_padded, max_num_blocks].
         block_tables = self.block_tables.gather_block_tables(
             input_batch.idx_mapping,
@@ -1063,8 +1077,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self, input_batch: InputBatch
     ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
         block_tables = self.block_tables.get_dummy_block_tables(input_batch.num_reqs)
-        slot_mappings = self.block_tables.get_dummy_slot_mappings(
-            input_batch.num_tokens
+        slot_mappings = pcp.maybe_get_pcp_dummy_slot_mappings(
+            self.pcp_manager, self.block_tables, input_batch.num_tokens
         )
         return block_tables, slot_mappings
 
@@ -1412,6 +1426,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             return ModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
 
         # Last rank: sample tokens
+        hidden_states, input_batch = pcp.maybe_restore_pcp_for_sampling(
+            self.pcp_manager, hidden_states, input_batch
+        )
+
         sampler_output, num_sampled, num_rejected = self.sample(
             hidden_states, input_batch, grammar_output
         )

vllm/v1/worker/gpu/model_states/default.py

@@ -151,7 +151,10 @@ class DefaultModelState(ModelState):
             # For piecewise cudagraphs and eager, use unpadded sizes.
             num_reqs = input_batch.num_reqs
             num_tokens = input_batch.num_tokens
-        query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
+        query_start_loc_cpu = torch.from_numpy(
+            input_batch.query_start_loc_np[: num_reqs + 1]
+        )
+        query_start_loc_gpu = input_batch.query_start_loc[: num_reqs + 1]
         max_query_len = input_batch.num_scheduled_tokens.max().item()
         seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound
         if for_capture:
@@ -174,7 +177,7 @@ class DefaultModelState(ModelState):
             attn_groups=attn_groups,
             num_reqs=num_reqs,
             num_tokens=num_tokens,
-            query_start_loc_gpu=input_batch.query_start_loc,
+            query_start_loc_gpu=query_start_loc_gpu,
             query_start_loc_cpu=query_start_loc_cpu,
             max_query_len=max_query_len,
             seq_lens=input_batch.seq_lens,
@@ -185,6 +188,7 @@ class DefaultModelState(ModelState):
             seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
             dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
             positions=input_batch.positions,
+            is_prefilling=torch.from_numpy(input_batch.is_prefilling_np),
             mm_req_doc_ranges=req_doc_ranges,
             for_cudagraph_capture=for_capture,
             rswa_prefix_lens=input_batch.prompt_lens,

vllm/v1/worker/gpu/pcp_manager.py

@@ -0,0 +1,680 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, replace
+
+import numpy as np
+import torch
+
+from vllm.config import CUDAGraphMode, VllmConfig
+from vllm.distributed.parallel_state import get_dcp_group, get_pcp_group
+from vllm.logger import init_logger
+from vllm.v1.attention.backends.utils import PAD_SLOT_ID
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.buffer_utils import async_copy_to_gpu
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
+from vllm.v1.worker.gpu.input_batch import (
+    InputBatch,
+    InputBuffers,
+    combine_sampled_and_draft_tokens,
+    prepare_pos_seq_lens,
+)
+from vllm.v1.worker.gpu.states import RequestState
+
+logger = init_logger(__name__)
+
+
+@dataclass(frozen=True)
+class RankSegment:
+    global_batch_req_idx: int
+    global_batch_slice: slice
+    rank_local_batch_slice: slice
+
+    @property
+    def num_tokens(self) -> int:
+        return self.global_batch_slice.stop - self.global_batch_slice.start
+
+
+class PCPManager:
+    """MRV2 PC batch manager.
+
+    The model runner keeps the global scheduled batch. This manager rewrites only
+    the per-step InputBatch into rank-local DualChunkSwap rows and keeps the
+    global-batch view private to restore to the global batch shape before
+    sampling/postprocess.
+    """
+
+    def __init__(
+        self,
+        pcp_world_size: int,
+        pcp_rank: int,
+        device: torch.device,
+        req_states: RequestState | None = None,
+        max_num_reqs: int | None = None,
+        max_num_tokens: int | None = None,
+        block_tables: BlockTables | None = None,
+        dcp_world_size: int = 1,
+        dcp_rank: int = 0,
+        cp_interleave: int = 1,
+    ) -> None:
+        self.pcp_world_size = pcp_world_size
+        self.pcp_rank = pcp_rank
+        self.device = device
+        self.dcp_world_size = dcp_world_size
+        self.dcp_rank = dcp_rank
+        self.cp_interleave = cp_interleave
+
+        self._global_batch: InputBatch | None = None
+        self._req_states = req_states
+        self._block_tables = block_tables
+        self._hidden_restore_idx: torch.Tensor | None = None
+        self._padded_gather_idx: torch.Tensor | None = None
+        self._gathered_kv_write_mask: torch.Tensor | None = None
+        self._pad_slot_id = torch.tensor(PAD_SLOT_ID, dtype=torch.int64, device=device)
+
+        max_num_local_reqs = 2 * max_num_reqs if max_num_reqs is not None else None
+        self._input_buffers = (
+            InputBuffers(max_num_local_reqs, max_num_tokens, device)
+            if max_num_local_reqs is not None and max_num_tokens is not None
+            else None
+        )
+        self._local_req_idx = (
+            torch.arange(max_num_local_reqs, dtype=torch.int32, device=device)
+            if max_num_local_reqs is not None
+            else None
+        )
+        self._local_block_tables: tuple[torch.Tensor, ...] | None
+        self._local_block_table_ptrs: torch.Tensor | None
+        if block_tables is not None and max_num_local_reqs is not None:
+            self._local_block_tables = tuple(
+                table.new_zeros((max_num_local_reqs, table.shape[1]))
+                for table in block_tables.input_block_tables
+            )
+            self._local_block_table_ptrs = torch.tensor(
+                [table.data_ptr() for table in self._local_block_tables],
+                dtype=torch.uint64,
+                device=device,
+            )
+        else:
+            self._local_block_tables = None
+            self._local_block_table_ptrs = None
+        num_kv_cache_groups = (
+            block_tables.num_kv_cache_groups if block_tables is not None else 0
+        )
+        self._global_batch_slot_mappings = (
+            torch.empty(
+                num_kv_cache_groups,
+                max_num_tokens,
+                dtype=torch.int64,
+                device=device,
+            )
+            if max_num_tokens is not None and num_kv_cache_groups > 0
+            else None
+        )
+        self._gathered_kv_slot_mappings = (
+            torch.empty(
+                num_kv_cache_groups,
+                max_num_tokens * pcp_world_size,
+                dtype=torch.int64,
+                device=device,
+            )
+            if max_num_tokens is not None and num_kv_cache_groups > 0
+            else None
+        )
+
+    @staticmethod
+    def validate_config(
+        vllm_config: VllmConfig,
+        supports_mm_inputs: bool,
+    ) -> None:
+        parallel_config = vllm_config.parallel_config
+        model_config = vllm_config.model_config
+        pcp_size = parallel_config.prefill_context_parallel_size
+        if pcp_size <= 1:
+            return
+
+        if not model_config.use_mla:
+            raise NotImplementedError("MRV2 PCP currently supports MLA models only.")
+        if parallel_config.pipeline_parallel_size > 1:
+            raise NotImplementedError("MRV2 PCP does not support PP yet.")
+        if model_config.is_encoder_decoder:
+            raise NotImplementedError(
+                "MRV2 PCP does not support encoder-decoder models yet."
+            )
+        if supports_mm_inputs:
+            raise NotImplementedError("MRV2 PCP does not support MM inputs yet.")
+        if vllm_config.lora_config is not None:
+            raise NotImplementedError("MRV2 PCP does not support LoRA yet.")
+        if vllm_config.speculative_config is not None:
+            raise NotImplementedError(
+                "MRV2 PCP does not support speculative decoding yet."
+            )
+        is_sparse_mla = hasattr(model_config.hf_text_config, "index_topk")
+        if (
+            is_sparse_mla
+            and vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
+        ):
+            raise NotImplementedError(
+                "MRV2 sparse MLA PCP does not support CUDA graphs yet. "
+                "Set -cc.cudagraph_mode=NONE."
+            )
+        if vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs():
+            raise NotImplementedError("MRV2 PCP supports PIECEWISE CUDA graphs only.")
+
+    @staticmethod
+    def _reorder_segments(
+        segments: list[RankSegment],
+        num_computed_tokens: np.ndarray,
+        is_prefilling: np.ndarray,
+        query_start_loc_np: np.ndarray,
+    ) -> list[RankSegment]:
+        """Move pure prefills last to match the batch ordering expected by
+        attention backends like MLA and sparse MLA.
+        """
+
+        def is_pure_prefill(segment: RankSegment) -> bool:
+            req_idx = segment.global_batch_req_idx
+            start_pos = (
+                num_computed_tokens[req_idx]
+                + segment.global_batch_slice.start
+                - query_start_loc_np[req_idx]
+            )
+            return is_prefilling[req_idx] and start_pos == 0
+
+        segments.sort(key=is_pure_prefill)
+        rank_offset = 0
+        for index, segment in enumerate(segments):
+            segments[index] = replace(
+                segment,
+                rank_local_batch_slice=slice(
+                    rank_offset, rank_offset + segment.num_tokens
+                ),
+            )
+            rank_offset += segment.num_tokens
+        return segments
+
+    def _get_rank_segments(
+        self,
+        rank: int,
+        num_scheduled_tokens: np.ndarray,
+        num_computed_tokens: np.ndarray,
+        is_prefilling: np.ndarray,
+        query_start_loc_np: np.ndarray,
+    ) -> list[RankSegment]:
+        """Build one rank's attention-compatible DualChunkSwap rows.
+
+        PCP=4 partitions each prefill into eight chunks:
+
+            full:  | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
+            rank 0:  0                           7
+            rank 1:      1                   6
+            rank 2:          2           5
+            rank 3:              3   4
+        """
+        rank_segments = []
+        rank_offset = 0
+        num_chunks = 2 * self.pcp_world_size
+        for global_batch_req_idx, num_tokens in enumerate(num_scheduled_tokens):
+            query_len = int(num_tokens)
+            if query_len == 0:
+                continue
+            global_batch_start = int(query_start_loc_np[global_batch_req_idx])
+            chunk_indices: tuple[int, ...]
+            if bool(is_prefilling[global_batch_req_idx]):
+                chunk_size = (query_len + num_chunks - 1) // num_chunks
+                chunk_indices = (rank, num_chunks - 1 - rank)
+            else:  # decodes are replicated
+                chunk_size = query_len
+                chunk_indices = (0,)
+
+            for chunk_idx in chunk_indices:
+                chunk_offset = chunk_idx * chunk_size
+                chunk_len = min(chunk_size, query_len - chunk_offset)
+                if chunk_len <= 0:
+                    continue
+                chunk_start = global_batch_start + chunk_offset
+                rank_segments.append(
+                    RankSegment(
+                        global_batch_req_idx=global_batch_req_idx,
+                        global_batch_slice=slice(chunk_start, chunk_start + chunk_len),
+                        rank_local_batch_slice=slice(
+                            rank_offset, rank_offset + chunk_len
+                        ),
+                    )
+                )
+                rank_offset += chunk_len
+        return self._reorder_segments(
+            rank_segments,
+            num_computed_tokens,
+            is_prefilling,
+            query_start_loc_np,
+        )
+
+    def _build_batch_layout(
+        self,
+        num_scheduled_tokens: np.ndarray,
+        num_computed_tokens: np.ndarray,
+        is_prefilling: np.ndarray,
+        query_start_loc_np: np.ndarray,
+    ) -> tuple[list[list[RankSegment]], list[int]]:
+        segments_by_rank = []
+        per_rank_num_tokens = []
+        for rank in range(self.pcp_world_size):
+            segments = self._get_rank_segments(
+                rank,
+                num_scheduled_tokens,
+                num_computed_tokens,
+                is_prefilling,
+                query_start_loc_np,
+            )
+            num_rank_tokens = sum(segment.num_tokens for segment in segments)
+            segments_by_rank.append(segments)
+            per_rank_num_tokens.append(num_rank_tokens)
+
+        # PCP=2 example:
+        #   global batch:       [A B C D E F G]
+        #   rank 0 / rank 1:    [A B G] / [C D E F]
+        #   padded gathered:    [A B G _ | C D E F]
+        #   hidden_restore_idx: [0, 1, 4, 5, 6, 7, 2]
+        #   padded_gather_idx:  [0, 1, 6, 0, 2, 3, 4, 5]
+        # Therefore global = gathered[hidden_restore_idx] and
+        # padded_gathered = global[padded_gather_idx].
+        hidden_restore_idx = np.empty(int(query_start_loc_np[-1]), dtype=np.int64)
+        padded_num_tokens = max(per_rank_num_tokens)
+        num_expanded_tokens = padded_num_tokens * self.pcp_world_size
+        padded_gather_idx = np.zeros(num_expanded_tokens, dtype=np.int64)
+        gathered_kv_write_mask = np.zeros(num_expanded_tokens, dtype=np.bool_)
+        for rank, segments in enumerate(segments_by_rank):
+            expanded_rank_offset = rank * padded_num_tokens
+            for segment in segments:
+                padded_gathered_slice = slice(
+                    expanded_rank_offset + segment.rank_local_batch_slice.start,
+                    expanded_rank_offset + segment.rank_local_batch_slice.stop,
+                )
+                padded_gather_idx[padded_gathered_slice] = np.arange(
+                    segment.global_batch_slice.start,
+                    segment.global_batch_slice.stop,
+                    dtype=np.int64,
+                )
+                # Cache insertion pairs one slot entry with each rank's local decode.
+                if not bool(is_prefilling[segment.global_batch_req_idx]) and rank != 0:
+                    continue
+                gathered_kv_write_mask[padded_gathered_slice] = True
+                hidden_restore_idx[segment.global_batch_slice] = np.arange(
+                    padded_gathered_slice.start,
+                    padded_gathered_slice.stop,
+                    dtype=np.int64,
+                )
+
+        self._hidden_restore_idx = async_copy_to_gpu(
+            hidden_restore_idx, device=self.device
+        )
+        self._padded_gather_idx = async_copy_to_gpu(
+            padded_gather_idx, device=self.device
+        )
+        self._gathered_kv_write_mask = async_copy_to_gpu(
+            gathered_kv_write_mask, device=self.device
+        )
+        return segments_by_rank, per_rank_num_tokens
+
+    def partition_batch(self, input_batch: InputBatch) -> InputBatch:
+        assert self._req_states is not None
+        assert self._input_buffers is not None
+        req_states = self._req_states
+        input_buffers = self._input_buffers
+        if input_batch.num_draft_tokens > 0:
+            raise NotImplementedError("MRV2 PCP does not support spec decode yet.")
+
+        global_batch = input_batch
+        self._global_batch = global_batch
+
+        num_scheduled_tokens = global_batch.num_scheduled_tokens
+        num_computed_tokens = global_batch.num_computed_tokens_np
+        is_prefilling = global_batch.is_prefilling_np
+
+        segments_by_rank, per_rank_num_tokens = self._build_batch_layout(
+            num_scheduled_tokens,
+            num_computed_tokens,
+            is_prefilling,
+            global_batch.query_start_loc_np,
+        )
+
+        local_segments = segments_by_rank[self.pcp_rank]
+        if not local_segments:
+            local_segments = [
+                RankSegment(
+                    global_batch_req_idx=0,
+                    global_batch_slice=slice(0, 0),
+                    rank_local_batch_slice=slice(0, 0),
+                )
+            ]
+
+        num_local_reqs = len(local_segments)
+        if num_local_reqs > 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}."
+            )
+
+        local_to_global_batch_req_idx_np = np.fromiter(
+            (segment.global_batch_req_idx for segment in local_segments),
+            dtype=np.int32,
+            count=num_local_reqs,
+        )
+        local_start_pos_np = np.fromiter(
+            (
+                num_computed_tokens[segment.global_batch_req_idx]
+                + segment.global_batch_slice.start
+                - global_batch.query_start_loc_np[segment.global_batch_req_idx]
+                for segment in local_segments
+            ),
+            dtype=np.int32,
+            count=num_local_reqs,
+        )
+        local_num_scheduled_tokens = np.fromiter(
+            (segment.num_tokens for segment in local_segments),
+            dtype=np.int32,
+            count=num_local_reqs,
+        )
+        local_to_global_req_idx_np = global_batch.idx_mapping_np[
+            local_to_global_batch_req_idx_np
+        ]
+        local_req_ids = [
+            global_batch.req_ids[global_batch_req_idx]
+            for global_batch_req_idx in local_to_global_batch_req_idx_np
+        ]
+
+        num_local_tokens = int(local_num_scheduled_tokens.sum())
+        num_local_tokens_padded = max(per_rank_num_tokens)
+        fresh_prefills = int(
+            np.count_nonzero(is_prefilling & (num_computed_tokens == 0))
+        )
+        continued_prefills = int(
+            np.count_nonzero(is_prefilling & (num_computed_tokens > 0))
+        )
+        logger.debug(
+            "PCP batch: rank=%d global_batch_reqs=%d fresh_prefills=%d "
+            "continued_prefills=%d decodes=%d local_reqs=%d "
+            "local_tokens=%d per_rank_tokens=%s",
+            self.pcp_rank,
+            global_batch.num_reqs,
+            fresh_prefills,
+            continued_prefills,
+            global_batch.num_reqs - fresh_prefills - continued_prefills,
+            num_local_reqs,
+            num_local_tokens,
+            per_rank_num_tokens,
+        )
+        if num_local_tokens_padded > input_buffers.max_num_tokens:
+            raise RuntimeError(
+                "PCP local token count exceeds the MRV2 input buffer size: "
+                f"{num_local_tokens_padded} > {input_buffers.max_num_tokens}."
+            )
+        rank_token_start = self.pcp_rank * num_local_tokens_padded
+        assert self._padded_gather_idx is not None
+        local_gather_idx = self._padded_gather_idx[
+            rank_token_start : rank_token_start + num_local_tokens_padded
+        ]
+        torch.index_select(
+            global_batch.input_ids,
+            0,
+            local_gather_idx,
+            out=input_buffers.input_ids[:num_local_tokens_padded],
+        )
+
+        local_query_start_loc_np = np.empty(
+            input_buffers.max_num_reqs + 1, dtype=np.int32
+        )
+        local_query_start_loc_np[0] = 0
+        local_query_start_loc_out = local_query_start_loc_np[1 : num_local_reqs + 1]
+        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_to_global_req_idx = async_copy_to_gpu(
+            local_to_global_req_idx_np, device=self.device
+        )
+        local_start_pos = async_copy_to_gpu(local_start_pos_np, device=self.device)
+
+        assert self._local_req_idx is not None
+        prepare_pos_seq_lens(
+            self._local_req_idx[:num_local_reqs],
+            local_query_start_loc,
+            local_start_pos,
+            input_buffers.positions,
+            input_buffers.seq_lens[:num_local_reqs],
+        )
+        seq_lens = input_buffers.seq_lens[:num_local_reqs]
+        is_padding = input_buffers.is_padding[:num_local_tokens_padded]
+        is_padding[:num_local_tokens].fill_(False)
+        is_padding[num_local_tokens:].fill_(True)
+        if num_local_tokens_padded > num_local_tokens:
+            input_buffers.input_ids[:num_local_tokens_padded].masked_fill_(
+                is_padding, 0
+            )
+            input_buffers.positions[:num_local_tokens_padded].masked_fill_(
+                is_padding, 0
+            )
+
+        total_num_logits = num_local_reqs if num_local_tokens > 0 else 0
+        if total_num_logits > 0:
+            cu_num_logits_np = np.arange(num_local_reqs + 1, dtype=np.int32)
+            cu_num_logits = torch.arange(
+                num_local_reqs + 1, device=self.device, dtype=torch.int32
+            )
+        else:
+            cu_num_logits_np = np.zeros(num_local_reqs + 1, dtype=np.int32)
+            cu_num_logits = torch.zeros(
+                num_local_reqs + 1, device=self.device, dtype=torch.int32
+            )
+        logits_indices = combine_sampled_and_draft_tokens(
+            input_buffers.input_ids,
+            local_to_global_req_idx,
+            req_states.last_sampled_tokens,
+            local_query_start_loc,
+            seq_lens,
+            req_states.prefill_len.gpu,
+            req_states.draft_tokens,
+            cu_num_logits,
+            total_num_logits,
+            1,
+        )
+
+        local_prefill_len_np = global_batch.prefill_len_np[
+            local_to_global_batch_req_idx_np
+        ]
+        local_num_computed_prefill_tokens_np = np.minimum(
+            local_start_pos_np, local_prefill_len_np
+        )
+        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
+
+        dcp_local_seq_lens = None
+        if self.dcp_world_size > 1:
+            prepare_dcp_local_seq_lens(
+                input_buffers.dcp_local_seq_lens,
+                seq_lens,
+                num_local_reqs,
+                self.dcp_world_size,
+                self.dcp_rank,
+                self.cp_interleave,
+            )
+            dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_local_reqs]
+
+        return replace(
+            input_batch,
+            req_ids=local_req_ids,
+            num_reqs=num_local_reqs,
+            num_reqs_after_padding=num_local_reqs,
+            idx_mapping=local_to_global_req_idx,
+            idx_mapping_np=local_to_global_req_idx_np,
+            expanded_idx_mapping=local_to_global_req_idx,
+            expanded_local_pos=torch.zeros(
+                num_local_reqs, dtype=torch.int32, device=self.device
+            ),
+            num_scheduled_tokens=local_num_scheduled_tokens,
+            num_tokens=num_local_tokens,
+            num_tokens_after_padding=num_local_tokens_padded,
+            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],
+            seq_lens=seq_lens,
+            seq_lens_cpu_upper_bound=torch.from_numpy(seq_lens_cpu_upper_bound_np),
+            dcp_local_seq_lens=dcp_local_seq_lens,
+            num_computed_tokens_np=local_start_pos_np,
+            prefill_len_np=local_prefill_len_np,
+            num_computed_prefill_tokens_np=local_num_computed_prefill_tokens_np,
+            is_prefilling_np=local_is_prefilling_np,
+            max_seq_len_np=global_batch.max_seq_len_np[local_to_global_batch_req_idx_np]
+            if global_batch.max_seq_len_np is not None
+            else None,
+            input_ids=input_buffers.input_ids[:num_local_tokens_padded],
+            positions=input_buffers.positions[:num_local_tokens_padded],
+            is_padding=is_padding,
+            logits_indices=logits_indices,
+            cu_num_logits=cu_num_logits,
+            cu_num_logits_np=cu_num_logits_np,
+            prompt_lens=None,
+        )
+
+    def prepare_attn(
+        self, input_batch: InputBatch
+    ) -> tuple[tuple[torch.Tensor, ...], torch.Tensor]:
+        assert self._block_tables is not None
+        assert self._local_block_tables is not None
+        assert self._local_block_table_ptrs is not None
+        block_tables = self._block_tables.gather_block_tables(
+            input_batch.idx_mapping,
+            input_batch.num_reqs_after_padding,
+            out=self._local_block_tables,
+            out_ptrs=self._local_block_table_ptrs,
+        )
+        slot_mappings = self.prepare_slot_mappings()
+        return block_tables, slot_mappings
+
+    def prepare_slot_mappings(self) -> torch.Tensor:
+        assert self._block_tables is not None
+        assert self._global_batch_slot_mappings is not None
+        assert self._global_batch is not None
+        global_batch = self._global_batch
+        global_batch_slot_mappings = self._block_tables.compute_slot_mappings(
+            global_batch.idx_mapping,
+            global_batch.query_start_loc,
+            global_batch.positions,
+            global_batch.num_tokens,
+            out=self._global_batch_slot_mappings,
+        )
+        return self._convert_to_gathered_slot_mappings(global_batch_slot_mappings)
+
+    def get_dummy_slot_mappings(self, num_tokens: int) -> torch.Tensor:
+        assert self._gathered_kv_slot_mappings is not None
+        self._gathered_kv_slot_mappings.fill_(PAD_SLOT_ID)
+        return self._gathered_kv_slot_mappings[:, : num_tokens * self.pcp_world_size]
+
+    def _convert_to_gathered_slot_mappings(
+        self,
+        global_batch_slot_mappings: torch.Tensor,
+    ) -> torch.Tensor:
+        assert self._padded_gather_idx is not None
+        assert self._gathered_kv_write_mask is not None
+        padded_gather_idx = self._padded_gather_idx
+        num_expanded_tokens = padded_gather_idx.shape[0]
+        if self._gathered_kv_slot_mappings is None:
+            self._gathered_kv_slot_mappings = global_batch_slot_mappings.new_empty(
+                global_batch_slot_mappings.shape[0], num_expanded_tokens
+            )
+        gathered_kv_slot_mappings = self._gathered_kv_slot_mappings[
+            :, :num_expanded_tokens
+        ]
+        torch.index_select(
+            global_batch_slot_mappings,
+            1,
+            padded_gather_idx,
+            out=gathered_kv_slot_mappings,
+        )
+        torch.where(
+            self._gathered_kv_write_mask.unsqueeze(0),
+            gathered_kv_slot_mappings,
+            self._pad_slot_id,
+            out=gathered_kv_slot_mappings,
+        )
+        return gathered_kv_slot_mappings
+
+    def restore_hidden_states(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        if self._hidden_restore_idx is None:
+            return hidden_states
+        gathered = get_pcp_group().all_gather(hidden_states, dim=0)
+        return gathered[self._hidden_restore_idx]
+
+    def restore_for_sampling(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> tuple[torch.Tensor, InputBatch]:
+        assert self._global_batch is not None
+        return self.restore_hidden_states(hidden_states), self._global_batch
+
+
+def maybe_partition_pcp_batch(
+    manager: PCPManager | None,
+    input_batch: InputBatch,
+) -> InputBatch:
+    if manager is None:
+        return input_batch
+    return manager.partition_batch(input_batch)
+
+
+def maybe_get_pcp_dummy_slot_mappings(
+    manager: PCPManager | None,
+    block_tables: BlockTables,
+    num_tokens: int,
+) -> torch.Tensor:
+    if manager is None:
+        return block_tables.get_dummy_slot_mappings(num_tokens)
+    return manager.get_dummy_slot_mappings(num_tokens)
+
+
+def maybe_restore_pcp_for_sampling(
+    manager: PCPManager | None,
+    hidden_states: torch.Tensor | None,
+    input_batch: InputBatch,
+) -> tuple[torch.Tensor, InputBatch]:
+    assert hidden_states is not None
+    if manager is None:
+        return hidden_states, input_batch
+    return manager.restore_for_sampling(hidden_states)
+
+
+def maybe_build_pcp_manager(
+    vllm_config: VllmConfig,
+    device: torch.device,
+    supports_mm_inputs: bool,
+    req_states: RequestState,
+    block_tables: BlockTables,
+) -> PCPManager | None:
+    parallel_config = vllm_config.parallel_config
+    pcp_size = parallel_config.prefill_context_parallel_size
+    if pcp_size <= 1:
+        return None
+
+    PCPManager.validate_config(vllm_config, supports_mm_inputs)
+
+    pcp_rank = get_pcp_group().rank_in_group
+    dcp_size = parallel_config.decode_context_parallel_size
+    dcp_rank = get_dcp_group().rank_in_group if dcp_size > 1 else 0
+
+    return PCPManager(
+        pcp_world_size=pcp_size,
+        pcp_rank=pcp_rank,
+        device=device,
+        req_states=req_states,
+        max_num_reqs=vllm_config.scheduler_config.max_num_seqs,
+        max_num_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
+        block_tables=block_tables,
+        dcp_world_size=dcp_size,
+        dcp_rank=dcp_rank,
+        cp_interleave=parallel_config.cp_kv_cache_interleave_size,
+    )