vLLM PCP/DCP 技术笔记

背景

vLLM 中引入了两个上下文并行维度——Prefill Context Parallel (PCP) 和 Decode Context Parallel (DCP)。本文是源码分析记录，覆盖配置入口、进程组初始化、通信模式、KV cache 布局、LSE 合并机制、各 attention backend 支持情况，以及一些数学推导。

设计动机

文档 context_parallel_deployment.md 给出了清晰的目标分工：

• Prefill 阶段：算力受限，目标是降低 TTFT —— 沿 query 序列维 T 切分。
• Decode 阶段：访存受限，单步 query 极少但 KV cache 很大 —— 沿 KV cache 的序列维 T 切分（不增加 GPU 数量，仅消除 TP 复用导致的 KV 复制）。

因此 vLLM 把它们拆成两个独立的并行维度：PCP 和 DCP。

一、配置层

ParallelConfig 中的关键字段：

校验逻辑

在 _validate_parallel_config 中：

tp_size % dcp_size == 0                                  # DCP 必须整除 TP
block_size % cp_kv_cache_interleave_size == 0             # 交错粒度对齐
if dcp_comm_backend == 'a2a': dcp_size > 1               # a2a 需要多卡

非 MLA 模型还有额外约束：

if decode_context_parallel_size > 1 and not use_mla:
    tp_size > total_num_kv_heads                          # TP 必须大于 KV head 数
    max_dcp_size = tp_size // total_num_kv_heads
    dcp_size <= max_dcp_size                              # DCP 不超过复制因子
    num_q_per_kv % dcp_size == 0                          # Q per KV 能被 DCP 整除

命令行入口

来自 arg_utils.py：

二、进程组与初始化

world_size 公式被改了

在 ParallelConfig.__post_init__ 中：

self.world_size = (
    self.pipeline_parallel_size
    * self.tensor_parallel_size
    * self.prefill_context_parallel_size   # ← 新增维度
)

注意：

• decode_context_parallel_size 不在乘积里 —— DCP 只是把已有 TP 组的卡再切一刀，不需要额外 GPU。
• PCP 是真实占用 GPU 的并行维度，而 DCP 不是。

rank 网格从 3D 扩展到 5D

在 initialize_model_parallel() 中，原本是 (DP, PP, TP) 三维，现在变成了：

# 布局顺序: ExternalDP × DP × PP × PCP × TP
all_ranks = torch.arange(world_size).reshape(
    -1,
    data_parallel_size,
    pipeline_model_parallel_size,
    prefill_context_model_parallel_size,   # ← 插在 PP 和 TP 之间
    tensor_model_parallel_size,
)

整体维序是 ExternalDP × DP × PP × PCP × TP。

各 group 的形成方式

所有 group 都是用 all_ranks 这个 5D 张量做 transpose + reshape 得到：

三、Group 关系图解

只看 transpose/reshape 的确难懂。下面用一个具体配置把所有 rank 摆出来。

假设配置

PP=2, PCP=2, TP=4, DCP=2 (DP=1, ExternalDP=1)
→ world_size = PP × PCP × TP = 2 × 2 × 4 = 16
→ 16 张 GPU，rank 编号 0..15

按 C-order（最后一维变化最快）展开，每个 rank 的坐标如下：

记忆要点：TP 在最里层（变化最快），PCP 包住 TP，PP 包住 PCP。

TP 组：固定 (PP, PCP)，把 TP 维收集起来

TP 组 0: [0, 1, 2, 3]    ← PP=0, PCP=0
TP 组 1: [4, 5, 6, 7]    ← PP=0, PCP=1
TP 组 2: [8, 9,10,11]    ← PP=1, PCP=0
TP 组 3: [12,13,14,15]   ← PP=1, PCP=1

同一个 TP 组内的 4 张卡组成一个张量并行域，模型权重在 head/hidden 维被切成 4 份。

DCP 组：TP 组内部再切

DCP 组 0: [0, 1]      ┐
DCP 组 1: [2, 3]      ┘  ← 都在 TP 组 0 内
DCP 组 2: [4, 5]      ┐
DCP 组 3: [6, 7]      ┘  ← 都在 TP 组 1 内
DCP 组 4: [8, 9]
DCP 组 5: [10,11]
DCP 组 6: [12,13]
DCP 组 7: [14,15]

关键关系：DCP ⊂ TP。一个 TP 组（4 卡）被切成 tp_size/dcp_size = 2 个 DCP 子组。每张卡同时是某个 TP 组的成员，也是其中一个 DCP 子组的成员。

PCP 组：固定 (PP, TP)，跨 PCP 维收集

PCP 组 0: [0, 1, 2, 3,  8, 9, 10, 11]    ← PCP=0 的所有 rank（跨 TP 和 PP）
PCP 组 1: [4, 5, 6, 7, 12, 13, 14, 15]   ← PCP=1 的所有 rank

PCP 组横跨多个 TP 组甚至 PP 组，通信范围比 DCP 大。这也解释了为什么 PCP 的通信模式（ring-attention 或 DualChunkSwap）比 DCP 更复杂——它需要跨更大范围的 rank 交换 KV 数据。

四、PCP 实现现状

main 分支：只有基本脚手架

main 分支上 PCP 的状态是"basic support"——只有输入准备部分搭好了：

• 进程组创建
• 配置字段和校验
• KV cache 切分（slot_mapping）
• 调度器适配

但 attention 计算还没接上：

class AttentionImplBase(ABC, Generic[T]):
    supports_pcp: bool = False    # 所有 backend 都是 False

if pcp_size > 1:
    assert layer_impl.supports_pcp, "PCP requires attention impls' support"

也就是说当前任何模型一旦真打开 --prefill-context-parallel-size 2，这个 assert 就会爆。

pcp-alt 分支：完整实现

origin/pcp-alt 分支（Lucas Wilkinson, 2026-03-02, commit a97bd607a）有完整 PCP 计算实现，采用 DualChunkSwap 范式而非 ring-attention。

commit message 原文：

Add PCP with DualChunkSwap token partitioning for balanced prefill computation 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)

为什么不是 ring-attention？ vLLM 走的是"双 chunk 互换 + KV 一次性合并"这条路，避免了 ring-attention 那种 N-1 轮 P2P 的复杂性，换取实现简单与确定的两种通信形态：

• Case 1（DCP = PCP）：每张卡负责自己的 query chunk + 自己的 KV chunk，partial out + LSE 用 all-reduce + LSE 合并搞定 —— 一次集合通信。
• Case 2（DCP = TP × PCP）：在 TP 内做完整 all-gather 拿到所有 KV，再 all-reduce + 切片 —— 两次集合通信但布局简单。

新增的核心组件：

docs/design/dcp_communication_patterns.md         | 350 +  ← 设计文档
vllm/v1/worker/cp_utils.py                        | 200+  ← PCPManager + DualChunkSwap
vllm/v1/attention/backends/flash_attn.py          | DCP+PCP 集成
vllm/v1/attention/backends/flashinfer.py          | DCP+PCP 集成
vllm/v1/attention/backends/mla/mla_attention.py   | DCP+PCP 集成
tests/distributed/test_context_parallel.py        | 端到端测试

五、DCP 完整流程

1. 配置阶段

CLI: --decode-context-parallel-size N
→ ParallelConfig.decode_context_parallel_size = N
→ 校验: tp_size % dcp_size == 0, block_size 对齐等

2. 进程组创建

_DCP 是 _TP 的子组。以 TP=4, DCP=2, 单节点为例：

TP group:  [0, 1, 2, 3]
DCP group: [0, 1] 和 [2, 3]   ← TP 组内连续切

每个 rank 仍是 TP rank 0/1/2/3，但同时也是某个 DCP group 的成员。所有 DCP rank 共享同一份 TP 参数，只在 KV cache 上分工。

3. KV cache 物理布局

DCP 不改变 KV cache 张量形状，改变的是"哪些 token 写到哪张卡"。

核心在 _compute_slot_mapping_kernel：

virtual_block_size = block_size * TOTAL_CP_WORLD_SIZE   # 虚拟块跨所有 CP rank
voff = token_pos_in_seq // virtual_block_size           # 虚拟块索引
cp_rank = (token_pos_in_seq // interleave) % cp_world_size  # 哪张 rank 负责
block_offset = voff * cp_world_size * block_size / interleave + ...

cp_kv_cache_interleave_size 控制交错粒度：

• = 1（token 级交错）：相邻 token 轮流写到不同 rank，负载最均衡
• = block_size（block 级交错）：整块连续 token 写到同一 rank，cache 友好

4. Decode 路径的 Attention 计算

为什么 KV 不需要 AllGather

DCP 的本质是利用在线 softmax 的可加性：

softmax(QK^T)V = LSE-merge[partial(Q, K0, V0), partial(Q, K1, V1), ...]

每个 rank 只用自己手里的 (Ki, Vi) 段算 partial output 和 partial LSE，在 head 维 reduce 把 partial 拼起来。KV 永远不出本卡。

数据量对比（设 B=batch, H=num_heads, D=head_dim, L=context_len, T=num_decode_tokens）：

Decode 时 L 通常几千到几十万，而 T 只有 B 那么大，Q 比 KV 小三四个数量级。DCP 的本质就是用"Q 的 AllGather + LSE 合并"换掉昂贵的"KV 的 AllGather"。

核心计算流程

DCP 在 decode 时的完整数据流（ag_rs 后端）：

# 1. Q AllGather：每张卡补齐完整 head
mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)

# 2. 用本地 KV cache 算 partial attention
#    注意：第二个参数直接传本地 kv_cache，没有 gather
attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)

# 3. LSE-aware 合并 partial 输出
attn_out = cp_lse_ag_out_rs(attn_out, lse, get_dcp_group(), ...)

5. Prefill 路径的 Attention 计算

Prefill 当前 token 要看到历史所有 KV，而历史 KV 散在 dcp_size 个 rank 上，所以必须把历史 KV 拿到本卡。但 vLLM 的实现避免了一次性 all_gather 整段历史：

# _context_parallel_compute_prefill_context
for chunk in split_kv_into_chunks(kv_cache, dcp_size):
    chunk_kv = all_gather(chunk)           # 流式拉取
    partial_out, lse = attn(q, chunk_kv)   # 算这一段
    running_out, running_lse = online_merge(running_out, running_lse, partial_out, lse)

这就是在线 LSE 合并——每拉一段就 merge，不等全部到齐。

6. 输出合并与后续层

完整数据流图（ag_rs 后端）：

每张 dcp rank 的 partial_out [B, 8, D] + partial_lse [B, 8]
      │
      │ ① AllGather LSE    ← 通信（量小：B*8 个 float）
      ▼
  all_lses [2, B, 8]
      │
      │ ② correct_attn_out ← 纯本地计算（Triton kernel）
      ▼
  out_weighted [B, 8, D]
      │
      │ ③ ReduceScatter (dim=1, op=SUM)
      ▼
  final_out [B, 4, D]       ← head 切回 TP 状态（即 allgather 之前的 head 数）
      │
      ├── @ W_o_local      ← 每张卡跟自己的 o_proj 切片相乘
      │
      └── TP AllReduce     ← 拼成完整 hidden_state

DCP 结束后，每张卡直接拿着 [B, 4, D] 跟自己那份 o_proj 权重相乘，然后参与正常 TP AllReduce。全过程不需要额外的 head AllGather。DCP 对后续层完全透明——后续层看到的就是"一张普通 TP rank 该有的 attention output"。

六、MLA 与 GQA/MHA 在 DCP 上的差异

MLA：天然适配

MLA 的 KV cache 形状为 [num_blocks, block_size, kv_lora_rank + qk_rope_head_dim]，没有 head 维。它存的是 latent KV（一段被压缩的低秩表示），所有 head 共享同一份 latent KV，通过权重矩阵 W_UK、W_UV 在线展开。

kernel 实际跑的就是 MQA（Multi-Query Attention）：多个 head 的 Q × 1 个共享的 latent KV。在 MLA 的 Triton kernel 中：

kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)  # 强行加一个 head=1 的维度
decode_attention_fwd(q, kv_c_and_k_pe_cache, ...)        # Run MQA

核心洞察：Q 的 head 拼齐（AllGather 后）与 KV 按 token 切（DCP 分工）这两个维度天然正交，不存在 head 不匹配问题。这也是为什么 11 种 MLA backend 全部通过统一的 forward_mqa 接口接入 DCP。

GQA/MHA：需要约束

GQA/MHA 的 KV cache 有 head 维。当 tp_size > num_kv_heads 时，KV head 在 TP 内被复制——这就是三个约束的根源。

举例：num_q=64, num_kv=8, TP=16, DCP=2

TP=16 > num_kv=8       → 复制因子 R = 16/8 = 2
                       → 每个 KV head 被 2 张 TP rank 复制
                       → 每张 TP rank 持有 1 个完整 KV head（没切 head）
                       → 每张 TP rank 持有 4 个 Q heads（都属于这 1 个 KV head）

DCP 组 = [tp0, tp1]    两张卡持有同一个 KV head 副本（都是 kv head 0）
                       但是不同的 token 段（S0、S1）
                       q heads 各 4 个（同 kv0 服务的 8 个 q 切一半）

Q AllGather(dim=1) →   每张卡拿到 8 个 q heads，全是 kv head 0 服务的
kernel 看到的：         num_qo=8, num_kv=1, group_size=8 ← 完美匹配

这个 group_size=8 就是模型原始定义的 GQA group_size。DCP 的 q allgather 干的事，本质是把 TP over-shard 拆碎的 GQA group 重新拼回模型原始的 group_size，然后让 kernel 跑一次"语义完整、token 不完整"的 attention，再用 LSE 把 token 维合回去。

三个视图下的形状对比

每张卡相当于做了 MQA

不管是 MLA 还是 GQA + DCP，每张 dcp rank 跑的 attention kernel 视角下都是 MQA 形态（一个 KV head 服务一组 Q heads）。但 MQA 形态不是 DCP 引入的，是 TP over-shard 的产物。DCP 在已有的 MQA 形态上做了两件事：

1. q allgather：把被 TP over-shard 撕碎的 group 拼回完整 group_size
2. token 维分工 + LSE 合并：让 dcp 同组的多张卡分担 KV 长度

这也回头解释了：

• 为什么 MLA（latent KV，天然 MQA）是 DCP 的 best fit
• 为什么 GQA + DCP 必须 tp > num_kv（强行制造 MQA 形态）
• 为什么纯 MHA 上 DCP 几乎用不起来——num_kv = num_q，要做到 MQA 形态需要 tp > num_q，通常做不到

数学上 DCP 是否真的要求 num_kv=1？

不一定。 num_kv=1 是 vLLM 当前工程实现的副产品。数学上 DCP 只要求 dcp 组内所有 rank 在 head 维上完全一致——即 dcp 组内的 rank 是"在 KV head 维上互为副本，在 KV token 维上互补"。vLLM 从 TP 维切 DCP 组 + 约束 tp_size > num_kv_heads 导致了 num_kv_per_card=1。

如果换个切法（比如让 DCP 组跨越 DP 维度而不是 TP 维度），那些 rank 天然持有相同 TP shard，KV head shard 也相同，就可以支持 num_kv_per_card > 1。

七、LSE（Log-Sum-Exp）合并机制

数学推导

从标准 attention 公式开始。设 query q（一个 token）对 N 个 KV token 做 attention：

score_j = q · k_j / sqrt(d)          j = 1..N

            exp(score_j)
alpha_j = ─────────────────          softmax 权重
           Σ_i exp(score_i)

output = Σ_j alpha_j · v_j          加权求和

展开写：

         Σ_j exp(score_j) · v_j
output = ───────────────────────
            Σ_j exp(score_j)

定义 LSE：

LSE = log( Σ_j exp(score_j) )

就是 softmax 分母取 log。名字就是字面意思：Log of Sum of Exp。

为什么需要 LSE

假设 N 个 KV token 切成 A（1..M）和 B（M+1..N）两段，分别在两张卡上算：

卡 0: output_A = Σ_{j∈A} exp(score_j)·v_j / Σ_{j∈A} exp(score_j)
     LSE_A    = log( Σ_{j∈A} exp(score_j) )

卡 1: output_B = Σ_{j∈B} exp(score_j)·v_j / Σ_{j∈B} exp(score_j)
     LSE_B    = log( Σ_{j∈B} exp(score_j) )

合并得到全局 output：

分子   = exp(LSE_A) · output_A + exp(LSE_B) · output_B
分母   = exp(LSE_A) + exp(LSE_B)

                 exp(LSE_A) · output_A + exp(LSE_B) · output_B
output_全局 = ─────────────────────────────────────────────
                      exp(LSE_A) + exp(LSE_B)

数值稳定版（防止 exp 溢出）：

m = max(LSE_A, LSE_B)
w_A = exp(LSE_A - m)
w_B = exp(LSE_B - m)

          w_A · output_A + w_B · output_B
output = ────────────────────────────────
                 w_A + w_B

静态合并 vs 在线合并

在线合并的算法（每来一段更新一次）：

初始状态（第一段）：
  running_out = out_0
  running_lse = lse_0

来了第二段 (out_1, lse_1)：
  m = max(running_lse, lse_1)
  w_old = exp(running_lse - m)
  w_new = exp(lse_1 - m)
  running_out = (w_old * running_out + w_new * out_1) / (w_old + w_new)
  running_lse = m + log(w_old + w_new)

来了第三段 (out_2, lse_2)：同上公式，running 当"旧的"。

这就是 FlashAttention 论文的核心创新——不存完整 attention matrix，一个 tile 一个 tile 地算，靠 LSE 在线合并。

切 Query 不切 KV：不需要 LSE

如果只切 Query 而 KV 完整：

rank 0: Attn(Q[0:T/2],  KV[0:N])  → output[0:T/2]   ← 完整的！
rank 1: Attn(Q[T/2:T],  KV[0:N])  → output[T/2:T]   ← 完整的！

每个 query token 对完整 KV 做 attention，softmax 分母涵盖了全部 N 个 KV token，没有被拆开。不需要 LSE，不需要合并，各 query 独立计算、结果独立正确。

这也是为什么"纯切 Q"的 parallel 方案（比如朴素的 sequence parallel prefill）在通信上比 DCP 简单。

LSE 的本质适用条件

只要 softmax 的分母被拆成多段分别计算，就需要 LSE 来合并。

形式化：对于某个 query token q，全局 softmax 分母是 Z = Σ_{j=1}^{N} exp(q·k_j)。如果 Z 被拆成 Z = Z_A + Z_B + …（不同段 KV 各算一部分指数和），单独每段的 output 不知道"自己占全局的比例"，必须用 LSE 还原这个比例。

LSE 的所有使用场景

八、Attention Backend 支持 DCP 一览

非 MLA（GQA/MHA）

这两个 backend 的 DCP 适用条件就是前文说的约束：tp_size > num_kv_heads 且 dcp_size ≤ tp_size / num_kv_heads。

MLA 类（全部支持，11 种 impl）

通过 forward_mqa(q, kv_cache, ...) -> (out, lse) 接口统一接入 DCP，MLA wrapper 统一处理 Q AllGather + LSE 合并，具体 impl 只需要实现这个接口，所以新增 MLA backend 基本可以"白嫖"DCP。

不支持 DCP 的 Backend

Triton Attn、Flex Attention、CPU Attn、ROCm AITER FA、TurboQuant Attn——这些完全找不到 DCP 相关引用。

总结

1. PCP 真实占用 GPU，使用 DualChunkSwap 计算范式，当前完整实现在 origin/pcp-alt 分支（main 分支只有基本脚手架）。PCP 不是 ring-attention，而是"双 chunk 互换 + KV 一次性合并"。

2. DCP 不增加 GPU，复用 TP 组 rank 在 token 维分工。核心是 LSE 合并机制——用"Q 的 AllGather + LSE 合并"换掉昂贵的"KV 的 AllGather"。DCP 对后续层完全透明。

3. MLA 对 DCP 天然友好；GQA/MHA 需要 TP over-shard + tp_size > num_kv_heads 配合才能干净工作，原因是 vLLM 从 TP 维切 DCP 组的工程决策。

4. LSE 合并是 attention 中 KV 维被切分时的通用数学工具，从 FlashAttention 到 DCP 到 Ring Attention，底层都是同一套在线 / 静态合并机制。