[Core] Add Prefill Context Parallelism (PCP) and All-to-All DCP communication

This PR adds Prefill Context Parallelism (PCP) support for splitting prefill
tokens across ranks using a DualChunkSwap pattern, and integrates an All-to-All
communication backend for Decode Context Parallelism (DCP).

Key changes:
- Add PCP with DualChunkSwap token partitioning for balanced prefill computation
- Add All-to-All DCP communication backend reducing NCCL calls from 3 to 2
- Restrict DCP+PCP to two clean configurations:
  - Case 1: DCP = PCP (same TP position, all-reduce only)
  - Case 2: DCP = TP × PCP (full TP all-gather, all-reduce + slice)
- Add PCPManager for buffer management and input partitioning
- Update attention backends (FlashAttention, FlashInfer, MLA) for PCP support
- Add comprehensive tests for DCP operations

Co-Authored-By: QiuChunshuo <qiuchunshuo@huawei.com>
Co-Authored-By: zhenwenqi2024 <zhenwenqi_2022@qq.com>
Co-Authored-By: FENP <yuanyongjie.yyj@antgroup.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

docs/design/dcp_communication_patterns.md

@@ -0,0 +1,350 @@
+# DCP Communication Patterns
+
+This document describes the communication patterns for Decode Context Parallelism (DCP) with various configurations of Tensor Parallelism (TP) and Prefill Context Parallelism (PCP).
+
+## Background
+
+- **TP (Tensor Parallelism)**: Splits attention heads across ranks. Each rank has `H/TP` heads.
+- **PCP (Prefill Context Parallelism)**: Splits prefill tokens across ranks. Each PCP slice has its own TP group.
+- **DCP (Decode Context Parallelism)**: Splits KV cache context across ranks for decode.
+
+### Rank Layout
+
+Ranks are laid out as `(ep, dp, pp, pcp, tp)`. For simplicity, we assume `ep=dp=pp=1`.
+
+For **PCP=2, TP=4** (8 ranks):
+```
+        TP=0  TP=1  TP=2  TP=3
+PCP=0    0     1     2     3
+PCP=1    4     5     6     7
+```
+
+For **PCP=2, TP=2** (4 ranks):
+```
+        TP=0  TP=1
+PCP=0    0     1
+PCP=1    2     3
+```
+
+For **PCP=1, TP=4** (4 ranks):
+```
+        TP=0  TP=1  TP=2  TP=3
+PCP=0    0     1     2     3
+```
+
+### DCP Group Formation
+
+DCP groups are formed by spanning **PCP first, then TP**:
+1. Transpose layout to `(tp, pcp)`
+2. Flatten and reshape to `(-1, dcp_size)`
+
+---
+
+## Case 1: PCP=1, TP=4, DCP=4
+
+**Groups:**
+| Group Type | Ranks |
+|------------|-------|
+| TP group | `[0, 1, 2, 3]` |
+| DCP group | `[0, 1, 2, 3]` (same as TP) |
+
+**Head distribution:**
+| Rank | TP Position | Heads |
+|------|-------------|-------|
+| 0 | 0 | `[0, H/4)` |
+| 1 | 1 | `[H/4, H/2)` |
+| 2 | 2 | `[H/2, 3H/4)` |
+| 3 | 3 | `[3H/4, H)` |
+
+**DCP Decode Communication:**
+
+```
+Step 1: TP All-Gather (query)
+┌─────────────────────────────────────────────────────┐
+│  All ranks gather with TP group [0,1,2,3]           │
+│  Each rank: H/4 heads → H heads                     │
+└─────────────────────────────────────────────────────┘
+
+Step 2: Attention
+┌─────────────────────────────────────────────────────┐
+│  Each rank computes attention with ALL H heads      │
+│  against its local KV slice (1/4 of context)        │
+└─────────────────────────────────────────────────────┘
+
+Step 3: DCP Reduce (reduce-scatter)
+┌─────────────────────────────────────────────────────┐
+│  DCP == TP, so reduce-scatter works directly        │
+│  Each rank gets its original H/4 heads back         │
+│  Rank 0: heads [0, H/4)                             │
+│  Rank 1: heads [H/4, H/2)                           │
+│  Rank 2: heads [H/2, 3H/4)                          │
+│  Rank 3: heads [3H/4, H)                            │
+└─────────────────────────────────────────────────────┘
+```
+
+**Optimization:** None needed. DCP == TP is optimal.
+
+---
+
+## Case 2: PCP=1, TP=4, DCP=2
+
+**Groups:**
+| Group Type | Ranks |
+|------------|-------|
+| TP group | `[0, 1, 2, 3]` |
+| DCP groups | `[0, 1]`, `[2, 3]` |
+
+**Head distribution:**
+| Rank | TP Position | DCP Group | Heads |
+|------|-------------|-----------|-------|
+| 0 | 0 | 0 | `[0, H/4)` |
+| 1 | 1 | 0 | `[H/4, H/2)` |
+| 2 | 2 | 1 | `[H/2, 3H/4)` |
+| 3 | 3 | 1 | `[3H/4, H)` |
+
+**DCP Decode Communication:**
+
+```
+Step 1: TP All-Gather (query)
+┌─────────────────────────────────────────────────────┐
+│  All ranks gather with TP group [0,1,2,3]           │
+│  Each rank: H/4 heads → H heads                     │
+│  (Gathers more than needed!)                        │
+└─────────────────────────────────────────────────────┘
+
+Step 2: Attention
+┌─────────────────────────────────────────────────────┐
+│  DCP group [0,1]: each computes with H heads        │
+│    against 1/2 of context                           │
+│  DCP group [2,3]: each computes with H heads        │
+│    against 1/2 of context                           │
+└─────────────────────────────────────────────────────┘
+
+Step 3: DCP Reduce (all-reduce + slice)
+┌─────────────────────────────────────────────────────┐
+│  DCP ⊂ TP, so need all-reduce + manual slice        │
+│  Within [0,1]: all-reduce, then each slices to H/4  │
+│  Within [2,3]: all-reduce, then each slices to H/4  │
+└─────────────────────────────────────────────────────┘
+```
+
+**Optimization:** Use partial-TP all-gather
+- DCP group `[0,1]` covers TP positions `{0, 1}` → only need `H/2` heads
+- DCP group `[2,3]` covers TP positions `{2, 3}` → only need `H/2` heads
+- Partial-TP groups: `[0,1]` and `[2,3]` (same as DCP groups in this case)
+
+```
+Optimized Step 1: Partial-TP All-Gather
+┌─────────────────────────────────────────────────────┐
+│  Ranks [0,1] gather within [0,1]: H/4 → H/2 heads   │
+│  Ranks [2,3] gather within [2,3]: H/4 → H/2 heads   │
+│  (Half the communication!)                          │
+└─────────────────────────────────────────────────────┘
+```
+
+---
+
+## Case 3: PCP=2, TP=2, DCP=4
+
+**Groups:**
+| Group Type | Ranks |
+|------------|-------|
+| TP groups | `[0, 1]` (PCP=0), `[2, 3]` (PCP=1) |
+| DCP group | `[0, 2, 1, 3]` (all ranks) |
+| PCP group | `[0, 1, 2, 3]` |
+
+**Head distribution:**
+| Rank | TP Position | PCP Slice | Heads |
+|------|-------------|-----------|-------|
+| 0 | 0 | 0 | `[0, H/2)` |
+| 1 | 1 | 0 | `[H/2, H)` |
+| 2 | 0 | 1 | `[0, H/2)` |
+| 3 | 1 | 1 | `[H/2, H)` |
+
+**DCP Decode Communication:**
+
+```
+Step 1: TP All-Gather (query)
+┌─────────────────────────────────────────────────────┐
+│  Rank 0 gathers with [0,1] → H heads                │
+│  Rank 1 gathers with [0,1] → H heads                │
+│  Rank 2 gathers with [2,3] → H heads                │
+│  Rank 3 gathers with [2,3] → H heads                │
+└─────────────────────────────────────────────────────┘
+
+Step 2: Attention
+┌─────────────────────────────────────────────────────┐
+│  Each rank computes with ALL H heads                │
+│  against its local KV slice (1/4 of context)        │
+└─────────────────────────────────────────────────────┘
+
+Step 3: DCP Reduce (all-reduce + slice)
+┌─────────────────────────────────────────────────────┐
+│  DCP=4 > TP=2, so need all-reduce + slice           │
+│  All-reduce across [0,2,1,3]                        │
+│  Each rank slices to its TP-local H/2 heads         │
+└─────────────────────────────────────────────────────┘
+```
+
+**Note:** DCP spans both PCP and TP dimensions. All ranks are in one DCP group.
+
+---
+
+## Case 4: PCP=2, TP=2, DCP=2
+
+**Groups:**
+| Group Type | Ranks |
+|------------|-------|
+| TP groups | `[0, 1]` (PCP=0), `[2, 3]` (PCP=1) |
+| DCP groups | `[0, 2]`, `[1, 3]` (span PCP, same TP!) |
+
+**Head distribution:**
+| Rank | TP Position | PCP Slice | DCP Group | Heads |
+|------|-------------|-----------|-----------|-------|
+| 0 | 0 | 0 | 0 | `[0, H/2)` |
+| 2 | 0 | 1 | 0 | `[0, H/2)` |
+| 1 | 1 | 0 | 1 | `[H/2, H)` |
+| 3 | 1 | 1 | 1 | `[H/2, H)` |
+
+**Key insight:** Ranks in the same DCP group have the **same TP position** (same heads)!
+
+**DCP Decode Communication:**
+
+```
+Step 1: TP All-Gather (query)
+┌─────────────────────────────────────────────────────┐
+│  Rank 0 gathers with [0,1] → H heads                │
+│  Rank 2 gathers with [2,3] → H heads                │
+│  (Both DCP peers do redundant work!)                │
+└─────────────────────────────────────────────────────┘
+
+Step 2: Attention
+┌─────────────────────────────────────────────────────┐
+│  DCP group [0,2]: each computes with H heads        │
+│    Rank 0: against KV slice A                       │
+│    Rank 2: against KV slice B                       │
+└─────────────────────────────────────────────────────┘
+
+Step 3: DCP Reduce (all-reduce, no scatter needed)
+┌─────────────────────────────────────────────────────┐
+│  All-reduce within [0,2] and [1,3]                  │
+│  Each rank keeps its original H/2 heads             │
+│  (No redistribution needed - same TP position!)    │
+└─────────────────────────────────────────────────────┘
+```
+
+**Optimization:** Skip TP all-gather entirely!
+- Ranks 0 and 2 already have the same heads `[0, H/2)`
+- They can compute attention with `H/2` heads directly
+- Partial-TP group size = 1 (no all-gather needed)
+
+```
+Optimized Step 1: No All-Gather!
+┌─────────────────────────────────────────────────────┐
+│  Each rank keeps its original H/2 heads             │
+│  (Zero communication!)                              │
+└─────────────────────────────────────────────────────┘
+```
+
+---
+
+## Case 5: PCP=2, TP=4, DCP=4
+
+**Groups:**
+| Group Type | Ranks |
+|------------|-------|
+| TP groups | `[0,1,2,3]` (PCP=0), `[4,5,6,7]` (PCP=1) |
+| DCP groups | `[0,4,1,5]`, `[2,6,3,7]` |
+
+**Head distribution:**
+| Rank | TP Position | PCP Slice | DCP Group | Heads |
+|------|-------------|-----------|-----------|-------|
+| 0 | 0 | 0 | 0 | `[0, H/4)` |
+| 4 | 0 | 1 | 0 | `[0, H/4)` |
+| 1 | 1 | 0 | 0 | `[H/4, H/2)` |
+| 5 | 1 | 1 | 0 | `[H/4, H/2)` |
+| 2 | 2 | 0 | 1 | `[H/2, 3H/4)` |
+| 6 | 2 | 1 | 1 | `[H/2, 3H/4)` |
+| 3 | 3 | 0 | 1 | `[3H/4, H)` |
+| 7 | 3 | 1 | 1 | `[3H/4, H)` |
+
+**Key insight:** Each DCP group covers **2 unique TP positions** (half of TP=4).
+
+**DCP Decode Communication (current):**
+
+```
+Step 1: TP All-Gather (query)
+┌─────────────────────────────────────────────────────┐
+│  Rank 0 gathers with [0,1,2,3] → H heads            │
+│  Rank 4 gathers with [4,5,6,7] → H heads            │
+│  (Gathers H heads but only needs H/2!)              │
+└─────────────────────────────────────────────────────┘
+
+Step 2: Attention
+┌─────────────────────────────────────────────────────┐
+│  Each rank computes with ALL H heads                │
+│  against its local KV slice (1/4 of context)        │
+└─────────────────────────────────────────────────────┘
+
+Step 3: DCP Reduce (all-reduce + slice)
+┌─────────────────────────────────────────────────────┐
+│  DCP=4, TP=4, but PCP>1 so can't reduce-scatter     │
+│  All-reduce within DCP group, slice to H/4 heads    │
+└─────────────────────────────────────────────────────┘
+```
+
+**Optimization:** Use partial-TP all-gather
+
+DCP group analysis:
+- `[0,4,1,5]` covers TP positions `{0, 1}` → needs `H/2` heads
+- `[2,6,3,7]` covers TP positions `{2, 3}` → needs `H/2` heads
+
+Partial-TP groups (per PCP slice):
+| Partial-TP Group | Ranks | TP Positions | For DCP Group |
+|------------------|-------|--------------|---------------|
+| `[0, 1]` | PCP=0 | {0, 1} | 0 |
+| `[4, 5]` | PCP=1 | {0, 1} | 0 |
+| `[2, 3]` | PCP=0 | {2, 3} | 1 |
+| `[6, 7]` | PCP=1 | {2, 3} | 1 |
+
+```
+Optimized Step 1: Partial-TP All-Gather
+┌─────────────────────────────────────────────────────┐
+│  Ranks [0,1] gather within [0,1]: H/4 → H/2 heads   │
+│  Ranks [4,5] gather within [4,5]: H/4 → H/2 heads   │
+│  Ranks [2,3] gather within [2,3]: H/4 → H/2 heads   │
+│  Ranks [6,7] gather within [6,7]: H/4 → H/2 heads   │
+│  (Half the communication vs full TP all-gather!)   │
+└─────────────────────────────────────────────────────┘
+```
+
+---
+
+## Summary: When to Use Each Pattern
+
+| Condition | All-Gather | Reduce |
+|-----------|------------|--------|
+| `DCP == TP` and `PCP == 1` | Full TP | reduce-scatter |
+| `DCP < TP` and `PCP == 1` | Partial-TP (DCP group) | all-reduce + slice |
+| `DCP == TP * PCP` | Full TP | all-reduce + slice |
+| `DCP < TP * PCP` and `PCP > 1` | Partial-TP | all-reduce + slice |
+
+### Partial-TP Group Formula
+
+```python
+unique_tp_per_dcp = dcp_size // pcp_size
+num_dcp_groups = (tp_size * pcp_size) // dcp_size
+
+# For each DCP group i, for each PCP slice p:
+#   Partial-TP group = ranks at TP positions [i * unique_tp_per_dcp, (i+1) * unique_tp_per_dcp)
+#                      within PCP slice p
+```
+
+### Communication Savings
+
+| Config | Full TP All-Gather | Partial-TP All-Gather | Savings |
+|--------|-------------------|----------------------|---------|
+| PCP=1, TP=4, DCP=4 | H | H | 0% |
+| PCP=1, TP=4, DCP=2 | H | H/2 | 50% |
+| PCP=2, TP=2, DCP=2 | H | 0 (skip!) | 100% |
+| PCP=2, TP=4, DCP=4 | H | H/2 | 50% |

tests/distributed/test_context_parallel.py

@@ -50,7 +50,8 @@ class ParallelSetup(NamedTuple):
     tp_size: int
     pp_size: int
     dcp_size: int
-    cp_kv_cache_interleave_size: int
+    pcp_size: int
+    dcp_kv_cache_interleave_size: int
     eager_mode: bool
     chunked_prefill: bool
 
@@ -73,7 +74,8 @@ class CPTestSettings:
         tp_base: int = 4,
         pp_base: int = 1,
         dcp_multipliers: list[float] | None = None,
-        cp_kv_cache_interleave_size: int = 1,
+        pcp_base: int = 1,
+        dcp_kv_cache_interleave_size: int = 1,
         multi_node_only: bool = False,
         runner: RunnerOption = "auto",
         attn_backend: str | None = None,
@@ -91,8 +93,9 @@ class CPTestSettings:
                             ParallelSetup(
                                 tp_size=tp_base,
                                 pp_size=pp_multiplier * pp_base,
-                                dcp_size=int(dcp_multiplier * tp_base),
-                                cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
+                                dcp_size=max(1, int(dcp_multiplier * tp_base)),
+                                pcp_size=pcp_base,
+                                dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
                                 eager_mode=eager_mode_val,
                                 chunked_prefill=chunked_prefill_val,
                             )
@@ -126,16 +129,18 @@ CP_TEXT_GENERATION_MODELS = {
         CPTestSettings.detailed(dcp_multipliers=[1]),
         CPTestSettings.detailed(
             dcp_multipliers=[0.5],
-            cp_kv_cache_interleave_size=64,
+            dcp_kv_cache_interleave_size=64,
             attn_backend="FLASHMLA",
         ),
+        CPTestSettings.detailed(tp_base=1, pcp_base=4, dcp_kv_cache_interleave_size=64),
+        CPTestSettings.detailed(tp_base=2, pcp_base=2, dcp_kv_cache_interleave_size=64),
     ],
     "Qwen/Qwen2.5-1.5B-Instruct": [
         CPTestSettings.detailed(
-            cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"
+            dcp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"
         ),
         CPTestSettings.detailed(
-            cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
+            dcp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
         ),
     ],
 }
@@ -156,7 +161,8 @@ def _test_cp_gsm8k(
         tp_size,
         pp_size,
         dcp_size,
-        cp_kv_cache_interleave_size,
+        pcp_size,
+        dcp_kv_cache_interleave_size,
         eager_mode,
         chunked_prefill,
     ) = parallel_setup
@@ -212,8 +218,10 @@ def _test_cp_gsm8k(
             str(pp_size),
             "--decode-context-parallel-size",
             str(dcp_size),
+            "--prefill-context-parallel-size",
+            str(pcp_size),
             "--dcp-kv-cache-interleave-size",
-            str(cp_kv_cache_interleave_size),
+            str(dcp_kv_cache_interleave_size),
             "--distributed-executor-backend",
             distributed_backend,
         ]

tests/distributed/test_dcp_a2a.py

@@ -0,0 +1,192 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for DCP A2A communication backend (no GPU required).
+
+Tests cover:
+1. DCP A2A config validation (--dcp-comm-backend)
+2. KVP group function exists
+3. LSE-weighted combination correctness
+"""
+
+import math
+
+import pytest
+import torch
+
+from vllm.config.parallel import ParallelConfig
+
+
+class TestDCPCommBackendConfig:
+    """Test --dcp-comm-backend config validation."""
+
+    def test_default_is_ag_rs(self):
+        """Default comm backend is ag_rs."""
+        config = ParallelConfig()
+        assert config.dcp_comm_backend == "ag_rs"
+
+    def test_a2a_requires_dcp_greater_than_1(self):
+        """A2A backend requires decode_context_parallel_size > 1."""
+        with pytest.raises(
+            ValueError, match="requires decode_context_parallel_size > 1"
+        ):
+            ParallelConfig(
+                dcp_comm_backend="a2a",
+                decode_context_parallel_size=1,
+            )
+
+    def test_a2a_with_dcp_valid(self):
+        """A2A backend is valid when DCP > 1."""
+        config = ParallelConfig(
+            dcp_comm_backend="a2a",
+            tensor_parallel_size=8,
+            decode_context_parallel_size=4,
+        )
+        assert config.dcp_comm_backend == "a2a"
+
+    def test_invalid_backend_rejected(self):
+        """Invalid backend values are rejected."""
+        with pytest.raises(ValueError, match="must be one of"):
+            ParallelConfig(
+                dcp_comm_backend="invalid",
+            )
+
+    def test_ag_rs_with_dcp_1_valid(self):
+        """ag_rs backend is valid with DCP=1 (no DCP)."""
+        config = ParallelConfig(
+            dcp_comm_backend="ag_rs",
+            decode_context_parallel_size=1,
+        )
+        assert config.dcp_comm_backend == "ag_rs"
+
+
+class TestLSEWeightedCombine:
+    """Test LSE-weighted combination logic (CPU only, no GPU).
+
+    The _lse_weighted_combine function is the reference implementation
+    that verifies the Triton kernel's correctness. It computes:
+
+        result[b,h,d] = sum_n(w_n * output_n[b,h,d])
+
+    where w_n = softmax(lse_n) = exp(lse_n) / sum_k(exp(lse_k))
+    """
+
+    def test_importable(self):
+        """Verify _lse_weighted_combine is importable."""
+        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
+
+        assert callable(_lse_weighted_combine)
+
+    def test_single_rank(self):
+        """Single rank: output unchanged."""
+        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
+
+        # N=1, B=2, H=4, D=8
+        outputs = torch.randn(1, 2, 4, 8)
+        lses = torch.randn(1, 2, 4)
+
+        result = _lse_weighted_combine(outputs, lses)
+
+        assert result.shape == (2, 4, 8)
+        torch.testing.assert_close(result, outputs.squeeze(0), rtol=1e-5, atol=1e-5)
+
+    def test_equal_lse(self):
+        """Equal LSE values: outputs averaged equally."""
+        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
+
+        _N, B, H, D = 2, 1, 1, 4
+        outputs = torch.tensor(
+            [
+                [[[1.0, 2.0, 3.0, 4.0]]],  # Rank 0
+                [[[5.0, 6.0, 7.0, 8.0]]],  # Rank 1
+            ]
+        )
+        lses = torch.tensor(
+            [
+                [[0.0]],  # Rank 0
+                [[0.0]],  # Rank 1
+            ]
+        )
+
+        result = _lse_weighted_combine(outputs, lses)
+
+        expected = (outputs[0] + outputs[1]) / 2
+        assert result.shape == (B, H, D)
+        torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
+
+    def test_dominant_rank(self):
+        """Different LSE values: larger LSE gets more weight."""
+        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
+
+        B, H, D = 1, 1, 2
+        outputs = torch.tensor(
+            [
+                [[[0.0, 0.0]]],  # Rank 0
+                [[[1.0, 1.0]]],  # Rank 1
+            ]
+        )
+        lses = torch.tensor(
+            [
+                [[-100.0]],  # Rank 0: negligible contribution
+                [[0.0]],  # Rank 1: dominant
+            ]
+        )
+
+        result = _lse_weighted_combine(outputs, lses)
+
+        assert result.shape == (B, H, D)
+        torch.testing.assert_close(result, outputs[1].squeeze(0), atol=1e-5, rtol=1e-5)
+
+    def test_mathematically_correct(self):
+        """Verify mathematical correctness of LSE combination."""
+        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
+
+        outputs = torch.tensor(
+            [
+                [[[2.0, 4.0]]],
+                [[[6.0, 8.0]]],
+            ]
+        )
+        lses = torch.tensor(
+            [
+                [[1.0]],  # exp(1) ≈ 2.718
+                [[2.0]],  # exp(2) ≈ 7.389
+            ]
+        )
+
+        result = _lse_weighted_combine(outputs, lses)
+
+        w0 = math.exp(1) / (math.exp(1) + math.exp(2))
+        w1 = math.exp(2) / (math.exp(1) + math.exp(2))
+        expected = torch.tensor([[[w0 * 2.0 + w1 * 6.0, w0 * 4.0 + w1 * 8.0]]])
+
+        torch.testing.assert_close(result, expected, rtol=1e-4, atol=1e-4)
+
+    def test_return_lse(self):
+        """return_lse=True returns global LSE (logsumexp of inputs)."""
+        from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine
+
+        B, H, D = 1, 1, 2
+        outputs = torch.tensor(
+            [
+                [[[1.0, 2.0]]],
+                [[[3.0, 4.0]]],
+            ]
+        )
+        lses = torch.tensor(
+            [
+                [[1.0]],
+                [[2.0]],
+            ]
+        )
+
+        result, global_lse = _lse_weighted_combine(outputs, lses, return_lse=True)
+
+        expected_global_lse = math.log(math.exp(1) + math.exp(2))
+
+        assert result.shape == (B, H, D)
+        assert global_lse.shape == (B, H)
+        assert abs(global_lse.item() - expected_global_lse) < 1e-5
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/quantization/test_torchao.py

@@ -20,7 +20,7 @@ TORCHAO_AVAILABLE = importlib.util.find_spec("torchao") is not None
 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
 def test_pre_quantized_model(vllm_runner):
     with vllm_runner(
-        "torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.15.0",
+        "drisspg/fp8-opt-125m",
         quantization="torchao",
         dtype="bfloat16",
         enforce_eager=True,
@@ -52,6 +52,22 @@ def test_opt_125m_int8wo_model_loading_with_params(vllm_runner, pt_load_map_loca
         assert output
 
 
+@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
+def test_opt_125m_int4wo_model_per_module_quant(vllm_runner):
+    torch._dynamo.reset()
+    model_name = "jerryzh168/opt-125m-int4wo-per-module"
+    with vllm_runner(
+        model_name=model_name,
+        quantization="torchao",
+        dtype="bfloat16",
+        pt_load_map_location="cuda:0",
+        enforce_eager=True,
+    ) as llm:
+        output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
+
+        assert output
+
+
 @pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
 def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
     torch._dynamo.reset()

tests/v1/attention/test_dcp_ops.py

@@ -0,0 +1,170 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for DCP (Decode Context Parallelism) operations."""
+
+from unittest.mock import patch
+
+import pytest
+import torch
+
+
+class MockGroupCoordinator:
+    """Mock GroupCoordinator for testing DCP functions without distributed setup."""
+
+    def __init__(self, world_size: int, rank: int):
+        self.world_size = world_size
+        self.rank_in_group = rank
+
+    def all_gather(self, tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        return tensor.repeat_interleave(self.world_size, dim=dim)
+
+    def reduce_scatter(self, tensor: torch.Tensor, dim: int = -1) -> torch.Tensor:
+        size = tensor.size(dim) // self.world_size
+        start = self.rank_in_group * size
+        return torch.narrow(tensor, dim, start, size)
+
+    def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
+        return tensor
+
+
+def _patch_groups(tp_ws, tp_rank, dcp_ws, dcp_rank, pcp_ws=1, pcp_rank=0):
+    """Context manager that patches TP, DCP and PCP group getters."""
+    tp = MockGroupCoordinator(tp_ws, tp_rank)
+    dcp = MockGroupCoordinator(dcp_ws, dcp_rank)
+    pcp = MockGroupCoordinator(pcp_ws, pcp_rank)
+    return (
+        (
+            patch("vllm.v1.attention.ops.common.get_tp_group", return_value=tp),
+            patch("vllm.v1.attention.ops.common.get_dcp_group", return_value=dcp),
+            patch("vllm.v1.attention.ops.common.get_pcp_group", return_value=pcp),
+        ),
+        tp,
+        dcp,
+    )
+
+
+@pytest.fixture
+def device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+@pytest.fixture
+def mock_groups():
+    """PCP=1, TP=DCP=2 mock groups."""
+    patches, tp, dcp = _patch_groups(tp_ws=2, tp_rank=0, dcp_ws=2, dcp_rank=0)
+    with patches[0], patches[1], patches[2]:
+        yield tp, dcp
+
+
+class TestDCPPrepareQuery:
+    def test_basic_shape(self, device, mock_groups):
+        from vllm.v1.attention.ops.common import dcp_prepare_query
+
+        tp_group, _ = mock_groups
+        B, H_local, D = 2, 4, 64
+        query = torch.randn(B, H_local, D, device=device)
+        result = dcp_prepare_query(query)
+        assert result.shape == (B, H_local * tp_group.world_size, D)
+
+    def test_single_rank_passthrough(self, device):
+        from vllm.v1.attention.ops.common import dcp_prepare_query
+
+        patches, _, _ = _patch_groups(1, 0, 1, 0)
+        with patches[0], patches[1], patches[2]:
+            B, H, D = 2, 8, 64
+            query = torch.randn(B, H, D, device=device)
+            result = dcp_prepare_query(query)
+            assert result.shape == query.shape
+            torch.testing.assert_close(result, query)
+
+
+class TestDCPReduceOutput:
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_basic_shape(self, device, mock_groups):
+        from vllm.v1.attention.ops.common import dcp_reduce_output
+
+        tp_group, _ = mock_groups
+        B, H_total, D = 2, 8, 64
+        attn_output = torch.randn(B, H_total, D, device=device)
+        attn_lse = torch.randn(B, H_total, device=device)
+        result = dcp_reduce_output(attn_output, attn_lse)
+        assert result.shape == (B, H_total // tp_group.world_size, D)
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_single_rank_both_groups(self, device):
+        from vllm.v1.attention.ops.common import dcp_reduce_output
+
+        patches, _, _ = _patch_groups(1, 0, 1, 0)
+        with patches[0], patches[1], patches[2]:
+            B, H, D = 2, 8, 64
+            attn_output = torch.randn(B, H, D, device=device)
+            attn_lse = torch.randn(B, H, device=device)
+            result = dcp_reduce_output(attn_output, attn_lse)
+            assert result.shape == attn_output.shape
+
+
+class TestEndToEndDCPFlow:
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_prepare_and_reduce_round_trip(self, device, mock_groups):
+        from vllm.v1.attention.ops.common import dcp_prepare_query, dcp_reduce_output
+
+        tp_group, _ = mock_groups
+        B, H_local, D = 4, 4, 64
+        query = torch.randn(B, H_local, D, device=device)
+        query_all_heads = dcp_prepare_query(query)
+        assert query_all_heads.shape == (B, H_local * tp_group.world_size, D)
+
+        attn_output = torch.randn_like(query_all_heads)
+        attn_lse = torch.randn(B, H_local * tp_group.world_size, device=device)
+        final_output = dcp_reduce_output(attn_output, attn_lse)
+        assert final_output.shape == query.shape
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_case1_dcp_equals_pcp(self, device):
+        """Case 1: DCP = PCP (TP=4, DCP=2, PCP=2) - no all-gather, all-reduce only."""
+        from vllm.v1.attention.ops.common import dcp_prepare_query, dcp_reduce_output
+
+        # DCP = PCP = 2, so no all-gather needed (same TP heads)
+        patches, tp, _ = _patch_groups(
+            tp_ws=4, tp_rank=0, dcp_ws=2, dcp_rank=0, pcp_ws=2, pcp_rank=0
+        )
+        with patches[0], patches[1], patches[2]:
+            B, H_local, D = 2, 2, 64
+            query = torch.randn(B, H_local, D, device=device)
+            query_prepared = dcp_prepare_query(query)
+            # Case 1: No all-gather, query unchanged
+            assert query_prepared.shape == query.shape
+
+            attn_output = torch.randn_like(query_prepared)
+            attn_lse = torch.randn(B, H_local, device=device)
+            final_output = dcp_reduce_output(attn_output, attn_lse)
+            # Case 1: All-reduce only, shape unchanged
+            assert final_output.shape == query.shape
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="Requires CUDA")
+    def test_case2_dcp_equals_tp_times_pcp(self, device):
+        """Case 2: DCP = TP × PCP (TP=2, DCP=4, PCP=2) - all-gather + reduce-scatter."""
+        from vllm.v1.attention.ops.common import dcp_prepare_query, dcp_reduce_output
+
+        # DCP = TP × PCP = 4, need all-gather and reduce-scatter
+        patches, tp, dcp = _patch_groups(
+            tp_ws=2, tp_rank=0, dcp_ws=4, dcp_rank=0, pcp_ws=2, pcp_rank=0
+        )
+        with patches[0], patches[1], patches[2]:
+            B, H_local, D = 2, 2, 64
+            query = torch.randn(B, H_local, D, device=device)
+            query_all_heads = dcp_prepare_query(query)
+            # Case 2: All-gather across TP
+            assert query_all_heads.shape == (B, H_local * tp.world_size, D)
+
+            attn_output = torch.randn_like(query_all_heads)
+            attn_lse = torch.randn(B, H_local * tp.world_size, device=device)
+            final_output = dcp_reduce_output(attn_output, attn_lse)
+            # Case 2: Reduce-scatter back to local heads
+            assert final_output.shape == query.shape
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/v1/worker/test_gpu_model_runner.py

@@ -972,7 +972,7 @@ def test_hybrid_block_table_initialization():
     max_num_reqs = 10
     max_num_blocks_per_req = 20
     max_num_batched_tokens = 512
-    cp_kv_cache_interleave_size = 8
+    dcp_kv_cache_interleave_size = 8
 
     block_table = BlockTable(
         block_size=block_size,
@@ -982,7 +982,7 @@ def test_hybrid_block_table_initialization():
         pin_memory=False,
         device=torch.device(DEVICE),
         kernel_block_size=kernel_block_sizes[0],
-        cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
+        dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
     )
 
     # Verify hybrid block configuration

vllm/config/parallel.py

@@ -36,6 +36,7 @@ ExpertPlacementStrategy = Literal["linear", "round_robin"]
 DistributedExecutorBackend = Literal["ray", "mp", "uni", "external_launcher"]
 DataParallelBackend = Literal["ray", "mp"]
 EPLBPolicyOption = Literal["default"]
+DCPCommBackend = Literal["ag_rs", "a2a"]
 All2AllBackend = Literal[
     "naive",
     "pplx",
@@ -287,6 +288,14 @@ class ParallelConfig:
     and will be deprecated when PCP is fully supported.
 
     """
+    dcp_comm_backend: DCPCommBackend = "ag_rs"
+    """Communication backend for Decode Context Parallel (DCP).
+    - "ag_rs": AllGather + ReduceScatter (default, existing behavior)
+    - "a2a": All-to-All exchange of partial outputs + LSE, then
+      combine with Triton kernel. Reduces NCCL calls from 3 to 2
+      per layer for MLA models.
+    """
+
     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`,
@@ -294,12 +303,11 @@ class ParallelConfig:
     store interleave_size tokens on total_cp_rank i,
     then store next interleave_size tokens on total_cp_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.
-    Block_size should be greater than or equal to cp_kv_cache_interleave_size.
-    Block_size should be divisible by cp_kv_cache_interleave_size.
+    Block_size should be >= dcp_kv_cache_interleave_size and divisible by it.
     """
 
     data_parallel_index: int = Field(init=False)
@@ -381,15 +389,32 @@ 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:
+        # DCP configuration with PCP is restricted to two clean cases:
+        #
+        # Case 1: DCP = PCP (e.g., PCP=2, TP=2, DCP=2)
+        #   - DCP groups ranks at same TP position, different PCP slices
+        #   - No Q all-gather needed (same heads)
+        #   - All-reduce across DCP to combine KV slices
+        #
+        # Case 2: DCP = TP × PCP (e.g., PCP=2, TP=2, DCP=4)
+        #   - DCP group contains all ranks
+        #   - Full TP all-gather for Q
+        #   - Reduce-scatter across DCP
+        tp = self.tensor_parallel_size
+        dcp = self.decode_context_parallel_size
+        pcp = self.prefill_context_parallel_size
+        if dcp > 1 and pcp > 1:
+            valid_dcp_sizes = {pcp, tp * pcp}
+            if dcp not in valid_dcp_sizes:
+                raise ValueError(
+                    f"When PCP > 1, DCP must be either PCP ({pcp}) or "
+                    f"TP×PCP ({tp * pcp}), but got DCP={dcp}. "
+                    f"Valid options: {valid_dcp_sizes}"
+                )
+
+        if self.dcp_comm_backend == "a2a" and self.decode_context_parallel_size <= 1:
             raise ValueError(
-                f"tp_size={self.tensor_parallel_size} must be divisible by"
-                f"dcp_size={self.decode_context_parallel_size}."
+                "dcp_comm_backend='a2a' requires decode_context_parallel_size > 1."
             )
 
         return self

vllm/config/vllm.py

@@ -992,30 +992,27 @@ class VllmConfig:
             )
         current_platform.check_and_update_config(self)
 
-        # If DCP, ensure the block size is right.
+        # If DCP, ensure the block size is compatible with interleave size.
         if self.parallel_config.decode_context_parallel_size > 1:
-            if self.parallel_config.dcp_kv_cache_interleave_size > 1 and (
-                self.parallel_config.cp_kv_cache_interleave_size
-                != self.parallel_config.dcp_kv_cache_interleave_size
+            # Migrate from deprecated cp_kv_cache_interleave_size
+            if self.parallel_config.cp_kv_cache_interleave_size > 1 and (
+                self.parallel_config.dcp_kv_cache_interleave_size
+                != self.parallel_config.cp_kv_cache_interleave_size
             ):
-                self.parallel_config.cp_kv_cache_interleave_size = (
-                    self.parallel_config.dcp_kv_cache_interleave_size
+                self.parallel_config.dcp_kv_cache_interleave_size = (
+                    self.parallel_config.cp_kv_cache_interleave_size
                 )
                 logger.warning_once(
-                    "cp_kv_cache_interleave_size is overridden by dcp_kv_cache"
-                    "_interleave_size. And dcp-kv-cache-interleave-size will be "
-                    "deprecated when PCP is fully supported."
+                    "cp_kv_cache_interleave_size is deprecated. "
+                    "Use dcp_kv_cache_interleave_size instead."
                 )
+            interleave = self.parallel_config.dcp_kv_cache_interleave_size
             assert (
-                self.parallel_config.cp_kv_cache_interleave_size
-                <= self.cache_config.block_size
-                and self.cache_config.block_size
-                % self.parallel_config.cp_kv_cache_interleave_size
-                == 0
+                interleave <= self.cache_config.block_size
+                and self.cache_config.block_size % interleave == 0
             ), (
-                f"Block_size({self.cache_config.block_size}) should be greater "
-                "than or equal to and divisible by cp_kv_cache_interleave_size "
-                f"({self.parallel_config.cp_kv_cache_interleave_size})."
+                f"block_size ({self.cache_config.block_size}) must be >= and "
+                f"divisible by dcp_kv_cache_interleave_size ({interleave})."
             )
 
         # Do this after all the updates to compilation_config.mode
@@ -1618,6 +1615,8 @@ class VllmConfig:
             f"tensor_parallel_size={self.parallel_config.tensor_parallel_size}, "  # noqa
             f"pipeline_parallel_size={self.parallel_config.pipeline_parallel_size}, "  # noqa
             f"data_parallel_size={self.parallel_config.data_parallel_size}, "  # noqa
+            f"decode_context_parallel_size={self.parallel_config.decode_context_parallel_size}, "  # noqa
+            f"dcp_comm_backend={self.parallel_config.dcp_comm_backend}, "  # noqa
             f"disable_custom_all_reduce={self.parallel_config.disable_custom_all_reduce}, "  # noqa
             f"quantization={self.model_config.quantization}, "
             f"enforce_eager={self.model_config.enforce_eager}, "

vllm/distributed/parallel_state.py

@@ -33,7 +33,7 @@ from contextlib import contextmanager, nullcontext
 from dataclasses import dataclass
 from datetime import timedelta
 from multiprocessing import shared_memory
-from typing import TYPE_CHECKING, Any, Protocol
+from typing import TYPE_CHECKING, Any, Optional, Protocol
 from unittest.mock import patch
 
 import torch
@@ -117,7 +117,7 @@ def _get_unique_name(name: str) -> str:
     return newname
 
 
-_groups: dict[str, Callable[[], "GroupCoordinator | None"]] = {}
+_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}
 
 
 def _register_group(group: "GroupCoordinator") -> None:
@@ -809,7 +809,7 @@ class GroupCoordinator:
         self,
         tensor_dict: dict[str, torch.Tensor | Any],
         dst: int | None = None,
-        all_gather_group: "GroupCoordinator | None" = None,
+        all_gather_group: Optional["GroupCoordinator"] = None,
         all_gather_tensors: dict[str, bool] | None = None,
     ) -> dict[str, torch.Tensor | Any] | None:
         """Send the input tensor dictionary.
@@ -903,7 +903,7 @@ class GroupCoordinator:
     def recv_tensor_dict(
         self,
         src: int | None = None,
-        all_gather_group: "GroupCoordinator | None" = None,
+        all_gather_group: Optional["GroupCoordinator"] = None,
         all_gather_tensors: dict[str, bool] | None = None,
     ) -> dict[str, torch.Tensor | Any] | None:
         """Recv the input tensor dictionary.
@@ -1225,9 +1225,6 @@ def get_dcp_group() -> GroupCoordinator:
     return _DCP
 
 
-# kept for backward compatibility
-get_context_model_parallel_group = get_dcp_group
-
 _PP: GroupCoordinator | None = None
 
 
@@ -1571,11 +1568,17 @@ 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 = decode_context_model_parallel_size or 1
+    if dcp > 1:
+        # DCP spans PCP dimension first, then TP dimension.
+        # E.g. tp=2, pcp=2: layout is [[0,1], [2,3]] (pcp x tp)
+        #      dcp=2: groups [0,2], [1,3] (same TP, span PCP)
+        #      dcp=4: group [0,2,1,3] (span both)
+        # Transpose to (tp, pcp) so PCP is innermost, then reshape.
+        r = all_ranks.transpose(-1, -2)
+        group_ranks = r.reshape(-1, dcp).unbind(0)
+    else:
+        group_ranks = all_ranks.reshape(-1, 1).unbind(0)
     group_ranks = [x.tolist() for x in group_ranks]
     if enable_elastic_ep:
         group_ranks = local_all_ranks.reshape(
@@ -1592,6 +1595,8 @@ def initialize_model_parallel(
 
     global _PCP
     assert _PCP is None, "prefill context parallel group is already initialized"
+    # PCP groups are essentially TP-sized groups across PCP dimension.
+    # For tp=6, pcp=2: PCP groups are [0,1,2,3,4,5] and [6,7,8,9,10,11]
     group_ranks = (
         all_ranks.transpose(3, 4)
         .reshape(-1, prefill_context_model_parallel_size)
@@ -1651,14 +1656,17 @@ def initialize_model_parallel(
     global _EP
     assert _EP is None, "expert parallel group is already initialized"
     # Don't create EP group for dense models.
-    if config.model_config is None or config.model_config.is_moe:
+    if config is None or config.model_config is None or config.model_config.is_moe:
+        # EP groups span DP and TP but NOT PCP. PCP ranks should have
+        # independent EP groups since they process different token chunks
+        # and run MoE all2all independently.
+        # Layout after transposes: (DCP_remain, PP, PCP, DP, TP)
         group_ranks = (
             all_ranks.transpose(1, 2)
+            .transpose(2, 3)
             .reshape(
                 -1,
-                data_parallel_size
-                * prefill_context_model_parallel_size
-                * tensor_model_parallel_size,
+                data_parallel_size * tensor_model_parallel_size,
             )
             .unbind(0)
         )

vllm/engine/arg_utils.py

@@ -85,6 +85,7 @@ from vllm.config.observability import DetailedTraceModules
 from vllm.config.parallel import (
     All2AllBackend,
     DataParallelBackend,
+    DCPCommBackend,
     DistributedExecutorBackend,
     ExpertPlacementStrategy,
 )
@@ -405,6 +406,7 @@ class EngineArgs:
     tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
     prefill_context_parallel_size: int = ParallelConfig.prefill_context_parallel_size
     decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
+    dcp_comm_backend: DCPCommBackend = ParallelConfig.dcp_comm_backend
     dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
     cp_kv_cache_interleave_size: int = ParallelConfig.cp_kv_cache_interleave_size
     data_parallel_size: int = ParallelConfig.data_parallel_size
@@ -820,6 +822,10 @@ class EngineArgs:
             "-dcp",
             **parallel_kwargs["decode_context_parallel_size"],
         )
+        parallel_group.add_argument(
+            "--dcp-comm-backend",
+            **parallel_kwargs["dcp_comm_backend"],
+        )
         parallel_group.add_argument(
             "--dcp-kv-cache-interleave-size",
             **parallel_kwargs["dcp_kv_cache_interleave_size"],
@@ -1720,6 +1726,7 @@ class EngineArgs:
             worker_cls=self.worker_cls,
             worker_extension_cls=self.worker_extension_cls,
             decode_context_parallel_size=self.decode_context_parallel_size,
+            dcp_comm_backend=self.dcp_comm_backend,
             dcp_kv_cache_interleave_size=self.dcp_kv_cache_interleave_size,
             cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
             _api_process_count=self._api_process_count,

vllm/model_executor/layers/attention/mla_attention.py

@@ -204,7 +204,11 @@ import vllm.envs as envs
 from vllm import _custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
-from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
+from vllm.distributed.parallel_state import (
+    get_dcp_group,
+    get_pcp_group,
+    is_global_first_rank,
+)
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp
@@ -247,12 +251,19 @@ from vllm.v1.attention.backend import (
 )
 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
 from vllm.v1.attention.backends.utils import (
+    fused_pcp_qkv_select,
     get_dcp_local_seq_lens,
+    get_pcp_query_restore_idx,
     get_per_layer_parameters,
     infer_global_hyperparameters,
+    pcp_kv_allgather_and_restore,
     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_rs,
+    dcp_prepare_query,
+)
+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
 from vllm.v1.kv_cache_interface import (
@@ -393,6 +404,12 @@ class MLAAttention(nn.Module, AttentionLayerBase):
 
         self.use_sparse = use_sparse
 
+        parallel_config = get_current_vllm_config().parallel_config
+        self.dcp_a2a = (
+            parallel_config.decode_context_parallel_size > 1
+            and parallel_config.dcp_comm_backend == "a2a"
+        )
+
         # Initialize q/k/v range constants.
         self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
         self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
@@ -508,6 +525,12 @@ class MLAAttention(nn.Module, AttentionLayerBase):
 
         if self.impl.dcp_world_size == -1:
             self.impl.dcp_world_size = get_dcp_group().world_size
+        if self.impl.pcp_world_size == -1:
+            self.impl.pcp_world_size = get_pcp_group().world_size
+        if self.impl.dcp_rank == -1:
+            self.impl.dcp_rank = get_dcp_group().rank_in_group
+        if self.impl.pcp_rank == -1:
+            self.impl.pcp_rank = get_pcp_group().rank_in_group
 
         fp8_attention = self.kv_cache_dtype.startswith("fp8")
 
@@ -520,6 +543,16 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         k_c_normed = k_c_normed[:num_actual_toks, ...]
         k_pe = k_pe[:num_actual_toks, ...]
 
+        if self.impl.pcp_world_size > 1:
+            assert attn_metadata.pcp_allgather_restore_idx is not None
+            k_c_normed, k_pe = pcp_kv_allgather_and_restore(
+                k_c_normed,
+                k_pe,
+                num_actual_toks,
+                attn_metadata.pcp_allgather_restore_idx,
+                get_pcp_group(),
+            )
+
         # write the latent and rope to kv cache
         if kv_cache.numel() > 0:
             ops.concat_and_cache_mla(
@@ -550,10 +583,13 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             num_mha_tokens = q.size(0) - num_mqa_tokens
 
         if num_mha_tokens > 0:
+            # After PCP all-gather, k_c_normed/k_pe have pcp_world_size copies of
+            # decode tokens, so skip num_mqa_tokens * pcp_world_size
+            kv_skip = num_mqa_tokens * self.impl.pcp_world_size
             self.impl.forward_mha(
                 q[num_mqa_tokens:],
-                k_c_normed[num_mqa_tokens:],
-                k_pe[num_mqa_tokens:],
+                k_c_normed[kv_skip:],
+                k_pe[kv_skip:],
                 kv_cache,
                 attn_metadata,
                 self._k_scale,
@@ -627,22 +663,32 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 assert not fp8_attention, "DCP not support fp8 kvcache now."
                 # 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)
+                # mqa_q do allgather in head dim across TP.
+                mqa_q = dcp_prepare_query(mqa_q)
 
             # call decode attn
             if not is_sparse_impl:
                 assert attn_metadata.decode is not None
             attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)
 
-            # correct dcp attn_out with lse.
+            # Only DCP shards KV cache. With PCP-only, each rank has the
+            # full KV cache during decode (gathered after prefill), so no
+            # collective needed.
             if self.impl.dcp_world_size > 1:
-                attn_out = cp_lse_ag_out_rs(
-                    attn_out,
-                    lse,
-                    get_dcp_group(),
-                    is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
-                )
+                if self.dcp_a2a:
+                    attn_out = dcp_a2a_lse_reduce(
+                        attn_out,
+                        lse,
+                        get_dcp_group(),
+                        is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
+                    )
+                else:
+                    attn_out = cp_lse_ag_out_rs(
+                        attn_out,
+                        lse,
+                        get_dcp_group(),
+                        is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),
+                    )
 
             # v_up projection
             self._v_up_proj(attn_out, out=mqa_output_slice)
@@ -1049,6 +1095,19 @@ class MLACommonPrefillMetadata:
         cu_seq_lens_lst: list[list[int]] | None = None
         chunk_size: int | None = None
 
+    @dataclass
+    class PCPMetadata:
+        @dataclass
+        class ChunkMetadata:
+            cu_seqlens_q: torch.Tensor
+            cu_seqlens_k: torch.Tensor
+            max_seqlen_q: int
+            max_seqlen_k: int
+
+        output_restore_idx: torch.Tensor
+        head: "MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata"
+        tail: "MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata"
+
     block_table: torch.Tensor
     query_start_loc: torch.Tensor
     max_query_len: int
@@ -1057,6 +1116,12 @@ class MLACommonPrefillMetadata:
     workspace_buffer: torch.Tensor | None = None
     q_data_type: torch.dtype | None = None
     output_dtype: torch.dtype | None = None
+    pcp_metadata: PCPMetadata | None = None
+
+
+PrefillKernelMetadata = (
+    MLACommonPrefillMetadata | MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata
+)
 
 
 @dataclass
@@ -1126,6 +1191,8 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]):
         | None
     ) = None
 
+    pcp_allgather_restore_idx: torch.Tensor | None = None
+
     def __post_init__(self):
         if self.head_dim is not None and not MLACommonBackend.supports_head_size(
             self.head_dim
@@ -1361,9 +1428,19 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             # 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
+        try:
+            self.pcp_world_size = get_pcp_group().world_size
+            self.pcp_rank = get_pcp_group().rank_in_group
+        except AssertionError:
+            # PCP might not be initialized in testing
+            self.pcp_world_size = 1
+            self.pcp_rank = 0
+        # DCP groups span PCP, so dcp_world_size is the effective CP world size.
+        self.dcp_local_block_size = parallel_config.dcp_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
+        self.dcp_kv_cache_interleave_size = parallel_config.dcp_kv_cache_interleave_size
+        # TODO(yyj) Remove this once the PCP bug for decode_length > 1 is fixed.
+        supports_dcp_with_varlen = supports_dcp_with_varlen and self.pcp_world_size == 1
 
         # Don't try to access the runner on AMD
         if self.aot_schedule:
@@ -1373,10 +1450,12 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             self.determine_chunked_prefill_workspace_size(vllm_config)
         )
 
+        # Only DCP shards KV cache, affecting workspace sizing.
+        # PCP gathers K/V after prefill so each rank has full sequence.
         if self.dcp_world_size > 1:
-            # Note(hc): The local kvcache is incomplete when DCP is triggered,
-            # an additional kvcache allgather across the DCP group is therefore
-            # required, so the workspace has to be enlarged by 1/DCP relative
+            # Note(hc): The local kvcache is incomplete when DCP or PCP is triggered,
+            # an additional kvcache allgather across the DCP&PCP group is therefore
+            # required, so the workspace has to be enlarged by 1/CP relative
             # to the original TP allocation.
             assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0
             self.chunked_prefill_workspace = torch.empty(
@@ -1574,6 +1653,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
         num_tokens = common_attn_metadata.num_actual_tokens
         max_query_len = common_attn_metadata.max_query_len
         max_seq_len = common_attn_metadata.max_seq_len
+        pcp_allgather_restore_idx = common_attn_metadata.pcp_allgather_restore_idx
 
         # Note(simon): be careful about the CPU <> GPU memory movement in this
         # function. We should avoid GPU -> CPU sync as much as possible because
@@ -1612,6 +1692,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             prefill_query_start_loc = (
                 query_start_loc[reqs_start:] - query_start_loc[reqs_start]
             )
+            prefill_query_start_loc_cpu = (
+                query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start]
+            )
 
             chunked_context_metadata = None
             if max_context_len_cpu > 0:
@@ -1775,6 +1858,35 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                     <= self.chunked_prefill_workspace_size
                 )
 
+            pcp_metadata = None
+            if self.pcp_world_size > 1:
+                output_res_idx = get_pcp_query_restore_idx(prefill_query_start_loc_cpu)
+                pcp_query_start_loc = prefill_query_start_loc // 2
+                max_query_len_half = max_query_len // 2
+
+                head_chunk = MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata(
+                    cu_seqlens_q=pcp_query_start_loc,
+                    cu_seqlens_k=pcp_query_start_loc * (self.pcp_rank + 1),
+                    max_seqlen_q=max_query_len_half,
+                    max_seqlen_k=max_query_len_half * (self.pcp_rank + 1),
+                )
+                tail_chunk = MLACommonPrefillMetadata.PCPMetadata.ChunkMetadata(
+                    cu_seqlens_q=pcp_query_start_loc,
+                    cu_seqlens_k=pcp_query_start_loc
+                    * (self.pcp_world_size * 2 - self.pcp_rank),
+                    max_seqlen_q=max_query_len_half,
+                    max_seqlen_k=max_query_len_half
+                    * (self.pcp_world_size * 2 - self.pcp_rank),
+                )
+
+                pcp_metadata = MLACommonPrefillMetadata.PCPMetadata(
+                    output_restore_idx=output_res_idx.to(
+                        device, dtype=torch.int32, non_blocking=True
+                    ),
+                    head=head_chunk,
+                    tail=tail_chunk,
+                )
+
             prefill_metadata = self.prefill_metadata_cls(
                 block_table=block_table_tensor[reqs_start:, ...],
                 query_start_loc=prefill_query_start_loc,
@@ -1782,6 +1894,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                 chunked_context=chunked_context_metadata,
                 output_dtype=self.model_config.dtype,
                 q_data_type=self.q_data_type,
+                pcp_metadata=pcp_metadata,
             )
 
             if self._use_cudnn_prefill:
@@ -1804,15 +1917,15 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                 dcp_tot_seq_lens_device = seq_lens[:num_decodes]
                 seq_lens = dcp_local_seq_lens
 
-                # After DCP distribution, the maximum number of tokens for any rank is
+                # After CP distribution, the maximum number of tokens for any rank is
                 # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
-                # and I is cp_kv_cache_interleave_size.
+                # and I is dcp_kv_cache_interleave_size.
                 # This eliminates GPU->CPU sync while minimizing workspace
                 # over-allocation.
-                num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
+                num_partitions = self.dcp_world_size * self.dcp_kv_cache_interleave_size
                 max_seq_len = (
                     (max_seq_len + num_partitions - 1) // num_partitions
-                ) * self.cp_kv_cache_interleave_size
+                ) * self.dcp_kv_cache_interleave_size
 
             decode_metadata = self._build_decode(
                 block_table_tensor=block_table_tensor[:num_decodes, ...],
@@ -1838,6 +1951,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             num_prefills=num_prefills,
             prefill=prefill_metadata,
             decode=decode_metadata,
+            pcp_allgather_restore_idx=pcp_allgather_restore_idx,
         )
 
         if self._use_fi_prefill and num_prefills > 0:
@@ -1859,7 +1973,7 @@ def reorg_kvcache(
     toks: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
-    reorg and unpad kvcache after cp local gather to tp layout for attn kernel.
+    reorg and unpad kvcache after dcp local gather to tp layout for attn kernel.
     e.g.
     allgatered_kv_c_normed = [T0_0, T0_1, T0_2, T0_3, T1_0, T1_1, ...,
                               T0_4, T0_5, pad, pad, T1_2, pad, ...]
@@ -1867,10 +1981,10 @@ def reorg_kvcache(
                                   T1_0, T1_1, T1_2, ...]
     Args:
         padded_local_chunk_seq_lens_lst: local chunk context lengths
-            under current CP rank.
-        local_context_lens_allranks: local context lengths on each CP rank.
-        sum_seq_len: the sum of cp_chunk_seq_lens_lst.
-        max_seq_len: the max value of cp_chunk_seq_lens_lst.
+            under current DCP rank.
+        local_context_lens_allranks: local context lengths on each DCP rank.
+        sum_seq_len: the sum of dcp_chunk_seq_lens_lst.
+        max_seq_len: the max value of dcp_chunk_seq_lens_lst.
         chunk_size: the local padded max context chunk from
             chunked_context_metadata building.
         chunk_idx: chunk idx of chunked_prefill.
@@ -2034,9 +2148,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             )
 
         self.dcp_world_size: int = -1
+        self.pcp_world_size: int = -1
+        self.dcp_rank: int = -1
+        self.pcp_rank: int = -1
 
-        self.cp_kv_cache_interleave_size: int = (
-            get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
+        self.dcp_kv_cache_interleave_size: int = (
+            get_current_vllm_config().parallel_config.dcp_kv_cache_interleave_size
         )
 
     def _flash_attn_varlen_diff_headdims(
@@ -2076,27 +2193,117 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             return attn_out, lse
         return attn_out
 
-    def _run_prefill_new_tokens_fa(
+    def _run_prefill_new_tokens_pcp(
         self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
     ):
+        """Run prefill attention with PCP (Prefill Context Parallelism) support.
+
+        This method handles PCP by splitting QKV into head/tail chunks,
+        running attention on each, then merging and restoring order.
+
+        NOTE: Only call this when pcp_world_size > 1.
+        """
+        assert self.pcp_world_size > 1
+        assert self.pcp_rank != -1
+
+        # NOTE When PCP is enabled, we split the queries keys and values into
+        # "head" and "tail" parts using the DualChunkSwap strategy to balance
+        # workload across PCP ranks. We run attention twice (once for the head
+        # part and once for the tail part), then concatenate the results and
+        # restore the original ordering.
+        #
+        # Considering pcp_world_size=2 and sequence is [0,1,2,3,4,5,6,7]
+        #
+        #   pcp_rank0: Q [0,1,6,7] KV [0,1,2,3,4,5,6,7]
+        #    Q\KV  0 1 2 3 4 5 6 7
+        # head 0   1 0 0 0 0 0 0 0
+        #      1   1 1 0 0 0 0 0 0
+        #      -------------------
+        # tail 6   1 1 1 1 1 1 1 0
+        #      7   1 1 1 1 1 1 1 1
+        #
+        #   pcp_rank1: Q[2,3,4,5] KV [0,1,2,3,4,5,6,7]
+        #    Q\KV  0 1 2 3 4 5 6 7
+        # head 2   1 1 1 0 0 0 0 0
+        #      3   1 1 1 1 0 0 0 0
+        #      -------------------
+        # tail 4   1 1 1 1 1 0 0 0
+        #      5   1 1 1 1 1 1 0 0
+
+        q_head, k_head, v_head, q_tail, k_tail, v_tail = fused_pcp_qkv_select(
+            q=q,
+            k=k,
+            v=v,
+            query_start_loc=prefill.query_start_loc,
+            pcp_rank=self.pcp_rank,
+            pcp_world_size=self.pcp_world_size,
+        )
+
+        pcp_metadata = prefill.pcp_metadata
+        assert pcp_metadata is not None
+
+        output_head, lse_head = self._run_prefill_new_tokens(
+            q=q_head,
+            k=k_head,
+            v=v_head,
+            prefill=pcp_metadata.head,
+            return_softmax_lse=True,
+        )
+
+        output_tail, lse_tail = self._run_prefill_new_tokens(
+            q=q_tail,
+            k=k_tail,
+            v=v_tail,
+            prefill=pcp_metadata.tail,
+            return_softmax_lse=True,
+        )
+
+        output = torch.cat([output_head, output_tail], dim=0)
+        output_restore_idx = pcp_metadata.output_restore_idx
+        if return_softmax_lse:
+            # FA returns LSE in shape [ H, B ]
+            lse = torch.cat([lse_head, lse_tail], dim=-1)
+            return (
+                torch.index_select(output, 0, output_restore_idx),
+                torch.index_select(lse, -1, output_restore_idx),
+            )
+        else:
+            return torch.index_select(output, 0, output_restore_idx)
+
+    def _run_prefill_new_tokens_fa(
+        self,
+        q,
+        k,
+        v,
+        prefill: PrefillKernelMetadata,
+        return_softmax_lse: bool,
+    ):
+        if isinstance(prefill, MLACommonPrefillMetadata):
+            cu_seqlens_q = cu_seqlens_k = prefill.query_start_loc
+            max_seqlen_q = max_seqlen_k = prefill.max_query_len
+        else:
+            cu_seqlens_q, cu_seqlens_k = prefill.cu_seqlens_q, prefill.cu_seqlens_k
+            max_seqlen_q, max_seqlen_k = prefill.max_seqlen_q, prefill.max_seqlen_k
+
         return self._flash_attn_varlen_diff_headdims(
             q=q,
             k=k,
             v=v,
-            cu_seqlens_q=prefill.query_start_loc,
-            cu_seqlens_k=prefill.query_start_loc,
-            max_seqlen_q=prefill.max_query_len,
-            max_seqlen_k=prefill.max_query_len,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
             softmax_scale=self.scale,
             causal=True,
             return_softmax_lse=return_softmax_lse,
         )
 
     def _run_prefill_new_tokens_fi(
-        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
+        self, q, k, v, prefill: PrefillKernelMetadata, return_softmax_lse
     ):
         assert isinstance(prefill, FlashInferPrefillMetadata)
         assert prefill.prefill_main is not None
+        assert self.pcp_world_size == 1, "PCP is not supported for FlashInfer Prefill."
 
         ret = prefill.prefill_main.run(
             q=q,
@@ -2110,10 +2317,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         return ret
 
     def _run_prefill_new_tokens_cudnn(
-        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
+        self, q, k, v, prefill: PrefillKernelMetadata, return_softmax_lse
     ):
         assert isinstance(prefill, CudnnPrefillMetadata)
         assert prefill.query_seq_lens is not None
+        assert self.pcp_world_size == 1, "PCP is not supported for CUDNN Prefill."
         from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
 
         output, lse = cudnn_batch_prefill_with_kv_cache(
@@ -2196,13 +2404,15 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         )
 
     def _run_prefill_new_tokens_trtllm_ragged(
-        self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse
+        self, q, k, v, prefill: PrefillKernelMetadata, return_softmax_lse
     ):
         """TRT-LLM ragged attention for new tokens (causal)."""
         from flashinfer.prefill import trtllm_ragged_attention_deepseek
 
+        assert isinstance(prefill, MLACommonPrefillMetadata)
         assert prefill.query_seq_lens is not None
         assert prefill.workspace_buffer is not None
+        assert self.pcp_world_size == 1, "PCP is not supported for TRT-LLM Prefill."
         # allocate BF16 / FP16 output tensor for TRT-LLM ragged attention
         out = torch.empty(
             q.shape[0],
@@ -2415,7 +2625,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         k_scale: torch.Tensor,
         dcp_world_size: int,
     ):
-        assert k_scale is None, "DCP not support scaled kvcache now."
+        assert k_scale is None, "PCP/DCP not support scaled kvcache now."
         assert attn_metadata.prefill is not None
         prefill_metadata = attn_metadata.prefill
         assert prefill_metadata.chunked_context is not None
@@ -2442,7 +2652,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
                 seq_starts=prefill_metadata.chunked_context.starts[i],
             )
             # workspace
-            # |------- N tokens --------|--------- N*dcp_size tokens ----------|
+            # |------- N tokens --------|-------- N*dcp_size tokens ----------|
             # |<- use for loca_gather ->|<--------- use for allgather -------->|
             allgather_offset = workspace.shape[0] // (dcp_world_size + 1)
             assert allgather_offset * (dcp_world_size + 1) == workspace.shape[0]
@@ -2453,8 +2663,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             ]
             assert toks * dcp_world_size <= cur_allgather_workspace.shape[0]
             cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size]
+            # TODO(yyj) Reduce to a single all-gather operation
             cur_allgather_kvcache.copy_(
-                get_dcp_group().all_gather(local_gathered_kvcache, dim=0)
+                get_pcp_group().all_gather(
+                    get_dcp_group().all_gather(local_gathered_kvcache, dim=0),
+                    dim=0,
+                )
             )
             assert (
                 cur_allgather_kvcache.shape[-1]
@@ -2524,6 +2738,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         # TODO (zyongye): Prefill function here
         assert attn_metadata.prefill is not None
         assert self.dcp_world_size != -1
+        assert self.pcp_world_size != -1
 
         prefill_metadata = attn_metadata.prefill
         use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
@@ -2544,16 +2759,26 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             k = k.to(prefill_metadata.q_data_type)
             v = v.to(prefill_metadata.q_data_type)
 
-        output_prefill = self._run_prefill_new_tokens(
-            prefill=prefill_metadata,
-            q=q,
-            k=k,
-            v=v,
-            return_softmax_lse=has_context,
-        )
+        if self.pcp_world_size > 1:
+            output_prefill = self._run_prefill_new_tokens_pcp(
+                prefill=attn_metadata.prefill,
+                q=q,
+                k=k,
+                v=v,
+                return_softmax_lse=has_context,
+            )
+        else:
+            output_prefill = self._run_prefill_new_tokens(
+                prefill=prefill_metadata,
+                q=q,
+                k=k,
+                v=v,
+                return_softmax_lse=has_context,
+            )
 
         if has_context:
             suffix_output, suffix_lse = output_prefill
+            # DCP groups span PCP, so dcp_world_size is the total CP world size
             if self.dcp_world_size > 1:
                 context_output, context_lse = (
                     self._context_parallel_compute_prefill_context(

vllm/model_executor/layers/quantization/mxfp4.py

@@ -1258,7 +1258,7 @@ class XpuMxfp4MoEMethod(Mxfp4MoEMethod):
             topk_weights=routing_weights,
             topk_ids=selected_experts,
             n_experts_per_token=layer.top_k,
-            activation=layer.activation.value,
+            activation=layer.activation,
             num_experts=layer.local_num_experts,
             is_mxfp4=True,
         )

vllm/platforms/cuda.py

@@ -282,6 +282,20 @@ class CudaPlatformBase(Platform):
                         "backend."
                     )
 
+        # lazy import to avoid circular import
+        from vllm.config import CUDAGraphMode
+
+        compilation_config = vllm_config.compilation_config
+        if (
+            compilation_config.cudagraph_mode.has_full_cudagraphs()
+            and parallel_config.prefill_context_parallel_size > 1
+        ):
+            logger.warning_once(
+                "Prefill context parallel (PCP) is enabled, which is "
+                "incompatible with full CUDA graphs. "
+                "Overriding cudagraph_mode to PIECEWISE."
+            )
+            compilation_config.cudagraph_mode = CUDAGraphMode.PIECEWISE
         scheduler_config = vllm_config.scheduler_config
         # Note: model_config may be None during testing
         if (

vllm/v1/attention/backend.py

@@ -333,6 +333,14 @@ class CommonAttentionMetadata:
     dcp_local_seq_lens_cpu: torch.Tensor | None = None
     """Sequence lengths of the local rank in decode context parallelism world"""
 
+    pcp_allgather_restore_idx: torch.Tensor | None = None
+    """Indices to restore the original order of KV in prefill context parallelism"""
+
+    global_num_scheduled_tokens: torch.Tensor | None = None
+    """(batch_size,), GLOBAL (pre-PCP-partition) scheduled token counts per request.
+    Used by PCP to correctly compute num_computed_tokens, since with PCP
+    query_start_loc is local but seq_lens is global."""
+
     # WARNING: Deprecated fields. Will be removed in a future release (v0.15.0)
     _seq_lens_cpu: torch.Tensor | None = None
     _num_computed_tokens_cpu: torch.Tensor | None = None
@@ -380,10 +388,22 @@ class CommonAttentionMetadata:
         return self._num_computed_tokens_cpu
 
     def compute_num_computed_tokens(self) -> torch.Tensor:
-        """Compute num_computed_tokens on device (seq_lens - query_lens)."""
+        """Compute num_computed_tokens on device.
+
+        With PCP, query_start_loc is local (partitioned) but seq_lens is
+        global, so ``seq_lens - query_lens`` gives the wrong result. When
+        global_num_scheduled_tokens is available we use
+        ``seq_lens - global_num_scheduled_tokens`` instead, which is always
+        correct since seq_lens = num_computed + num_scheduled (both global).
+        """
         if self._num_computed_tokens_cache is None:
-            query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
-            self._num_computed_tokens_cache = self.seq_lens - query_lens
+            if self.global_num_scheduled_tokens is not None:
+                self._num_computed_tokens_cache = (
+                    self.seq_lens - self.global_num_scheduled_tokens
+                )
+            else:
+                query_lens = self.query_start_loc[1:] - self.query_start_loc[:-1]
+                self._num_computed_tokens_cache = self.seq_lens - query_lens
         return self._num_computed_tokens_cache
 
     # TODO(lucas): remove once we have FULL-CG spec-decode support
@@ -401,6 +421,9 @@ class CommonAttentionMetadata:
             _num_computed_tokens_cpu=self._num_computed_tokens_cpu[:num_actual_reqs]
             if self._num_computed_tokens_cpu is not None
             else None,
+            global_num_scheduled_tokens=maybe_slice_reqs(
+                self.global_num_scheduled_tokens
+            ),
             num_reqs=num_actual_reqs,
             num_actual_tokens=num_actual_tokens,
             max_query_len=self.max_query_len,
@@ -626,7 +649,7 @@ class AttentionImplBase(ABC, Generic[T]):
     # Whether the attention impl supports Prefill Context Parallelism.
     supports_pcp: bool = False
     # Whether the attention impl(or ops) supports MTP
-    # when cp_kv_cache_interleave_size > 1
+    # when dcp_kv_cache_interleave_size > 1
     supports_mtp_with_cp_non_trivial_interleave_size: bool = False
 
     # some attention backends might not always want to return lse
@@ -649,9 +672,6 @@ class AttentionImplBase(ABC, Generic[T]):
     pcp_world_size: int
     pcp_rank: int
 
-    total_cp_world_size: int
-    total_cp_rank: int
-
     def __new__(cls, *args, **kwargs):
         # use __new__ so that all subclasses will call this
         self = super().__new__(cls)
@@ -672,8 +692,6 @@ 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.need_to_return_lse_for_decode = (
             self.dcp_world_size > 1 and self.can_return_lse_for_decode

vllm/v1/attention/backends/flash_attn.py

@@ -22,7 +22,11 @@ from vllm.v1.attention.backends.fa_utils import (
     get_flash_attn_version,
     is_flash_attn_varlen_func_available,
 )
-from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
+from vllm.v1.attention.ops.common import (
+    cp_lse_ag_out_rs,
+    dcp_prepare_query,
+)
+from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
 
 if is_flash_attn_varlen_func_available():
@@ -34,7 +38,6 @@ if is_flash_attn_varlen_func_available():
     )
 from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vllm_config
 from vllm.config.cache import CacheDType
-from vllm.distributed.parallel_state import get_dcp_group
 from vllm.logger import init_logger
 from vllm.model_executor.layers.batch_invariant import (
     vllm_is_batch_invariant,
@@ -296,17 +299,19 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
         self.aot_schedule = get_flash_attn_version() == 3
 
         try:
-            from vllm.distributed.parallel_state import get_dcp_group
+            from vllm.distributed.parallel_state import get_dcp_group, get_tp_group
 
             self.dcp_world_size = get_dcp_group().world_size
             self.dcp_rank = get_dcp_group().rank_in_group
+            self.tp_world_size = get_tp_group().world_size
         except AssertionError:
-            # DCP might not be initialized in testing
+            # DCP/TP might not be initialized in testing
             self.dcp_world_size = 1
             self.dcp_rank = 0
+            self.tp_world_size = 1
 
-        self.cp_kv_cache_interleave_size = (
-            self.parallel_config.cp_kv_cache_interleave_size
+        self.dcp_kv_cache_interleave_size = (
+            self.parallel_config.dcp_kv_cache_interleave_size
         )
 
         self.use_full_cuda_graph = (
@@ -409,7 +414,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                     batch_size=batch_size,
                     max_seqlen_q=max_query_len,
                     max_seqlen_k=max_seq_len,
-                    num_heads_q=self.num_heads_q * self.dcp_world_size,
+                    num_heads_q=self.num_heads_q * self.tp_world_size,
                     num_heads_kv=self.num_heads_kv,
                     headdim=self.headdim,
                     cache_seqlens=seqlens,
@@ -439,16 +444,16 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
                 dcp_context_kv_lens,
                 self.dcp_world_size,
                 self.dcp_rank,
-                self.cp_kv_cache_interleave_size,
+                self.dcp_kv_cache_interleave_size,
             )
             # After DCP distribution, the maximum number of tokens for any rank is
             # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
-            # and I is cp_kv_cache_interleave_size.
+            # and I is dcp_kv_cache_interleave_size.
             # This eliminates GPU->CPU sync while minimizing workspace over-allocation.
-            num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
+            num_partitions = self.dcp_world_size * self.dcp_kv_cache_interleave_size
             max_dcp_context_kv_len = (
                 (max_seq_len + num_partitions - 1) // num_partitions
-            ) * self.cp_kv_cache_interleave_size
+            ) * self.dcp_kv_cache_interleave_size
 
             scheduler_metadata = schedule(
                 batch_size=num_reqs,
@@ -609,6 +614,13 @@ class FlashAttentionImpl(AttentionImpl):
 
         self.supports_quant_query_input = True
 
+        parallel_config = get_current_vllm_config().parallel_config
+        dcp_a2a = (
+            parallel_config.decode_context_parallel_size > 1
+            and parallel_config.dcp_comm_backend == "a2a"
+        )
+        self.dcp_combine = dcp_a2a_lse_reduce if dcp_a2a else cp_lse_ag_out_rs
+
     def forward(
         self,
         layer: torch.nn.Module,
@@ -830,12 +842,12 @@ class FlashAttentionImpl(AttentionImpl):
         block_table = attn_metadata.block_table
 
         query = query.contiguous()
-        query_across_dcp = get_dcp_group().all_gather(query, dim=1)
+        query_all_heads = dcp_prepare_query(query)
         sliding_window_size = (
             list(self.sliding_window) if self.sliding_window is not None else None
         )
         context_attn_out, context_lse = flash_attn_varlen_func(
-            q=query_across_dcp,
+            q=query_all_heads,
             k=key_cache,
             v=value_cache,
             out=None,
@@ -857,11 +869,10 @@ class FlashAttentionImpl(AttentionImpl):
             v_descale=v_descale,
             num_splits=attn_metadata.max_num_splits,
         )
-        # FA returns LSE in shape [ H, B ] but cp_lse_ag_out_rs wants [ B, H ]
-        context_attn_out_cor, context_lse_cor = cp_lse_ag_out_rs(
+        # FA returns LSE in shape [ H, B ] but DCP combine wants [ B, H ]
+        context_attn_out_cor, context_lse_cor = self.dcp_combine(
             context_attn_out,
             context_lse.transpose(0, 1),
-            get_dcp_group(),
             return_lse=True,
         )
         context_lse_cor = context_lse_cor.transpose(0, 1).contiguous()

vllm/v1/attention/backends/flashinfer.py

@@ -3,6 +3,7 @@
 """Attention layer with FlashInfer."""
 
 from dataclasses import dataclass
+from functools import partial
 from typing import ClassVar
 
 import numpy as np
@@ -58,7 +59,11 @@ from vllm.v1.attention.backends.utils import (
     infer_global_hyperparameters,
     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_rs,
+    dcp_prepare_query,
+)
+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.kv_cache_interface import AttentionSpec, UniformTypeKVCacheSpecs
 from vllm.v1.utils import CpuGpuBuffer
@@ -170,7 +175,12 @@ class BatchDCPPrefillWrapper:
     def __init__(
         self,
         workspace_buffer: torch.Tensor | None = None,
+        dcp_a2a: bool = False,
     ):
+        if dcp_a2a:
+            self._dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
+        else:
+            self._dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
         self._context = BatchPrefillWithPagedKVCacheWrapper(
             workspace_buffer, get_kv_cache_layout()
         )
@@ -239,22 +249,18 @@ class BatchDCPPrefillWrapper:
         value: torch.Tensor,
         out: torch.Tensor,
     ):
-        prefill_query_across_dcp = get_dcp_group().all_gather(
-            prefill_query.contiguous(), dim=1
-        )
+        prefill_query_all_heads = dcp_prepare_query(prefill_query.contiguous())
         output_context_tmp, lse_context_tmp = self._context.run(
-            prefill_query_across_dcp,
+            prefill_query_all_heads,
             kv_cache_permute,
             k_scale=layer._k_scale_float,
             v_scale=layer._v_scale_float,
             return_lse=True,
         )
-        output_context, lse_context = cp_lse_ag_out_rs(
+        output_context, lse_context = self._dcp_combine(
             output_context_tmp,
             lse_context_tmp,
-            get_dcp_group(),
             return_lse=True,
-            is_lse_base_on_e=False,
         )
         lse_context = lse_context.transpose(0, 1).contiguous()
 
@@ -552,6 +558,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             self.dcp_rank = 0
             self.dcp_kv_cache_interleave_size = 1
         self.use_dcp = self.dcp_world_size > 1
+        self.dcp_a2a = (
+            self.use_dcp and vllm_config.parallel_config.dcp_comm_backend == "a2a"
+        )
 
         self.num_qo_heads = self.model_config.get_num_attention_heads(
             self.vllm_config.parallel_config
@@ -701,6 +710,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
             if self.use_dcp:
                 self._prefill_wrapper = BatchDCPPrefillWrapper(
                     workspace_buffer=self._get_workspace_buffer(),
+                    dcp_a2a=self.dcp_a2a,
                 )
             else:
                 self._prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper(
@@ -1219,6 +1229,19 @@ class FlashInferImpl(AttentionImpl):
         self.bmm2_scale: float | None = None
         self.o_sf_scale: float | None = None
 
+        try:
+            parallel_config = vllm_config.parallel_config
+            dcp_a2a = (
+                parallel_config.decode_context_parallel_size > 1
+                and parallel_config.dcp_comm_backend == "a2a"
+            )
+        except AttributeError:
+            dcp_a2a = False
+        if dcp_a2a:
+            self.dcp_combine = partial(dcp_a2a_lse_reduce, is_lse_base_on_e=False)
+        else:
+            self.dcp_combine = partial(cp_lse_ag_out_rs, is_lse_base_on_e=False)
+
     def fused_output_quant_supported(self, quant_key: QuantKey):
         return (
             self.support_trtllm_attn
@@ -1487,9 +1510,7 @@ class FlashInferImpl(AttentionImpl):
                 assert decode_wrapper._sm_scale == self.scale
 
                 if use_dcp:
-                    decode_query = get_dcp_group().all_gather(
-                        decode_query.contiguous(), dim=-2
-                    )
+                    decode_query = dcp_prepare_query(decode_query.contiguous())
                     output_tmp = torch.empty_like(decode_query)
                     lse = torch.empty(
                         (decode_query.size(0), decode_query.size(1)),
@@ -1505,11 +1526,10 @@ class FlashInferImpl(AttentionImpl):
                         lse=lse,
                         return_lse=True,
                     )
-                    output[:num_decode_tokens] = cp_lse_ag_out_rs(
+                    output[:num_decode_tokens] = self.dcp_combine(
                         output_tmp,
                         lse,
                         get_dcp_group(),
-                        is_lse_base_on_e=False,
                     )
                 else:
                     decode_wrapper.run(

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

@@ -112,7 +112,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
         vllm_config: VllmConfig,
         device: torch.device,
     ):
-        interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
+        interleave_size = vllm_config.parallel_config.dcp_kv_cache_interleave_size
         super().__init__(
             kv_cache_spec,
             layer_names,
@@ -251,6 +251,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
 
 class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
     can_return_lse_for_decode: bool = True
+    supports_pcp: bool = True
 
     def __init__(
         self,

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

@@ -186,6 +186,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
 
 class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
     can_return_lse_for_decode: bool = True
+    supports_pcp: bool = True
 
     def __init__(
         self,

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

@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, ClassVar
+from typing import TYPE_CHECKING, ClassVar, Optional
 
 import numpy as np
 import torch
@@ -549,7 +549,7 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
         kv_sharing_target_layer_name: str | None,
         # MLA Specific Arguments
         topk_indice_buffer: torch.Tensor | None = None,
-        indexer: "Indexer | None" = None,
+        indexer: Optional["Indexer"] = None,
         **mla_args,
     ) -> None:
         self.num_heads = num_heads

vllm/v1/attention/backends/utils.py

@@ -16,6 +16,7 @@ import torch
 from typing_extensions import runtime_checkable
 
 from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import cdiv
 from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
 
@@ -27,6 +28,7 @@ import vllm.envs as envs
 from vllm.distributed.kv_transfer.kv_connector.utils import (
     get_kv_connector_cache_layout,
 )
+from vllm.distributed.parallel_state import GroupCoordinator
 from vllm.logger import init_logger
 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
 from vllm.v1.attention.backend import (
@@ -508,18 +510,48 @@ def split_decodes_and_prefills(
         num_decode_tokens: The number of tokens in the decode requests.
         num_prefill_tokens: The number of tokens in the prefill requests.
     """
-    max_query_len = common_attn_metadata.max_query_len
     num_reqs = common_attn_metadata.num_reqs
     num_tokens = common_attn_metadata.num_actual_tokens
     query_start_loc = common_attn_metadata.query_start_loc_cpu
 
-    if max_query_len <= decode_threshold and (
-        not require_uniform or decode_threshold <= 1
-    ):
-        return num_reqs, 0, num_tokens, 0
+    # For PCP, use global_num_scheduled_tokens for classification since
+    # query_start_loc contains local (partitioned) counts
+    global_scheduled = common_attn_metadata.global_num_scheduled_tokens
+    if global_scheduled is not None:
+        global_query_lens = global_scheduled[:num_reqs].cpu()
+        max_query_len = int(global_query_lens.max().item())
+        # Also get num_computed_tokens to distinguish decode from new prefill
+        num_computed = common_attn_metadata.compute_num_computed_tokens()[
+            :num_reqs
+        ].cpu()
+    else:
+        max_query_len = common_attn_metadata.max_query_len
+        global_query_lens = None
+        num_computed = None
 
-    query_lens = query_start_loc[1:] - query_start_loc[:-1]
-    if query_lens[0].item() > decode_threshold:
+    # For the "all decode" fast path, also need to check num_computed_tokens
+    # (new prefills have num_computed_tokens=0 and should not be classified as decode)
+    all_decode = max_query_len <= decode_threshold and (
+        not require_uniform or decode_threshold <= 1
+    )
+    if all_decode:
+        # With PCP, check if any request is a new prefill (num_computed_tokens=0)
+        if num_computed is not None and torch.any(num_computed == 0):
+            all_decode = False
+        if all_decode:
+            return num_reqs, 0, num_tokens, 0
+
+    # Use global lens if available (PCP), otherwise compute from local query_start_loc
+    if global_query_lens is not None:
+        query_lens = global_query_lens
+    else:
+        query_lens = query_start_loc[1:] - query_start_loc[:-1]
+    # Check if first request is prefill (cannot be decode if query_lens > threshold
+    # OR if it's a new prefill with num_computed_tokens == 0)
+    first_is_prefill = query_lens[0].item() > decode_threshold
+    if num_computed is not None and num_computed[0].item() == 0:
+        first_is_prefill = True
+    if first_is_prefill:
         # first request is not decode, so no decode requests
         return 0, num_reqs, 0, num_tokens
 
@@ -527,17 +559,25 @@ 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)
+        # Note: skip the total tokens check for PCP since num_tokens is local
         if torch.all((query_lens == query_lens[0]) | (query_lens == 0)):
-            assert num_reqs * query_lens[0] == num_tokens, "tokens not padded correctly"
+            if global_query_lens is None:
+                assert num_reqs * query_lens[0] == num_tokens, (
+                    "tokens not padded correctly"
+                )
             return num_reqs, 0, num_tokens, 0  # all decodes
         is_prefill = query_lens != query_lens[0]
     else:
+        # Prefill = query_lens > threshold OR num_computed_tokens == 0 (new prefill)
         is_prefill = query_lens > decode_threshold
+        if num_computed is not None:
+            is_prefill = is_prefill | (num_computed == 0)
 
     if not torch.any(is_prefill):
         return num_reqs, 0, num_tokens, 0
 
     first_prefill = is_prefill.int().argmax(dim=-1).item()
+    # Classification is based on query_lens (global for PCP), but assertion should match
     assert torch.all(query_lens[:first_prefill] <= decode_threshold)
     num_decodes = first_prefill
     num_prefills = num_reqs - num_decodes
@@ -789,7 +829,7 @@ def get_dcp_local_seq_lens(
     seq_lens: torch.Tensor,
     dcp_size: int = 1,
     dcp_rank: int | None = None,
-    cp_kv_cache_interleave_size: int = 1,
+    dcp_kv_cache_interleave_size: int = 1,
 ) -> torch.Tensor:
     """While using dcp, kv_cache size stored on each rank may be different,
     use this function to calculate split decode seq_lens of each dcp rank.
@@ -811,20 +851,115 @@ def get_dcp_local_seq_lens(
     )
     base = (
         seq_lens_tiled
-        // cp_kv_cache_interleave_size
+        // dcp_kv_cache_interleave_size
         // dcp_size
-        * cp_kv_cache_interleave_size
+        * dcp_kv_cache_interleave_size
     )
     remainder = seq_lens_tiled - base * dcp_size
     remainder = torch.clip(
-        remainder - rank_offsets * cp_kv_cache_interleave_size,
+        remainder - rank_offsets * dcp_kv_cache_interleave_size,
         0,
-        cp_kv_cache_interleave_size,
+        dcp_kv_cache_interleave_size,
     )
     dcp_local_seq_lens = base + remainder
     return dcp_local_seq_lens.squeeze(1)
 
 
+def pcp_kv_allgather_and_restore(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    num_actual_tokens: int,
+    pcp_allgather_restore_idx: torch.Tensor,
+    pcp_group: GroupCoordinator,
+):
+    """
+    All-gather key and value tensors across PCP ranks and restore the original order.
+    Args:
+        key: key tensor for the current pcp rank.
+        value: value tensor for the current pcp rank.
+        num_actual_tokens: number of actual tokens (Exclude graph padding tokens).
+        pcp_allgather_restore_idx: indices to restore the original order.
+        pcp_group: PCP group coordinator.
+    Returns:
+        key: all-gathered and restored key tensor.
+        value: all-gathered and restored value tensor.
+    """
+    # NOTE(yyj): we must `slice` key and value because pcp_allgather_restore_idx
+    # ignores the padding from CUDA Graph.
+    # TODO(yyj) Batch all-gather operations to reduce launch overhead.
+    # Be careful about the dimensions of key and value.
+    key_across_cp = pcp_group.all_gather(key[:num_actual_tokens].contiguous(), dim=0)
+    value_across_cp = pcp_group.all_gather(
+        value[:num_actual_tokens].contiguous(), dim=0
+    )
+
+    # Reorder kv after pcp allgather.
+    # Note that there are duplicate decoding tokens after allgather.
+    key = torch.index_select(key_across_cp, 0, pcp_allgather_restore_idx)
+    value = torch.index_select(value_across_cp, 0, pcp_allgather_restore_idx)
+
+    return key, value
+
+
+def get_pcp_query_restore_idx(cu_num_tokens: torch.Tensor) -> torch.Tensor:
+    """
+    Get restore index for PCP query splitting.
+
+    When queries are split into head/tail halves for PCP, this returns
+    the argsort index to restore original order after processing.
+
+    Args:
+        cu_num_tokens: cumulative token counts, shape [num_reqs + 1]
+
+    Returns:
+        restore_idx: tensor to reorder concatenated [head, tail] back to original
+    """
+    cu = cu_num_tokens.cpu().numpy()
+    starts, ends = cu[:-1], cu[1:]
+    half_lens = (ends - starts) // 2
+    total = half_lens.sum()
+
+    seq_ids = np.repeat(np.arange(len(half_lens)), half_lens)
+    cu_half = np.concatenate([[0], np.cumsum(half_lens)[:-1]])
+    offsets = np.arange(total) - cu_half[seq_ids]
+
+    head = starts[seq_ids] + offsets
+    tail = (ends - half_lens)[seq_ids] + offsets
+    return torch.from_numpy(np.concatenate([head, tail]).argsort().astype(np.int32))
+
+
+def extend_all_queries_by_1(
+    common_attn_metadata: CommonAttentionMetadata,
+    arange: torch.Tensor,
+    new_slot_mapping: torch.Tensor,
+) -> CommonAttentionMetadata:
+    """
+    Creates a new CommonAttentionMetadata with all query lengths increased by 1.
+    Also all seq lens are increased by 1.
+    This is useful e.g. in speculative decoding with draft models, where we
+    extend each sequence by 1 token.
+    The slot mapping is computed externally, as it requires more information.
+    """
+    cad = common_attn_metadata
+    # query start loc must be increased by [+0, +1, +2, ..., +batch_size]
+    new_query_start_loc = cad.query_start_loc + arange[: len(cad.query_start_loc)]
+    new_query_start_loc_cpu = cad.query_start_loc_cpu + torch.arange(
+        len(cad.query_start_loc_cpu), dtype=torch.int32
+    )
+    new_cad = cad.replace(
+        query_start_loc=new_query_start_loc,
+        query_start_loc_cpu=new_query_start_loc_cpu,
+        seq_lens=cad.seq_lens + 1,
+        # each request is extended by 1 token -> batch_size tokens are added
+        num_actual_tokens=cad.num_actual_tokens + cad.batch_size(),
+        # All query lens increase by 1, so max query len increases by 1
+        max_query_len=cad.max_query_len + 1,
+        max_seq_len=cad.max_seq_len + 1,
+        slot_mapping=new_slot_mapping,
+    )
+    return new_cad
+
+
 def mamba_get_block_table_tensor(
     block_table: torch.Tensor,
     seq_lens: torch.Tensor,
@@ -864,3 +999,220 @@ def mamba_get_block_table_tensor(
         )
         indices_to_gather = (start_indices.unsqueeze(1) + offsets).to(torch.int64)
         return torch.gather(block_table, 1, indices_to_gather)
+
+
+@triton.jit
+def _fused_pcp_qkv_select_kernel(
+    q_ptr,
+    q_stride_B,
+    q_stride_H,
+    k_ptr,
+    k_stride_B,
+    k_stride_H,
+    v_ptr,
+    v_stride_B,
+    v_stride_H,
+    query_start_ptr,
+    out_q_head_ptr,
+    out_q_tail_ptr,
+    out_k_head_ptr,
+    out_k_tail_ptr,
+    out_v_head_ptr,
+    out_v_tail_ptr,
+    pcp_world_size: tl.constexpr,
+    pcp_rank: tl.constexpr,
+    n_head: tl.constexpr,
+    q_head_dim: tl.constexpr,
+    k_head_dim: tl.constexpr,
+    v_head_dim: tl.constexpr,
+    SEQ_BLOCK_SIZE: tl.constexpr,
+    DIM_BLOCK_SIZE: tl.constexpr,
+):
+    req_id = tl.program_id(0) // (2 * pcp_world_size)
+    seq_block_id = tl.program_id(0) % (2 * pcp_world_size)
+    head_id = tl.program_id(1)
+    dim_block_id = tl.program_id(2)
+    dim_off = tl.arange(0, DIM_BLOCK_SIZE) + dim_block_id * DIM_BLOCK_SIZE
+
+    q_start_loc = tl.load(query_start_ptr + req_id)
+    q_end_loc = tl.load(query_start_ptr + req_id + 1)
+    q_select_len = (q_end_loc - q_start_loc) // 2
+
+    # Select Q
+    if seq_block_id < 2:
+        block_q_start_loc = q_start_loc + seq_block_id * q_select_len
+        out_ptr = out_q_head_ptr if seq_block_id == 0 else out_q_tail_ptr
+        for qi in range(tl.cdiv(q_select_len, SEQ_BLOCK_SIZE)):
+            q_offset = tl.arange(0, SEQ_BLOCK_SIZE) + qi * SEQ_BLOCK_SIZE
+            mask = (dim_off[None, :] < q_head_dim) & (q_offset[:, None] < q_select_len)
+            q_src_idx = block_q_start_loc + q_offset[:, None]
+            q_dst_idx = q_start_loc // 2 + q_offset[:, None]
+            q_val = tl.load(
+                q_ptr
+                + q_src_idx * q_stride_B
+                + head_id * q_stride_H
+                + dim_off[None, :],
+                mask=mask,
+            )
+            tl.store(
+                out_ptr
+                + q_dst_idx * n_head * q_head_dim
+                + head_id * q_head_dim
+                + dim_off[None, :],
+                q_val,
+                mask=mask,
+            )
+
+    # Select KV
+    kv_start_loc = q_start_loc * pcp_world_size
+    kv_select_len = q_select_len
+    k_d_mask = dim_off[None, :] < k_head_dim
+    v_d_mask = dim_off[None, :] < v_head_dim
+    block_src_kv_start_loc = kv_start_loc + seq_block_id * kv_select_len
+    block_dst_kv_head_start_loc = (
+        kv_start_loc // 2 // pcp_world_size * (pcp_rank + 1)
+        + seq_block_id * kv_select_len
+    )
+    block_dst_kv_tail_start_loc = (
+        kv_start_loc // 2 // pcp_world_size * (2 * pcp_world_size - pcp_rank)
+        + seq_block_id * kv_select_len
+    )
+    for ki in range(tl.cdiv(kv_select_len, SEQ_BLOCK_SIZE)):
+        kv_offset = tl.arange(0, SEQ_BLOCK_SIZE) + ki * SEQ_BLOCK_SIZE
+        kv_block_mask = kv_offset[:, None] < kv_select_len
+        kv_src_idx = block_src_kv_start_loc + kv_offset[:, None]
+        kv_dst_idx_head = block_dst_kv_head_start_loc + kv_offset[:, None]
+        kv_dst_idx_tail = block_dst_kv_tail_start_loc + kv_offset[:, None]
+        k_val = tl.load(
+            k_ptr + kv_src_idx * k_stride_B + head_id * k_stride_H + dim_off[None, :],
+            mask=k_d_mask & kv_block_mask,
+        )
+        v_val = tl.load(
+            v_ptr + kv_src_idx * v_stride_B + head_id * v_stride_H + dim_off[None, :],
+            mask=v_d_mask & kv_block_mask,
+        )
+        if seq_block_id < pcp_rank + 1:
+            tl.store(
+                out_k_head_ptr
+                + kv_dst_idx_head * n_head * k_head_dim
+                + head_id * k_head_dim
+                + dim_off[None, :],
+                k_val,
+                mask=k_d_mask & kv_block_mask,
+            )
+            tl.store(
+                out_v_head_ptr
+                + kv_dst_idx_head * n_head * v_head_dim
+                + head_id * v_head_dim
+                + dim_off[None, :],
+                v_val,
+                mask=v_d_mask & kv_block_mask,
+            )
+        if seq_block_id < 2 * pcp_world_size - pcp_rank:
+            tl.store(
+                out_k_tail_ptr
+                + kv_dst_idx_tail * n_head * k_head_dim
+                + head_id * k_head_dim
+                + dim_off[None, :],
+                k_val,
+                mask=k_d_mask & kv_block_mask,
+            )
+            tl.store(
+                out_v_tail_ptr
+                + kv_dst_idx_tail * n_head * v_head_dim
+                + head_id * v_head_dim
+                + dim_off[None, :],
+                v_val,
+                mask=v_d_mask & kv_block_mask,
+            )
+
+
+def fused_pcp_qkv_select(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    query_start_loc: torch.Tensor,
+    pcp_world_size: int,
+    pcp_rank: int,
+):
+    """
+    Select the query and kv tensors for PCP. Instead of calling
+    `torch.index_select` multiple times, this function fuses the
+    selection for Q, K, and V into a single kernel to reduce
+    kernel launch overhead.
+    Args:
+        q: query tensor on the current PCP rank.
+        k: key tensor across PCP ranks.
+        v: value tensor across PCP ranks.
+        query_start_loc: start location of each query.
+        pcp_world_size: number of PCP ranks.
+        pcp_rank: rank of the current PCP rank.
+    Returns:
+        q_head: selected query tensor for pcp head.
+        k_head: selected key tensor for pcp head.
+        v_head: selected value tensor for pcp head.
+        q_tail: selected query tensor for pcp tail.
+        k_tail: selected key tensor for pcp tail.
+        v_tail: selected value tensor for pcp tail.
+
+    """
+    q_head = torch.empty(
+        (q.size(0) // 2,) + q.shape[1:], device=q.device, dtype=q.dtype
+    )
+    q_tail = torch.empty_like(q_head)
+    k_head = torch.empty(
+        (q.size(0) // 2 * (pcp_rank + 1),) + k.shape[1:], device=k.device, dtype=k.dtype
+    )
+    v_head = torch.empty(
+        (q.size(0) // 2 * (pcp_rank + 1),) + v.shape[1:], device=v.device, dtype=v.dtype
+    )
+    k_tail = torch.empty(
+        (q.size(0) // 2 * (2 * pcp_world_size - pcp_rank),) + k.shape[1:],
+        device=k.device,
+        dtype=k.dtype,
+    )
+    v_tail = torch.empty(
+        (q.size(0) // 2 * (2 * pcp_world_size - pcp_rank),) + v.shape[1:],
+        device=v.device,
+        dtype=v.dtype,
+    )
+    BS = len(query_start_loc) - 1
+    DIM_BLOCK_SIZE: int = 64
+    SEQ_BLOCK_SIZE: int = 256
+    assert q.shape[1] == k.shape[1] == v.shape[1]
+    n_head = q.shape[1]
+    n_dim_block = (
+        max(q.shape[2], k.shape[2], v.shape[2]) + DIM_BLOCK_SIZE
+    ) // DIM_BLOCK_SIZE
+    grid = (
+        2 * pcp_world_size * BS,
+        n_head,
+        n_dim_block,
+    )
+    _fused_pcp_qkv_select_kernel[grid](
+        q,
+        q.stride(0),
+        q.stride(1),
+        k,
+        k.stride(0),
+        k.stride(1),
+        v,
+        v.stride(0),
+        v.stride(1),
+        query_start_loc,
+        q_head,
+        q_tail,
+        k_head,
+        k_tail,
+        v_head,
+        v_tail,
+        pcp_world_size,
+        pcp_rank,
+        n_head,
+        q.shape[2],
+        k.shape[2],
+        v.shape[2],
+        SEQ_BLOCK_SIZE,
+        DIM_BLOCK_SIZE,
+    )
+    return q_head, k_head, v_head, q_tail, k_tail, v_tail

vllm/v1/attention/ops/common.py

@@ -2,7 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import torch
 
-from vllm.distributed.parallel_state import GroupCoordinator
+from vllm.distributed.parallel_state import (
+    GroupCoordinator,
+    get_dcp_group,
+    get_pcp_group,
+    get_tp_group,
+)
 from vllm.triton_utils import tl, triton
 
 
@@ -190,7 +195,7 @@ def _cp_lse_common(
     cp_attn_lse: [ B, H ]
     """
     if cp_group.world_size == 1:
-        return cp_attn_out
+        return cp_attn_out, cp_attn_lse
 
     if ctx is None:
         ctx = CPTritonContext()
@@ -256,6 +261,85 @@ def cp_lse_ag_out_ar(
     return out
 
 
+# Backward compatibility aliases for DCP
+DCPTritonContext = CPTritonContext
+
+
+def dcp_prepare_query(query: torch.Tensor) -> torch.Tensor:
+    """Prepare query for DCP decode attention.
+
+    Two cases based on DCP configuration:
+    - Case 1 (DCP = PCP): No all-gather needed, ranks already have same heads
+    - Case 2 (DCP = TP × PCP): All-gather across TP to get all heads
+    """
+    dcp_group = get_dcp_group()
+    tp_group = get_tp_group()
+
+    try:
+        pcp_world_size = get_pcp_group().world_size
+    except AssertionError:
+        pcp_world_size = 1
+
+    # Case 1: DCP = PCP (same TP position, no all-gather needed)
+    # Case 2: DCP = TP × PCP (spans TP, need all-gather)
+    dcp_spans_tp = dcp_group.world_size > pcp_world_size
+
+    if dcp_spans_tp:
+        return tp_group.all_gather(query, dim=1)
+    else:
+        return query
+
+
+def dcp_reduce_output(
+    attn_output: torch.Tensor,
+    attn_lse: torch.Tensor,
+    ctx: CPTritonContext | None = None,
+    return_lse: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Reduce DCP partial attention outputs across the DCP group.
+
+    Two cases based on DCP configuration:
+    - Case 1 (DCP = PCP): All-reduce only (no scatter, same heads)
+    - Case 2 (DCP = TP × PCP): Reduce-scatter (all-reduce + scatter to TP heads)
+    """
+    dcp_group = get_dcp_group()
+    tp_group = get_tp_group()
+
+    try:
+        pcp_world_size = get_pcp_group().world_size
+    except AssertionError:
+        pcp_world_size = 1
+
+    # Case 1: DCP = PCP (same TP position)
+    # Case 2: DCP = TP × PCP (spans TP)
+    dcp_spans_tp = dcp_group.world_size > pcp_world_size
+
+    if not dcp_spans_tp:
+        # Case 1: All-reduce only (no scatter needed, ranks have same heads)
+        return cp_lse_ag_out_ar(
+            attn_output,
+            attn_lse,
+            dcp_group,
+            ctx=ctx,
+            return_lse=return_lse,
+            is_lse_base_on_e=True,
+        )
+    else:
+        # Case 2: All-reduce across DCP, then slice to TP-local heads
+        out, lse = _cp_lse_common(
+            attn_output, attn_lse, dcp_group, ctx=ctx, is_lse_base_on_e=True
+        )
+        out = dcp_group.all_reduce(out)
+        # Slice to TP-local heads (already reduced, just need to select)
+        tp_rank = tp_group.rank_in_group
+        tp_num_heads = out.shape[1] // tp_group.world_size
+        out = out[:, tp_num_heads * tp_rank : tp_num_heads * (tp_rank + 1), :]
+        if return_lse:
+            lse = lse[:, tp_num_heads * tp_rank : tp_num_heads * (tp_rank + 1)]
+            return out, lse
+        return out
+
+
 @triton.jit
 def _pack_seq_kernel(
     x_ptr,  # [N, D]

vllm/v1/attention/ops/dcp_alltoall.py

@@ -0,0 +1,363 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+DCP All-to-All communication backend for attention.
+
+Provides All-to-All (A2A) communication as an alternative to
+AllGather + ReduceScatter (AG+RS) for Decode Context Parallel (DCP).
+Instead of gathering the full Q tensor and scattering partial outputs,
+A2A exchanges partial attention outputs and their LSE values across
+ranks, then combines them with exact LSE-weighted reduction.
+
+This reduces the number of NCCL calls per attention layer from 3
+(AG for Q, AG for K metadata, RS for output) to 2 (A2A for output,
+A2A for LSE), lowering per-step communication overhead for long-context
+decode where NCCL latency is a significant fraction of step time.
+
+Usage:
+    vllm serve model --tp 16 --dcp 16 --dcp-comm-backend a2a
+
+Reference: https://arxiv.org/abs/2507.07120
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+
+from vllm.triton_utils import tl, triton
+
+if TYPE_CHECKING:
+    from vllm.distributed.parallel_state import GroupCoordinator
+    from vllm.v1.attention.ops.common import CPTritonContext
+
+
+def _lse_weighted_combine(
+    outputs: torch.Tensor,
+    lses: torch.Tensor,
+    return_lse: bool = False,
+    is_lse_base_on_e: bool = True,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """
+    CPU reference implementation for LSE-weighted combination.
+
+    This is a pure PyTorch implementation used for testing and validation.
+    For GPU execution, use dcp_lse_combine_triton instead.
+
+    Args:
+        outputs: Partial attention outputs [N, B, H, D]
+                 N = number of KV shards (ranks)
+                 B = batch size (num_tokens)
+                 H = number of heads per rank
+                 D = head dimension
+        lses: Log-sum-exp values [N, B, H]
+        return_lse: If True, also return the global LSE
+        is_lse_base_on_e: If True, LSE is base e; if False, base 2
+
+    Returns:
+        Combined output [B, H, D], and optionally global LSE [B, H]
+    """
+    N, B, H, D = outputs.shape
+
+    # Handle NaN and inf in LSEs
+    lses = torch.where(
+        torch.isnan(lses) | torch.isinf(lses),
+        torch.tensor(float("-inf"), device=lses.device, dtype=lses.dtype),
+        lses,
+    )
+
+    # Compute max LSE for numerical stability
+    lse_max, _ = lses.max(dim=0)  # [B, H]
+    lse_max = torch.where(
+        lse_max == float("-inf"),
+        torch.zeros_like(lse_max),
+        lse_max,
+    )
+
+    # Compute weights: softmax over the N dimension
+    if is_lse_base_on_e:
+        weights = torch.exp(lses - lse_max.unsqueeze(0))  # [N, B, H]
+    else:
+        weights = torch.pow(2.0, lses - lse_max.unsqueeze(0))  # [N, B, H]
+
+    # Handle NaN weights
+    weights = torch.where(torch.isnan(weights), torch.zeros_like(weights), weights)
+
+    # Normalize weights
+    weight_sum = weights.sum(dim=0, keepdim=True)  # [1, B, H]
+    weights = weights / weight_sum.clamp(min=1e-10)  # [N, B, H]
+
+    # Weighted combination: sum over N dimension
+    result = (outputs * weights.unsqueeze(-1)).sum(dim=0)  # [B, H, D]
+
+    if return_lse:
+        if is_lse_base_on_e:
+            global_lse = torch.log(weight_sum.squeeze(0)) + lse_max  # [B, H]
+        else:
+            global_lse = torch.log2(weight_sum.squeeze(0)) + lse_max  # [B, H]
+        return result, global_lse
+
+    return result
+
+
+@triton.jit
+def _dcp_lse_combine_kernel(
+    # Input pointers
+    recv_output_ptr,
+    recv_lse_ptr,
+    # Output pointers
+    out_ptr,
+    out_lse_ptr,
+    # Strides for recv_output [N, B, H_local, D]
+    ro_stride_N,
+    ro_stride_B,
+    ro_stride_H,
+    ro_stride_D,
+    # Strides for recv_lse [N, B, H_local]
+    rl_stride_N,
+    rl_stride_B,
+    rl_stride_H,
+    # Strides for output [B, H_local, D]
+    o_stride_B,
+    o_stride_H,
+    o_stride_D,
+    # Constants
+    N: tl.constexpr,
+    HEAD_DIM: tl.constexpr,
+    IS_BASE_E: tl.constexpr,
+    RETURN_LSE: tl.constexpr,
+):
+    """
+    Triton kernel for LSE-weighted combination of partial attention outputs.
+
+    After All-to-All, each rank has:
+    - recv_output [N, B, H_local, D]: partial outputs from all KV shards
+    - recv_lse [N, B, H_local]: partial LSEs from all KV shards
+
+    This kernel computes the weighted combination locally (no communication).
+
+    Grid: (B, H_local)
+    Each program handles one (batch, head) and processes all D elements.
+    """
+    batch_idx = tl.program_id(0).to(tl.int64)
+    head_idx = tl.program_id(1).to(tl.int64)
+
+    # Base offset for this (batch, head)
+    base_lse_offset = batch_idx * rl_stride_B + head_idx * rl_stride_H
+    base_out_offset = batch_idx * ro_stride_B + head_idx * ro_stride_H
+
+    # First pass: find max LSE for numerical stability
+    lse_max = -float("inf")
+    for n in tl.static_range(N):
+        lse_offset = n * rl_stride_N + base_lse_offset
+        lse_val = tl.load(recv_lse_ptr + lse_offset)
+        lse_val = tl.where(
+            (lse_val != lse_val) | (lse_val == float("inf")),
+            -float("inf"),
+            lse_val,
+        )
+        lse_max = tl.maximum(lse_max, lse_val)
+
+    lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max)
+
+    # Second pass: compute sum of exp(lse - max)
+    lse_sum = 0.0
+    for n in tl.static_range(N):
+        lse_offset = n * rl_stride_N + base_lse_offset
+        lse_val = tl.load(recv_lse_ptr + lse_offset)
+        lse_val = tl.where(
+            (lse_val != lse_val) | (lse_val == float("inf")),
+            -float("inf"),
+            lse_val,
+        )
+        if IS_BASE_E:
+            lse_sum += tl.exp(lse_val - lse_max)
+        else:
+            lse_sum += tl.exp2(lse_val - lse_max)
+
+    # Compute global LSE
+    if IS_BASE_E:  # noqa: SIM108
+        global_lse = tl.log(lse_sum) + lse_max
+    else:
+        global_lse = tl.log2(lse_sum) + lse_max
+
+    # Third pass: weighted combination across D dimension
+    d_offsets = tl.arange(0, HEAD_DIM)
+    acc = tl.zeros([HEAD_DIM], dtype=tl.float32)
+
+    for n in tl.static_range(N):
+        lse_offset = n * rl_stride_N + base_lse_offset
+        lse_val = tl.load(recv_lse_ptr + lse_offset)
+        lse_val = tl.where(
+            (lse_val != lse_val) | (lse_val == float("inf")),
+            -float("inf"),
+            lse_val,
+        )
+        if IS_BASE_E:
+            weight = tl.exp(lse_val - global_lse)
+        else:
+            weight = tl.exp2(lse_val - global_lse)
+        weight = tl.where(weight != weight, 0.0, weight)
+
+        out_offsets = n * ro_stride_N + base_out_offset + d_offsets * ro_stride_D
+        out_vals = tl.load(recv_output_ptr + out_offsets)
+        acc += out_vals.to(tl.float32) * weight
+
+    # Store result
+    final_offsets = (
+        batch_idx * o_stride_B + head_idx * o_stride_H + d_offsets * o_stride_D
+    )
+    tl.store(out_ptr + final_offsets, acc)
+
+    if RETURN_LSE:
+        tl.store(out_lse_ptr + base_lse_offset, global_lse)
+
+
+def dcp_lse_combine_triton(
+    recv_output: torch.Tensor,
+    recv_lse: torch.Tensor,
+    return_lse: bool = False,
+    is_lse_base_on_e: bool = True,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """
+    Triton-accelerated LSE-weighted combination for DCP A2A.
+
+    Args:
+        recv_output: [N, B, H_local, D] - partial outputs from all KV shards
+        recv_lse: [N, B, H_local] - partial LSEs from all KV shards
+        return_lse: If True, also return the global LSE
+        is_lse_base_on_e: If True, LSE is base e; if False, base 2
+
+    Returns:
+        Combined output [B, H_local, D]
+        If return_lse=True, also returns global_lse [B, H_local]
+    """
+    N, B, H_local, D = recv_output.shape
+
+    out = torch.empty(
+        (B, H_local, D), device=recv_output.device, dtype=recv_output.dtype
+    )
+
+    if return_lse:
+        out_lse = torch.empty(
+            (B, H_local), device=recv_lse.device, dtype=recv_lse.dtype
+        )
+    else:
+        out_lse = torch.empty(1, device=recv_lse.device, dtype=recv_lse.dtype)
+
+    ro_stride_N, ro_stride_B, ro_stride_H, ro_stride_D = recv_output.stride()
+    rl_stride_N, rl_stride_B, rl_stride_H = recv_lse.stride()
+    o_stride_B, o_stride_H, o_stride_D = out.stride()
+
+    grid = (B, H_local, 1)
+
+    _dcp_lse_combine_kernel[grid](
+        recv_output,
+        recv_lse,
+        out,
+        out_lse,
+        ro_stride_N,
+        ro_stride_B,
+        ro_stride_H,
+        ro_stride_D,
+        rl_stride_N,
+        rl_stride_B,
+        rl_stride_H,
+        o_stride_B,
+        o_stride_H,
+        o_stride_D,
+        N=N,
+        HEAD_DIM=D,
+        IS_BASE_E=is_lse_base_on_e,
+        RETURN_LSE=return_lse,
+    )
+
+    if return_lse:
+        return out, out_lse
+    return out
+
+
+def dcp_a2a_lse_reduce(
+    cp_attn_out: torch.Tensor,
+    cp_attn_lse: torch.Tensor,
+    cp_group: GroupCoordinator,
+    ctx: CPTritonContext | None = None,
+    return_lse: bool = False,
+    is_lse_base_on_e: bool = True,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """
+    Combine partial attention outputs across DCP ranks using All-to-All.
+
+    Each rank holds attention output for all heads but only a local shard
+    of the KV cache. This function:
+    1. Exchanges partial outputs across ranks via All-to-All
+    2. Exchanges LSE values via All-to-All
+    3. Combines them with exact LSE-weighted reduction (Triton kernel)
+
+    Tensor flow:
+        Input:  cp_attn_out [B, H, D] - all heads, local KV shard
+        Reshape: [N, B, H/N, D] - split heads across ranks
+        A2A:    Two all_to_all_single calls (output and LSE)
+        Combine: recv [N, B, H/N, D] + lse [N, B, H/N] -> [B, H/N, D]
+
+    Args:
+        cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim
+        cp_attn_lse: [B, H] log-sum-exp values (fp32)
+        cp_group: GroupCoordinator for DCP communication
+        ctx: CPTritonContext (unused, for signature compatibility)
+        return_lse: If True, also return the combined global LSE
+        is_lse_base_on_e: If True, LSE is base e; if False, base 2
+
+    Returns:
+        Combined output [B, H/N, D] (head-scattered)
+        If return_lse=True, also returns global_lse [B, H/N]
+    """
+    world_size = cp_group.world_size
+
+    if world_size == 1:
+        if return_lse:
+            return cp_attn_out, cp_attn_lse
+        return cp_attn_out
+
+    local_output = cp_attn_out.contiguous()
+    local_lse = cp_attn_lse.contiguous()
+
+    B, H, D = local_output.shape
+    H_per_rank = H // world_size
+
+    # Reshape for All-to-All: [B, H, D] -> [N, B, H/N, D]
+    # Split heads into N chunks, each destined for a different rank
+    send_output = (
+        local_output.view(B, world_size, H_per_rank, D).permute(1, 0, 2, 3).contiguous()
+    )
+    recv_output = torch.empty_like(send_output)
+
+    # Same for LSE: [B, H] -> [N, B, H/N]
+    send_lse = local_lse.view(B, world_size, H_per_rank).permute(1, 0, 2).contiguous()
+    recv_lse = torch.empty_like(send_lse)
+
+    # All-to-All for partial attention outputs and LSE values (async overlap)
+    work_output = dist.all_to_all_single(
+        recv_output.view(-1),
+        send_output.view(-1),
+        group=cp_group.device_group,
+        async_op=True,
+    )
+    work_lse = dist.all_to_all_single(
+        recv_lse.view(-1),
+        send_lse.view(-1),
+        group=cp_group.device_group,
+        async_op=True,
+    )
+    work_output.wait()
+    work_lse.wait()
+
+    # LSE-weighted combination via Triton kernel (local, no communication)
+    return dcp_lse_combine_triton(
+        recv_output,
+        recv_lse,
+        return_lse=return_lse,
+        is_lse_base_on_e=is_lse_base_on_e,
+    )

vllm/v1/core/kv_cache_coordinator.py

@@ -335,8 +335,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

@@ -1300,14 +1300,10 @@ def _report_kv_cache_config(
         * min_block_size
     )
     dcp_size = vllm_config.parallel_config.decode_context_parallel_size
-    pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
-    if pcp_size * dcp_size > 1:
-        num_tokens *= pcp_size * dcp_size
+    if dcp_size > 1:
+        num_tokens *= dcp_size
         logger.info(
-            "Multiplying the GPU KV cache size by the cp_world_size %d "
-            "(pcp_world_size %d * dcp_world_size %d).",
-            pcp_size * dcp_size,
-            pcp_size,
+            "Multiplying the GPU KV cache size by the dcp_world_size %d.",
             dcp_size,
         )
     num_tokens_str = f"{num_tokens:,}"

vllm/v1/core/single_type_kv_cache_manager.py

@@ -50,8 +50,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
@@ -429,8 +429,8 @@ class FullAttentionManager(SingleTypeKVCacheManager):
             [] for _ in range(len(kv_cache_group_ids))
         )
         block_size = kv_cache_spec.block_size
-        if dcp_world_size * pcp_world_size > 1:
-            block_size *= dcp_world_size * pcp_world_size
+        if dcp_world_size > 1:
+            block_size *= dcp_world_size
         max_num_blocks = max_length // block_size
         for block_hash in itertools.islice(block_hashes, max_num_blocks):
             # block_hashes is a chain of block hashes. If a block hash is not

vllm/v1/engine/core.py

@@ -137,11 +137,8 @@ class EngineCore:
                 logger.warning("Disabling chunked prefill for model without KVCache")
                 vllm_config.scheduler_config.enable_chunked_prefill = False
 
-        scheduler_block_size = (
-            vllm_config.cache_config.block_size
-            * vllm_config.parallel_config.decode_context_parallel_size
-            * vllm_config.parallel_config.prefill_context_parallel_size
-        )
+        dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+        scheduler_block_size = vllm_config.cache_config.block_size * dcp_size
 
         self.scheduler: SchedulerInterface = Scheduler(
             vllm_config=vllm_config,

vllm/v1/kv_cache_interface.py

@@ -113,11 +113,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/spec_decode/eagle.py

@@ -978,6 +978,8 @@ class SpecDecodeBaseProposer:
             slot_mapping=common_attn_metadata.slot_mapping[:total_num_tokens],
             causal=True,
             dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
+            dcp_local_seq_lens_cpu=common_attn_metadata.dcp_local_seq_lens_cpu,
+            pcp_allgather_restore_idx=common_attn_metadata.pcp_allgather_restore_idx,
         )
 
         return (
@@ -1257,6 +1259,8 @@ class SpecDecodeBaseProposer:
             slot_mapping=common_attn_metadata.slot_mapping[token_indices],
             causal=True,
             dcp_local_seq_lens=common_attn_metadata.dcp_local_seq_lens,
+            dcp_local_seq_lens_cpu=common_attn_metadata.dcp_local_seq_lens_cpu,
+            pcp_allgather_restore_idx=common_attn_metadata.pcp_allgather_restore_idx,
         )
 
         return spec_common_attn_metadata, token_indices

vllm/v1/worker/block_table.py

@@ -4,11 +4,10 @@
 import numpy as np
 import torch
 
-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.utils.math_utils import cdiv
 from vllm.v1.utils import CpuGpuBuffer
-from vllm.v1.worker.cp_utils import get_total_cp_world_size
 
 logger = init_logger(__name__)
 
@@ -23,7 +22,7 @@ class BlockTable:
         pin_memory: bool,
         device: torch.device,
         kernel_block_size: int,
-        cp_kv_cache_interleave_size: int,
+        dcp_kv_cache_interleave_size: int,
     ):
         """
         Args:
@@ -81,13 +80,6 @@ class BlockTable:
         else:
             self._kernel_block_arange = None
 
-        try:
-            self.pcp_world_size = get_pcp_group().world_size
-            self.pcp_rank = get_pcp_group().rank_in_group
-        except AssertionError:
-            # PCP might not be initialized in testing
-            self.pcp_world_size = 1
-            self.pcp_rank = 0
         try:
             self.dcp_world_size = get_dcp_group().world_size
             self.dcp_rank = get_dcp_group().rank_in_group
@@ -95,7 +87,7 @@ class BlockTable:
             # DCP might not be initialized in testing
             self.dcp_world_size = 1
             self.dcp_rank = 0
-        self.cp_kv_cache_interleave_size = cp_kv_cache_interleave_size
+        self.dcp_kv_cache_interleave_size = dcp_kv_cache_interleave_size
 
     def append_row(
         self,
@@ -139,16 +131,19 @@ class BlockTable:
         # NOTE(woosuk): We can't simply use `token_indices // block_size`
         # here because M (max_model_len) is not necessarily divisible by
         # block_size.
-        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
-        if total_cp_world_size > 1:
+        # Only DCP shards the KV cache. With PCP, tokens are split during
+        # prefill but K/V are gathered (pcp_kv_allgather_and_restore) so
+        # each rank inserts the FULL sequence into its cache.
+        dcp_world_size = self.dcp_world_size
+        dcp_rank = self.dcp_rank
+        if dcp_world_size > 1:
             # Note(hc): The DCP implement store kvcache with an interleave
             # style, the kvcache for the token whose token_idx is i is
-            # always stored on the GPU whose dcp_rank equals i % cp_world_size:
+            # always stored on the GPU whose dcp_rank equals i % dcp_world_size:
 
             # Use a "virtual block" which equals to world_size * block_size
             # for block_table_indices calculation.
-            virtual_block_size = self.block_size * total_cp_world_size
+            virtual_block_size = self.block_size * dcp_world_size
             block_table_indices = (
                 req_indices * self.max_num_blocks_per_req
                 + positions // virtual_block_size
@@ -160,16 +155,16 @@ class BlockTable:
             virtual_block_offsets = positions % virtual_block_size
             mask = (
                 virtual_block_offsets
-                // self.cp_kv_cache_interleave_size
-                % total_cp_world_size
-                == total_cp_rank
+                // self.dcp_kv_cache_interleave_size
+                % dcp_world_size
+                == dcp_rank
             )
             # Calculate local block_offsets
             block_offsets = (
                 virtual_block_offsets
-                // (total_cp_world_size * self.cp_kv_cache_interleave_size)
-                * self.cp_kv_cache_interleave_size
-                + virtual_block_offsets % self.cp_kv_cache_interleave_size
+                // (dcp_world_size * self.dcp_kv_cache_interleave_size)
+                * self.dcp_kv_cache_interleave_size
+                + virtual_block_offsets % self.dcp_kv_cache_interleave_size
             )
             # Calculate slot_mapping
             slot_mapping = block_numbers * self.block_size + block_offsets
@@ -263,7 +258,7 @@ class MultiGroupBlockTable:
         block_sizes: list[int],
         kernel_block_sizes: list[int],
         max_num_blocks: list[int] | None = None,
-        cp_kv_cache_interleave_size: int = 1,
+        dcp_kv_cache_interleave_size: int = 1,
     ) -> None:
         if len(kernel_block_sizes) != len(block_sizes):
             raise ValueError(
@@ -275,9 +270,9 @@ class MultiGroupBlockTable:
             # (max_model_len//dcp_world_size) tokens in kvcache,
             # so the block_size which used for calc max_num_blocks_per_req
             # must be multiplied by dcp_world_size.
-            total_cp_world_size = get_total_cp_world_size()
+            dcp_world_size = get_dcp_group().world_size
             max_num_blocks = [
-                cdiv(max_model_len, block_size * total_cp_world_size)
+                cdiv(max_model_len, block_size * dcp_world_size)
                 for block_size in block_sizes
             ]
 
@@ -296,7 +291,7 @@ class MultiGroupBlockTable:
                 pin_memory,
                 device,
                 kernel_block_size,
-                cp_kv_cache_interleave_size,
+                dcp_kv_cache_interleave_size,
             )
             for block_size, kernel_block_size, max_num_blocks_per_req in zip(
                 block_sizes, kernel_block_sizes, max_num_blocks

vllm/v1/worker/cp_utils.py

@@ -2,8 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from typing import TYPE_CHECKING, Any, cast
 
+import numpy as np
+import torch
+
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.distributed import get_dcp_group, get_pcp_group
+from vllm.v1.utils import CpuGpuBuffer
 
 if TYPE_CHECKING:
     from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -11,10 +15,263 @@ else:
     AttentionLayerBase = object
 
 
+class PCPManager:
+    """Manager for Prefill Context Parallelism (PCP) buffers and partitioning."""
+
+    def __init__(
+        self,
+        pcp_world_size: int,
+        pcp_rank: int,
+        max_num_reqs: int,
+        max_padded_num_tokens: int,
+        device: torch.device,
+        pin_memory: bool = False,
+    ) -> None:
+        self.pcp_world_size = pcp_world_size
+        self.pcp_rank = pcp_rank
+        self.device = device
+
+        self.pcp_allgather_restore_idx = CpuGpuBuffer(
+            max_padded_num_tokens,
+            dtype=torch.int64,
+            device=device,
+            pin_memory=pin_memory,
+        )
+        self._pcp_unpad_mask_tensor = torch.zeros(
+            max_padded_num_tokens, device="cpu", dtype=torch.bool
+        )
+        self.pcp_unpad_mask = self._pcp_unpad_mask_tensor.numpy()
+        self.pcp_padded_slot_mapping = torch.empty(
+            max_padded_num_tokens, dtype=torch.int64, device=device
+        )
+        self.global_num_scheduled_tokens = CpuGpuBuffer(
+            max_num_reqs, dtype=torch.int32, device=device, pin_memory=pin_memory
+        )
+        # Cached values from partition_inputs
+        self.local_num_scheduled: np.ndarray = np.array([], dtype=np.int32)
+        self.local_total: int = 0
+        self.padded_total: int = 0
+        self.global_total: int = 0  # Total tokens before partitioning
+
+    def partition_inputs(
+        self,
+        positions_np: np.ndarray,
+        req_indices: np.ndarray,
+        num_scheduled_tokens: np.ndarray,
+        num_computed_tokens: np.ndarray,
+        arange_np: np.ndarray,
+        reorder_batch_threshold: int | None,
+    ) -> tuple[int, np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Partition inputs for this PCP rank using DualChunkSwap splitting.
+
+        Each request's prefill tokens are split across PCP ranks using a
+        "DualChunkSwap" pattern: tokens are padded to a multiple of 2*world_size,
+        then split into head/tail chunks assigned to ranks in an interleaved
+        pattern to balance computation across the sequence.
+
+        For decode requests (tokens <= reorder_batch_threshold AND has context),
+        tokens are duplicated across all ranks instead of split. This matches
+        the reorder_batch_to_split_decodes_and_prefills definition.
+
+        This method:
+        1. Computes which tokens this rank processes
+        2. Gathers local positions/req_indices from the global arrays
+        3. Builds pcp_allgather_restore_idx for reordering after all-gather
+        4. Builds pcp_unpad_mask for removing padding after restore
+
+        Args:
+            positions_np: Global position values, modified in-place to local
+            req_indices: Global request indices, modified in-place to local
+            num_scheduled_tokens: Per-request token counts
+            num_computed_tokens: Per-request computed token counts (0 = prefill)
+            arange_np: Pre-allocated arange buffer
+            reorder_batch_threshold: Threshold distinguishing decode vs prefill
+
+        Returns:
+            (local_total, positions_np[:local_total], req_indices[:local_total],
+             gathered_positions[:local_total])
+
+        Note:
+            positions_np contains position values for RoPE encoding (DualChunkSwap).
+            gathered_positions contains positions gathered from the global array,
+            with padding positions clamped to the request start. These are used
+            for computing token_indices for input token gathering.
+        """
+        assert reorder_batch_threshold is not None
+        num_reqs = len(num_scheduled_tokens)
+        # Decode = tokens <= threshold AND has computed context (not a new prefill)
+        # This matches reorder_batch_to_split_decodes_and_prefills definition
+        is_decode = (num_scheduled_tokens <= reorder_batch_threshold) & (
+            num_computed_tokens > 0
+        )
+        num_decode_reqs = int(is_decode.sum())
+        num_decode_tokens = int(num_scheduled_tokens[:num_decode_reqs].sum())
+        ws = self.pcp_world_size
+
+        # Pad to multiple of 2*ws; decode reqs are duplicated instead
+        padded = np.ceil(num_scheduled_tokens / (2 * ws)).astype(np.int32) * (2 * ws)
+        padded[:num_decode_reqs] = num_scheduled_tokens[:num_decode_reqs] * ws
+
+        # Cumsum and arange for padded tokens
+        cu_padded = np.cumsum(padded)
+        padded_total = cu_padded[-1]
+        padded_arange = arange_np[:padded_total] - np.repeat(cu_padded - padded, padded)
+
+        # Unpad mask: True for real tokens
+        self.pcp_unpad_mask[:padded_total] = padded_arange < np.repeat(
+            num_scheduled_tokens, padded
+        )
+
+        # Tokens per request for this rank
+        pcp_tokens = padded // ws
+        chunk = (pcp_tokens // 2).clip(min=1)
+        # Decodes are replicated across PCP ranks, so chunk == pcp_tokens
+        chunk[:num_decode_reqs] = pcp_tokens[:num_decode_reqs]
+
+        # Aranges for pcp_tokens and chunks
+        cu_pcp = np.cumsum(pcp_tokens)
+        pcp_arange = arange_np[: cu_pcp[-1]] - np.repeat(
+            cu_pcp - pcp_tokens, pcp_tokens
+        )
+        cu_chunk = np.cumsum(chunk)
+        chunk_arange = arange_np[: cu_chunk[-1]] - np.repeat(cu_chunk - chunk, chunk)
+
+        # Head/tail mask
+        head_mask = pcp_arange < np.repeat(chunk, pcp_tokens)
+
+        def get_positions(start_loc: int | np.ndarray, rank: int) -> np.ndarray:
+            """Get positions for given rank with DualChunkSwap interleaving."""
+            pos = np.zeros(len(head_mask), dtype=np.int32)
+            head_start = start_loc + rank * chunk
+            tail_start = start_loc + (2 * ws - rank - 1) * chunk
+            pos[head_mask] = chunk_arange + np.repeat(head_start, chunk)
+            pos[~head_mask] = (
+                chunk_arange[num_decode_tokens:]
+                + np.repeat(tail_start, chunk)[num_decode_tokens:]
+            )
+            return pos
+
+        # Positions for this rank
+        positions = get_positions(0, self.pcp_rank)
+        if num_decode_reqs > 0:
+            cu_dec = np.cumsum(num_scheduled_tokens[:num_decode_reqs])
+            positions[:num_decode_tokens] = arange_np[:num_decode_tokens] - np.repeat(
+                cu_dec - num_scheduled_tokens[:num_decode_reqs],
+                num_scheduled_tokens[:num_decode_reqs],
+            )
+
+        # Build restore index for all-gather
+        padded_start = np.concatenate([[0], cu_padded[:-1]])
+        all_pos = np.concatenate([get_positions(padded_start, r) for r in range(ws)])
+        self.pcp_allgather_restore_idx.np[: len(all_pos)] = all_pos.argsort()
+        self.pcp_allgather_restore_idx.copy_to_gpu(len(all_pos))
+
+        # Convert position values to indices into original batch
+        # (for input token gathering)
+        cu_orig = np.cumsum(num_scheduled_tokens)
+        orig_start = np.concatenate([[0], cu_orig[:-1]])
+        pcp_total = int(pcp_tokens.sum())
+        orig_lens = np.repeat(num_scheduled_tokens, pcp_tokens)
+        orig_starts = np.repeat(orig_start, pcp_tokens)
+        # For padding positions (position >= seq_len), clamp to request's first token
+        # to avoid indexing into other requests' positions
+        local_indices = np.where(
+            positions[:pcp_total] >= orig_lens,
+            orig_starts,  # Clamp to request's start, not global 0
+            orig_starts + positions[:pcp_total],
+        ).astype(np.int64)
+
+        # Gather positions for token indices (clamped for padding, same as old behavior)
+        gathered_positions = positions_np[local_indices].copy()
+
+        # Compute actual position values for RoPE encoding.
+        # For padding positions, clamp to 0 (relative) so RoPE matches the
+        # clamped token content. This ensures padding K/V have consistent
+        # RoPE with the token they were cloned from.
+        is_padding = positions[:pcp_total] >= orig_lens
+        clamped_positions = np.where(is_padding, 0, positions[:pcp_total])
+        req_computed_tokens = np.repeat(num_computed_tokens, pcp_tokens)
+        positions_np[:pcp_total] = req_computed_tokens[:pcp_total] + clamped_positions
+
+        # Gather req_indices using local_indices (needed for input token gathering)
+        req_indices[:pcp_total] = req_indices[local_indices]
+
+        self.local_num_scheduled = pcp_tokens[:num_reqs]
+        self.local_total = pcp_total
+        self.padded_total = int(padded_total)
+        self.global_total = int(num_scheduled_tokens.sum())
+
+        gns = self.global_num_scheduled_tokens
+        gns.np[:num_reqs] = num_scheduled_tokens
+        gns.np[num_reqs:].fill(0)
+        gns.copy_to_gpu()
+
+        return (
+            pcp_total,
+            positions_np[:pcp_total],
+            req_indices[:pcp_total],
+            gathered_positions,
+        )
+
+    def restore_hidden_states(
+        self, hidden_states: torch.Tensor, num_tokens: int
+    ) -> torch.Tensor:
+        """All-gather hidden states, restore order, and remove padding."""
+        hidden_states = get_pcp_group().all_gather(hidden_states[:num_tokens], 0)
+        # Use padded_total for correct restore index size
+        restore_size = (
+            self.padded_total if self.padded_total > 0 else hidden_states.shape[0]
+        )
+        restore_idx = self.pcp_allgather_restore_idx.gpu[:restore_size]
+        hidden_states = hidden_states.index_select(0, restore_idx)
+        mask = self._pcp_unpad_mask_tensor[:restore_size].to(
+            hidden_states.device, non_blocking=True
+        )
+        return hidden_states[mask]
+
+    def pad_slot_mapping(self, slot_mapping: torch.Tensor) -> torch.Tensor:
+        """
+        Expand slot_mapping for the all-gathered KV cache.
+
+        After KV all-gather, slot_mapping needs to account for per-request
+        PCP padding. This places real slot values at unpadded positions
+        and -1 at padding positions.
+
+        Args:
+            slot_mapping: Unpadded slot mapping (global real tokens only)
+        """
+        if self.padded_total == 0:
+            return slot_mapping
+
+        out = self.pcp_padded_slot_mapping[: self.padded_total]
+        out.fill_(-1)
+        mask = self._pcp_unpad_mask_tensor[: self.padded_total].to(
+            slot_mapping.device, non_blocking=True
+        )
+        out[mask] = slot_mapping
+        return out
+
+
+def get_total_cp_world_size() -> int:
+    """Get the total context parallel world size (PCP * DCP)."""
+    try:
+        pcp_world_size = get_pcp_group().world_size
+    except AssertionError:
+        # PCP might not be initialized in testing
+        pcp_world_size = 1
+    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
+
+
 def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
     pcp_size = vllm_config.parallel_config.prefill_context_parallel_size
     dcp_size = vllm_config.parallel_config.decode_context_parallel_size
-    interleave_size = vllm_config.parallel_config.cp_kv_cache_interleave_size
+    interleave_size = vllm_config.parallel_config.dcp_kv_cache_interleave_size
     if pcp_size * dcp_size > 1:
         layer_type = cast(type[Any], AttentionLayerBase)
         layers = get_layers_from_vllm_config(vllm_config, layer_type)
@@ -24,7 +281,7 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
                 continue
             if vllm_config.speculative_config is not None and interleave_size > 1:
                 assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, (
-                    "MTP with cp_kv_cache_interleave_size > 1 is not "
+                    "MTP with dcp_kv_cache_interleave_size > 1 is not "
                     f"supported in {layer_impl.__class__.__name__}."
                 )
             if dcp_size > 1:
@@ -41,17 +298,3 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
                     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
-    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

vllm/v1/worker/gpu_input_batch.py

@@ -94,7 +94,7 @@ class InputBatch:
         logitsprocs_need_output_token_ids: bool = False,
         is_spec_decode: bool = False,
         is_pooling_model: bool = False,
-        cp_kv_cache_interleave_size: int = 1,
+        dcp_kv_cache_interleave_size: int = 1,
     ):
         self.is_pooling_model = is_pooling_model
         self.is_spec_decode = is_spec_decode
@@ -147,7 +147,7 @@ class InputBatch:
             block_sizes=block_sizes,
             kernel_block_sizes=kernel_block_sizes,
             max_num_blocks=max_num_blocks_per_req,
-            cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
+            dcp_kv_cache_interleave_size=dcp_kv_cache_interleave_size,
         )
 
         # Sampling-related.

vllm/v1/worker/gpu_model_runner.py

@@ -37,6 +37,7 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_
 from vllm.distributed.kv_transfer.kv_connector.utils import copy_kv_blocks
 from vllm.distributed.parallel_state import (
     get_dcp_group,
+    get_pcp_group,
     get_pp_group,
     get_tp_group,
     graph_capture,
@@ -166,6 +167,7 @@ from vllm.v1.structured_output.utils import apply_grammar_bitmask
 from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
 from vllm.v1.worker import mamba_utils
 from vllm.v1.worker.cp_utils import (
+    PCPManager,
     check_attention_cp_compatibility,
     get_total_cp_world_size,
 )
@@ -416,7 +418,9 @@ class GPUModelRunner(
         # Always set to false after the first forward pass
         self.calculate_kv_scales = self.cache_config.calculate_kv_scales
         self.dcp_world_size = self.parallel_config.decode_context_parallel_size
+        self.pcp_world_size = self.parallel_config.prefill_context_parallel_size
         self.dcp_rank = 0 if self.dcp_world_size <= 1 else get_dcp_group().rank_in_group
+        self.pcp_rank = 0 if self.pcp_world_size <= 1 else get_pcp_group().rank_in_group
         self.max_num_tokens = scheduler_config.max_num_batched_tokens
         self.max_num_reqs = scheduler_config.max_num_seqs
 
@@ -577,7 +581,7 @@ class GPUModelRunner(
             # uses output token ids so we set this conservatively.
             logitsprocs_need_output_token_ids=bool(custom_logitsprocs),
             is_pooling_model=self.is_pooling_model,
-            cp_kv_cache_interleave_size=self.parallel_config.cp_kv_cache_interleave_size,
+            dcp_kv_cache_interleave_size=self.parallel_config.dcp_kv_cache_interleave_size,
         )
 
         # Separate cuda stream for overlapping transfer of sampled token ids from
@@ -606,9 +610,19 @@ class GPUModelRunner(
         self.encoder_timing_registry: dict[str, EncoderTimingStats] = {}
         self._encoder_timing_lock = threading.Lock()
 
+        # For PCP, buffers need extra capacity for DualChunkSwap padding
+        if self.pcp_world_size > 1:
+            ws = self.pcp_world_size
+            max_padded_num_tokens = max(
+                self.max_num_tokens * ws,
+                self.max_num_tokens + self.max_num_reqs * 2 * ws,
+            )
+        else:
+            max_padded_num_tokens = self.max_num_tokens
+
         # Persistent buffers for CUDA graphs.
-        self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32)
-        self.positions = self._make_buffer(self.max_num_tokens, dtype=torch.int64)
+        self.input_ids = self._make_buffer(max_padded_num_tokens, dtype=torch.int32)
+        self.positions = self._make_buffer(max_padded_num_tokens, dtype=torch.int64)
         self.query_start_loc = self._make_buffer(
             self.max_num_reqs + 1, dtype=torch.int32
         )
@@ -622,9 +636,12 @@ class GPUModelRunner(
         # version of this tensor, avoid a RuntimeError by not creating a
         # numpy buffer.
         self.inputs_embeds = self._make_buffer(
-            self.max_num_tokens, self.inputs_embeds_size, dtype=self.dtype, numpy=False
+            max_padded_num_tokens,
+            self.inputs_embeds_size,
+            dtype=self.dtype,
+            numpy=False,
         )
-        self.is_token_ids = self._make_buffer(self.max_num_tokens, dtype=torch.bool)
+        self.is_token_ids = self._make_buffer(max_padded_num_tokens, dtype=torch.bool)
         self.discard_request_mask = self._make_buffer(
             self.max_num_reqs, dtype=torch.bool
         )
@@ -640,11 +657,22 @@ class GPUModelRunner(
             # Double buffer to avoid race condition: previous iteration's async
             # copy may still be reading from CPU while current iteration writes.
             self.is_mm_embed_buffers = [
-                self._make_buffer(self.max_num_tokens, dtype=torch.bool),
-                self._make_buffer(self.max_num_tokens, dtype=torch.bool),
+                self._make_buffer(max_padded_num_tokens, dtype=torch.bool),
+                self._make_buffer(max_padded_num_tokens, dtype=torch.bool),
             ]
             self.is_mm_embed_idx = 0
 
+        # Manager for Prefill Context Parallelism
+        if self.pcp_world_size > 1:
+            self.pcp_manager = PCPManager(
+                self.pcp_world_size,
+                self.pcp_rank,
+                self.max_num_reqs,
+                max_padded_num_tokens,
+                self.device,
+                self.pin_memory,
+            )
+
         # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
         if self.uses_mrope:
             # NOTE: `mrope_positions` is implemented with one additional dummy
@@ -658,7 +686,7 @@ class GPUModelRunner(
             # 1D-RoPE.
             # See page 5 of https://arxiv.org/abs/2409.12191
             self.mrope_positions = self._make_buffer(
-                (3, self.max_num_tokens + 1), dtype=torch.int64
+                (3, max_padded_num_tokens + 1), dtype=torch.int64
             )
 
         # Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
@@ -674,7 +702,7 @@ class GPUModelRunner(
         # OPTIMIZATION: Cache the tensors rather than creating them every step.
         # Keep in int64 to avoid overflow with long context
         self.arange_np = np.arange(
-            max(self.max_num_reqs + 1, self.max_model_len, self.max_num_tokens),
+            max(self.max_num_reqs + 1, self.max_model_len, max_padded_num_tokens),
             dtype=np.int64,
         )
 
@@ -688,7 +716,7 @@ class GPUModelRunner(
         self.kv_sharing_fast_prefill_logits_indices = None
         if self.cache_config.kv_sharing_fast_prefill:
             self.kv_sharing_fast_prefill_logits_indices = torch.zeros(
-                self.max_num_tokens, dtype=torch.int32, device=self.device
+                max_padded_num_tokens, dtype=torch.int32, device=self.device
             )
 
         self.uniform_decode_query_len = 1 + self.num_spec_tokens
@@ -1519,14 +1547,15 @@ class GPUModelRunner(
     ) -> tuple[
         torch.Tensor,
         SpecDecodeMetadata | None,
+        int,  # num_tokens for model forward
     ]:
         """
         :return: tuple[
-            logits_indices, spec_decode_metadata,
+            logits_indices, spec_decode_metadata, num_tokens
         ]
         """
-        total_num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
-        assert total_num_scheduled_tokens > 0
+        total_num_tokens = scheduler_output.total_num_scheduled_tokens
+        assert total_num_tokens > 0
         num_reqs = self.input_batch.num_reqs
         assert num_reqs > 0
 
@@ -1543,13 +1572,34 @@ class GPUModelRunner(
         cu_num_tokens, arange = self._get_cumsum_and_arange(num_scheduled_tokens)
 
         # Get positions.
-        positions_np = self.positions.np[:total_num_scheduled_tokens]
+        positions_np = self.positions.np[:total_num_tokens]
         np.add(
             self.input_batch.num_computed_tokens_cpu[req_indices],
             arange,
             out=positions_np,
         )
 
+        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
+        self.input_batch.block_table.commit_slot_mapping(total_num_tokens)
+
+        logits_indices = torch.from_numpy(cu_num_tokens - 1).to(
+            self.device, non_blocking=True
+        )
+
+        local_token_indices = None
+        if self.pcp_world_size > 1:
+            total_num_tokens, positions_np, req_indices, local_token_indices = (
+                self.pcp_manager.partition_inputs(
+                    positions_np,
+                    req_indices,
+                    num_scheduled_tokens,
+                    self.input_batch.num_computed_tokens_cpu[:num_reqs],
+                    self.arange_np,
+                    self.reorder_batch_threshold,
+                )
+            )
+            cu_num_tokens = np.cumsum(self.pcp_manager.local_num_scheduled)
+
         # Calculate M-RoPE positions.
         # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
         if self.uses_mrope:
@@ -1564,9 +1614,15 @@ class GPUModelRunner(
         # E.g., [0, 1, 0, 1, 2, 3, 4, 0, 1, 2]
         # -> [0, 1, M, M + 1, M + 2, M + 3, M + 4, 2 * M, 2 * M + 1, 2 * M + 2]
         # where M is the max_model_len.
-        token_indices = (
-            positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
-        )
+        if local_token_indices is not None:
+            token_indices = (
+                local_token_indices
+                + req_indices * self.input_batch.token_ids_cpu.shape[1]
+            )
+        else:
+            token_indices = (
+                positions_np + req_indices * self.input_batch.token_ids_cpu.shape[1]
+            )
         token_indices_tensor = torch.from_numpy(token_indices)
 
         # NOTE(woosuk): We use torch.index_select instead of np.take here
@@ -1576,7 +1632,7 @@ class GPUModelRunner(
             self.input_batch.token_ids_cpu_tensor.flatten(),
             0,
             token_indices_tensor,
-            out=self.input_ids.cpu[:total_num_scheduled_tokens],
+            out=self.input_ids.cpu[:total_num_tokens],
         )
         if self.enable_prompt_embeds:
             is_token_ids = self.input_batch.is_token_ids_tensor.flatten()
@@ -1584,7 +1640,7 @@ class GPUModelRunner(
                 is_token_ids,
                 0,
                 token_indices_tensor,
-                out=self.is_token_ids.cpu[:total_num_scheduled_tokens],
+                out=self.is_token_ids.cpu[:total_num_tokens],
             )
 
         # Because we did not pre-allocate a massive prompt_embeds CPU tensor on
@@ -1625,9 +1681,6 @@ class GPUModelRunner(
 
                 output_idx += num_sched
 
-        self.input_batch.block_table.compute_slot_mapping(req_indices, positions_np)
-        self.input_batch.block_table.commit_slot_mapping(total_num_scheduled_tokens)
-
         # Prepare the attention metadata.
         self.query_start_loc.np[0] = 0
         self.query_start_loc.np[1 : num_reqs + 1] = cu_num_tokens
@@ -1635,7 +1688,6 @@ class GPUModelRunner(
         # like FlashAttention requires that
         self.query_start_loc.np[num_reqs + 1 :].fill(cu_num_tokens[-1])
         self.query_start_loc.copy_to_gpu()
-        query_start_loc = self.query_start_loc.gpu[: num_reqs + 1]
 
         self.seq_lens.np[:num_reqs] = (
             self.input_batch.num_computed_tokens_cpu[:num_reqs] + num_scheduled_tokens
@@ -1657,25 +1709,25 @@ class GPUModelRunner(
         # Copy the tensors to the GPU.
         self._prepare_input_ids(
             scheduler_output,
-            total_num_scheduled_tokens,
+            total_num_tokens,
             cu_num_tokens,
         )
 
         if self.uses_mrope:
             # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
-            self.mrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(
-                self.mrope_positions.cpu[:, :total_num_scheduled_tokens],
+            self.mrope_positions.gpu[:, :total_num_tokens].copy_(
+                self.mrope_positions.cpu[:, :total_num_tokens],
                 non_blocking=True,
             )
         elif self.uses_xdrope_dim > 0:
             # Only relevant for models using XD-RoPE (e.g, HunYuan-VL)
-            self.xdrope_positions.gpu[:, :total_num_scheduled_tokens].copy_(
-                self.xdrope_positions.cpu[:, :total_num_scheduled_tokens],
+            self.xdrope_positions.gpu[:, :total_num_tokens].copy_(
+                self.xdrope_positions.cpu[:, :total_num_tokens],
                 non_blocking=True,
             )
         else:
             # Common case (1D positions)
-            self.positions.copy_to_gpu(total_num_scheduled_tokens)
+            self.positions.copy_to_gpu(total_num_tokens)
 
         use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
         if not use_spec_decode:
@@ -1684,10 +1736,10 @@ class GPUModelRunner(
             # from these partial requests, we do so for simplicity.
             # We will ignore the sampled tokens from the partial requests.
             # TODO: Support prompt logprobs.
-            logits_indices = query_start_loc[1:] - 1
             spec_decode_metadata = None
             num_sampled_tokens = np.ones(num_reqs, dtype=np.int32)
         else:
+            assert self.pcp_world_size == 1, "PCP not support spec decode now"
             # Get the number of draft tokens for each request.
             # Iterate over the dictionary rather than all requests since not all
             # requests have draft tokens.
@@ -1729,6 +1781,7 @@ class GPUModelRunner(
         return (
             logits_indices,
             spec_decode_metadata,
+            total_num_tokens,
         )
 
     def _build_attention_metadata(
@@ -1824,7 +1877,7 @@ class GPUModelRunner(
                 self.seq_lens.cpu[:num_reqs],
                 self.dcp_world_size,
                 self.dcp_rank,
-                self.parallel_config.cp_kv_cache_interleave_size,
+                self.parallel_config.dcp_kv_cache_interleave_size,
             )
             self.dcp_local_seq_lens.cpu[num_reqs:].fill_(0)
             self.dcp_local_seq_lens.copy_to_gpu(num_reqs_padded)
@@ -1834,6 +1887,15 @@ class GPUModelRunner(
                 :num_reqs_padded
             ]
 
+        if self.pcp_world_size > 1:
+            pm = self.pcp_manager
+            cm_base.pcp_allgather_restore_idx = pm.pcp_allgather_restore_idx.gpu[
+                : pm.padded_total
+            ]
+            cm_base.global_num_scheduled_tokens = pm.global_num_scheduled_tokens.gpu[
+                :num_reqs_padded
+            ]
+
         if logits_indices is not None and self.cache_config.kv_sharing_fast_prefill:
             cm_base.num_logits_indices = logits_indices.size(0)
             cm_base.logits_indices_padded = self._prepare_kv_sharing_fast_prefill(
@@ -3304,7 +3366,7 @@ class GPUModelRunner(
         self,
         num_tokens_padded: int,
         num_reqs_padded: int,
-        num_tokens_unpadded: int,
+        num_tokens: int,
         ubatch_slices: "UBatchSlices | None" = None,
     ) -> tuple[
         dict[int, torch.Tensor] | None,
@@ -3316,7 +3378,7 @@ class GPUModelRunner(
         Args:
             num_tokens_padded: Total number of tokens (padded)
             num_reqs_padded: Total number of requests (padded)
-            num_tokens_unpadded: Actual number of tokens (unpadded)
+            num_tokens: Actual number of tokens for this rank
             ubatch_slices: Optional ubatch slicing info for DBO
 
         Returns:
@@ -3331,6 +3393,10 @@ class GPUModelRunner(
         ):
             return None, None
 
+        global_num_tokens = (
+            self.pcp_manager.global_total if self.pcp_world_size > 1 else num_tokens
+        )
+
         def _get_slot_mapping(kv_cache_gid: int):
             assert num_reqs_padded is not None and num_tokens_padded is not None
             kv_cache_spec = self.kv_cache_config.kv_cache_groups[
@@ -3348,7 +3414,7 @@ class GPUModelRunner(
 
             # Fill unused with -1. Needed for reshape_and_cache in full cuda
             # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID
-            slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)
+            slot_mapping[global_num_tokens:num_tokens_padded].fill_(-1)
 
             return slot_mapping
 
@@ -3357,6 +3423,12 @@ class GPUModelRunner(
             for gid, _ in enumerate(self.kv_cache_config.kv_cache_groups)
         }
 
+        if self.pcp_world_size > 1:
+            slot_mappings_by_gid = {
+                gid: self.pcp_manager.pad_slot_mapping(slot_mapping[:global_num_tokens])
+                for gid, slot_mapping in slot_mappings_by_gid.items()
+            }
+
         slot_mappings_by_layer: dict[str, torch.Tensor] = {}
         for gid, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups):
             slot_mapping = slot_mappings_by_gid[gid]
@@ -3443,14 +3515,19 @@ class GPUModelRunner(
             req_ids = self.input_batch.req_ids
             tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids]
             num_scheduled_tokens_np = np.array(tokens, dtype=np.int32)
-            max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
-            num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens
 
-            logits_indices, spec_decode_metadata = self._prepare_inputs(
-                scheduler_output,
-                num_scheduled_tokens_np,
+            # num_tokens is local when PCP > 1, global otherwise
+            logits_indices, spec_decode_metadata, num_tokens = self._prepare_inputs(
+                scheduler_output, num_scheduled_tokens_np
             )
 
+            if self.pcp_world_size > 1:
+                max_num_scheduled_tokens = int(
+                    self.pcp_manager.local_num_scheduled.max()
+                )
+            else:
+                max_num_scheduled_tokens = int(num_scheduled_tokens_np.max())
+
             cascade_attn_prefix_lens = None
             # Disable cascade attention when using microbatching (DBO)
             if self.cascade_attn_enabled and not self.parallel_config.use_ubatching:
@@ -3468,7 +3545,7 @@ class GPUModelRunner(
                 num_tokens_across_dp,
                 cudagraph_stats,
             ) = self._determine_batch_execution_and_padding(
-                num_tokens=num_tokens_unpadded,
+                num_tokens=num_tokens,
                 num_reqs=num_reqs,
                 num_scheduled_tokens_np=num_scheduled_tokens_np,
                 max_num_scheduled_tokens=max_num_scheduled_tokens,
@@ -3535,18 +3612,21 @@ class GPUModelRunner(
             slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
                 num_tokens_padded=num_tokens_padded
                 if pad_attn or has_separate_kv_update
-                else num_tokens_unpadded,
+                else num_tokens,
                 num_reqs_padded=(
                     num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs
                 ),
-                num_tokens_unpadded=num_tokens_unpadded,
+                num_tokens=num_tokens,
                 ubatch_slices=ubatch_slices_padded,
             )
 
+            attn_num_tokens_padded = (
+                num_tokens_padded if pad_attn or self.pcp_world_size > 1 else None
+            )
             attn_metadata, spec_decode_common_attn_metadata = (
                 self._build_attention_metadata(
-                    num_tokens=num_tokens_unpadded,
-                    num_tokens_padded=num_tokens_padded if pad_attn else None,
+                    num_tokens=num_tokens,
+                    num_tokens_padded=attn_num_tokens_padded,
                     num_reqs=num_reqs,
                     num_reqs_padded=num_reqs_padded if pad_attn else None,
                     max_query_len=max_num_scheduled_tokens,
@@ -3624,6 +3704,11 @@ class GPUModelRunner(
                 hidden_states = model_output
                 aux_hidden_states = None
 
+            if self.pcp_world_size > 1:
+                hidden_states = self.pcp_manager.restore_hidden_states(
+                    hidden_states,
+                    num_tokens,
+                )
             if not self.broadcast_pp_output:
                 # Common case.
                 if not get_pp_group().is_last_rank:
@@ -4822,7 +4907,7 @@ class GPUModelRunner(
         slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
             num_tokens_padded=num_tokens,
             num_reqs_padded=num_reqs_padded,
-            num_tokens_unpadded=num_tokens_unpadded,
+            num_tokens=num_tokens_unpadded,
             ubatch_slices=ubatch_slices_padded,
         )
 
@@ -5233,7 +5318,7 @@ class GPUModelRunner(
 
         # Add `is_profile` here to pre-allocate communication buffers
         hidden_states, last_hidden_states = self._dummy_run(
-            self.max_num_tokens, is_profile=True
+            self.max_num_tokens // self.pcp_world_size, is_profile=True
         )
         if get_pp_group().is_last_rank:
             if self.is_pooling_model: