DeepInfer 源码拆解（二十五）：上下文并行 CP 全链路——zigzag 切分、mirror 布局与注意力里的 KV 收集

从哪里接着讲

第二篇讲过异步调度 ZOS，第三篇讲过 KV 池与 YoCo 截断，第十一篇讲过注意力族与压缩 KV，第十三篇把 DP Attention、Wide-EP、CP、ENP 四条并行轴一次摆开。那篇对 CP 的结论是「代码里存在、生产暂不用」：建组、缓冲区、zigzag 切分、前向里三处 gather 都说到了，但每个机制都只点到为止。

这一篇把 CP 单独拎出来，沿着一次开启 CP 的 prefill 从入口走到出口：

1. 元数据准备阶段，cp_mirror_len 是怎么算出来的，每个 rank 为什么要向上取整。
2. prepare_cp_mirror 这一个内核同时产出四样东西——mirror 索引、逆索引、ks、ke，它的并行化思路是什么。
3. 模型入口处的重排、出口处的逆重排，以及 positions、ngram_token_ids、valid_token_mask 这些字段为什么要跟着换布局。
4. 注意力层的 CP 分支：三种 KV 来源、KV gather 在层内的发起顺序、_make_cp_gather_getter 返回闭包的含义。
5. 压缩器怎么在 CP 下保证「原始相邻 token 配对」的正确性，为什么逆索引被传进了内核而不是在上层先重排。
6. YoCo 截断与 CP 叠加时，为什么还要再做一次「compact mirror」。
7. 这条路径的通信账、显存账，以及实现里几处容易被误读的细节。

（图源：DeepInfer 官方仓库 doc/images/architecture_overview.png）

读之前需要的最小背景：上下文并行（Context Parallel，CP）是把一条请求的 token 序列沿序列维切开，交给多个 rank 各算一段 query；每个 rank 只算自己那段的 query，但 causal attention 要求它能看见自己之前所有位置的 K/V。本站 MLA 完全拆解 讲过 MLA 的 KV 压缩结构，YoCo 与 KV 内存 讲过 DeepInfer 的两本块账，DeepEP 完全拆解 讲过 BucketBuffer 的异步通信接口——CP 的缓冲区直接复用后者的实现，下文会反复用到这三篇的结论。

为什么是「mirror」而不是「连续切段」

最朴素的 CP 切法是连续切：长度 16K 的 prompt、CP=4，rank 0 拿第 0–3K 个 token，rank 1 拿第 4–7K，依此类推。每段长度一样，看起来均衡。

问题出在因果注意力上。rank 0 处理的是序列最前面的 query，它只需要看这 4K 个 token 的 K/V；rank 3 处理的是最后 4K 个 query，却要看全部 16K 的 K/V。当各 rank 在注意力里把 K/V 收集成完整序列时，rank 0 要收集 16K 行却只用前 4K 行有效，rank 3 要收集 16K 行且全部有效——收集的数据量相同，但浪费程度不同。更深的问题在 chunked prefill 与长序列的混合：如果每条序列长度不同，连续切段会让某个 rank 拿到几条完整的长序列，其他 rank 全是短序列，负载差异被放大。

zigzag（也叫 load-balanced 切分）是常见解法：把序列切成 $2p$ 块（$p$ 是 CP 大小），rank $r$ 拿第 $r$ 块和第 $2p-1-r$ 块。rank 0 拿最前一块和最后一块，rank 1 拿第二块和倒数第二块。每个 rank 恰好拿到「一段靠前、一段靠后」的 query，于是每个 rank 的有效 K 覆盖范围——从自己最靠前那块 query 的起点到序列末尾——彼此接近，收集回来的 K/V 行里有效比例也更接近。

DeepInfer 把这种布局叫 mirror，因为原始 varlen 布局里的 token 被「镜像」到另一个张量上：源（原始 token 序列）和目标（本 rank 的 mirror 张量）靠一组索引对应，而不是靠指针切片。这么做的直接好处是 本 rank 的 token 在 mirror 张量里仍然是连续的——rank 0 的两块 token 被拼成一个连续区间，放在同一个张量的相邻行，于是后续每一层都只需处理一个首尾相接的 token 区间，不需要在 kernel 里维护「两段」的概念。

（图源：lucidrains/ring-attention-pytorch 仓库 README 的 ring.png，对应 Ring Attention 论文的块轮转流程。注意这张图对应的是「块沿环流转」的做法，与下面 DeepInfer 采用的「一次 gather 完整 KV、query 固定不动」不同，两者是同一个问题的两种解法，后续会在对照里说明。）

Ring Attention 的出发点与 DeepInfer 的 mirror 布局不同：它不把完整 KV 收集一次，而是让每个 rank 的 K/V 块沿环依次传给下一个 rank，每一步做一次块级注意力并在线更新 softmax 统计量，从而把序列长度扩展到几乎不受单卡显存限制的规模。代价是通信与计算必须严格流水，且每一步都要维护在线 softmax 的 m、l 与输出累加器。

沿用 varlen（变长）批次的表达方式：所有请求的 token 首尾相接拼成一个长度为 total_q 的一维张量，用 cu_seqlens_q 标记每条请求的边界。mirror 布局在这个一维轴上重新排：对每条请求，按 zigzag 切块，再按 rank 分组拼接。

元数据：cp_mirror_len 与向上取整

CP 打开后，每个 rank 要处理的 query 行数不再等于 total_q，而是 cp_mirror_len。它在 C++ 侧 prepare_prompt_base 里与 total_q 一起算出来：

1
int cp_size = args.cp_size;
2
int total_q = 0;
3
int cp_mirror_len = 0;
4
for (auto p_seq : args.seq_metadata_list) {
5
    total_q += p_seq->seqlen_q;
6
    cp_mirror_len += (p_seq->seqlen_q + cp_size * 2 - 1) / (cp_size * 2);
7
}
8
cp_mirror_len *= 2;
9
result["total_q"] = total_q;
10
result["cp_mirror_len"] = cp_mirror_len;

逐条看：每条请求先按 $2p$ 向上取整，算出一块的长度 $\mathrm{ps}$，再乘以 2——这正是本 rank 拿到的两块。把所有请求加起来，就得到本 rank 需要处理的全部 query 行数。

向上取整是整条链路的基石。它带来两个直接推论：

• 每个 rank 的 cp_mirror_len 只与批内请求的长度有关，与本 rank 的序号无关。 所有 rank 拿到的元数据完全相同，这也是为什么 CP 组里只需要 rank 0 算一份元数据、其余 rank 收广播（见后文）。
• 序列尾部短于一个块时，各 rank 拿到的行数会不相等。 均匀分割是理想的参考，但代码里对每条请求独立向上取整，所以当某条请求的长度不足以填满它分到的两块时，这些 rank 实际写入的有效行就变少，而 cp_mirror_len 仍按 $\mathrm{ps}$ 计满。典型例子是短请求混在长请求的批里：一条只有 3 个 token 的请求，CP=2 时 $\mathrm{ps}_i = \lceil 3/4 \rceil = 1$，$2p = 4$ 个块里只有块 0、1、2 非空，块 3 全空——于是 rank 1 的尾块那一行在 mirror_cp_index 里写 -1（padding），排在内侧、不参与计算。反过来，序列尾部若恰好凑不满最后一块，尾部行也会变成 padding。块数固定为 $2p$，超过序列长度的块就是空的。

cp_mirror_len 在 Python 侧通过 PrepareArgs.cp_size 传入（prepare_inputmetadata_dict 里只在 is_prompt 时赋值），decode 阶段不涉及 CP，因为 decode 时每个请求只生成一个 token，序列维切分的意义和代价都不划算。这一点和第十三篇的结论一致：CP 是 prefill 侧的优化。

prepare_cp_mirror：一个内核产出四张索引表

InputMetadata.prepare_for_mirror_cp() 在模型 forward 的最开头调用一次，它调用 prepare_cp_mirror 内核，返回四个张量：

1
mirror_cp_index, mirror_cp_reverse_index, ks_all, ke_all = prepare_cp_mirror(
2
    self.cu_seqlens_q,
3
    self.cu_seqlens_k,
4
    PG.context_parallel_group().size(),
5
    PG.context_parallel_group().rank(),
6
    kernel_ratios,
7
    self.total_q,
8
    self.cp_mirror_len,
9
)

文档字符串把四样东西的语义写得清楚：

• mirror_cp_index：长度 cp_mirror_len，第 $j$ 项是本 rank mirror 布局第 $j$ 行对应的原始 varlen 位置，-1 表示 padding。用 gather(x, mirror_cp_index) 就能把任何与 query 对齐的 varlen 张量重排成 mirror 布局。
• mirror_cp_reverse_index：长度 total_q，第 $i$ 项是原始位置 $i$ 的 token 在 gather 回来的完整张量里的位置。注意它的形状是 total_q，不是 cp_mirror_len：它是用来在「所有 rank 的 mirror 张量首尾拼接」之后重新排序的。
• ks / ke：对每个压缩率，记录 mirror 布局每一行 query 对应的 key 区间起止。

1
cu_q = [0, *accumulate(seq_lens)]
2
ps_all = [(sq + 2 * cp_size - 1) // (2 * cp_size) for sq in seq_lens]
3
cu_elems = [0, *accumulate(2 * ps for ps in ps_all)]
4
total_per_rank = cu_elems[-1]
5
cu_k_cn = [[0, *accumulate(max(1, sk // ratio) for sk in seq_ks)] for ratio in compress_ratios]
6
mirror_index: list[int] = []
7
reverse_index = [0] * cu_q[-1]
8
ks: list[list[int]] = [[] for _ in compress_ratios]
9
ke: list[list[int]] = [[] for _ in compress_ratios]
10
for i, (sq, sk) in enumerate(zip(seq_lens, seq_ks, strict=True)):
11
    ps = ps_all[i]
12
    for rank in range(cp_size):
13
        for half, block in enumerate((rank, 2 * cp_size - 1 - rank)):
14
            for j in range(ps):
15
                local_idx = ps * block + j
16
                if local_idx < sq:
17
                    reverse_index[cu_q[i] + local_idx] = rank * total_per_rank + cu_elems[i] + half * ps + j
18
                if rank == cp_rank:
19
                    mirror_index.append(cu_q[i] + local_idx if local_idx < sq else -1)
20
                    delta = sk - sq + local_idx + 1 if local_idx < sq else 0
21
                    for ri, ratio in enumerate(compress_ratios):
22
                        ks[ri].append(cu_k_cn[ri][i])
23
                        ke[ri].append(cu_k_cn[ri][i] + delta // ratio)

1
local_idx = t - cu_seqlens_q[seg]
2
blk = local_idx // ps
3
half = blk // cp_size
4
rank = blk + half * (cp2 - 1 - 2 * blk)
5
mirror_cp_reverse_index[t] = rank * cu_lens[num_ratios, batch_size] + cu_lens[num_ratios, seg] + half * ps + (local_idx - ps * blk)

1
if hasattr(self, "positions_cp"):
2
    self.positions_cp = self.positions
3
self.positions = gather(self.positions, self.mirror_cp_index)
4
if self.ngram_token_ids is not None:
5
    self.ngram_token_ids = gather(self.ngram_token_ids, self.mirror_cp_index)
6
if self.gate_force_random is not None:
7
    self.gate_force_random_cp = gather(self.gate_force_random, self.mirror_cp_index)
8
self.valid_token_mask = self.mirror_cp_index != -1

模型入口与出口：两次 gather 夹住整张网络

HFModel.forward 的形状很简洁：

1
if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
2
    input_metadata.prepare_for_mirror_cp()
3
    hidden_states = gather(hidden_states, input_metadata.mirror_cp_index)
4
# ... 逐层前向 ...
5
if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
6
    with userbuffer.cp_session():
7
        gathered_hidden_states = userbuffer.cp.all_gather(hidden_states)()
8
        hidden_states = gather(gathered_hidden_states, input_metadata.mirror_cp_reverse_index)

这里需要区分两个「gather」：

1. 入口与出口处的 gather 是 TileLang 实现的按索引重排（单卡内的 elementwise 拷贝），不是集合通信。入口处 gather(x, mirror_cp_index) 把本 rank 负责的两段 token 从原始位置搬到相邻行；出口处 gather(x_gathered, mirror_cp_reverse_index) 把收集回来的结果还原到原始 varlen 顺序。
2. 出口处的 all_gather 才是集合通信，它通过 CPManager 把每个 rank 的 mirror 张量按 rank 顺序拼成一个长张量，共 $p \cdot \texttt{cp\_mirror\_len}$ 行。

为什么只在入口出口各做一次集合通信，而不是每层都做？各层只需要算本 rank 的 query，query 之间的横向依赖只在注意力里存在，而注意力的 KV 收集已经单独处理（见下一节）。如果每层都把 hidden 切成原始布局、收集、再切开，通信次数会乘以层数——对于 40 层的模型就是 80 次，完全不可接受。所以中间整张网络的残差流保持在 mirror 布局，只在首尾跨越 mirror 与原始布局的边界。

这里有一个值得展开的计算：输出 hidden states 为什么需要在出口还原成原始顺序？ 因为采样层（站内 采样与惩罚）按请求取最后一个 token 的 logits，需要按原 varlen 顺序找到每条请求的末行；而且输出行数与 total_q 对齐，下游的 detokenizer 与 ZOS 后处理都假定这个形状。mirror 布局里请求的末行位置被打散了，所以必须还原。

all_gather 的返回值是闭包而不是张量，这里调用了 userbuffer.cp.all_gather(hidden_states)()，立即等待。为什么这里不推迟等待？因为出口处没有可重叠的计算：隐藏状态出来就是要采样了。推迟等待的价值在于注意力里那些「发起后立刻去做别的计算」的场景。

注意力层：一段代码覆盖三种 KV 来源

V4AttentionBase.forward 的 CP 分支集中在 with userbuffer.cp_session(enabled=use_cp): 块里。要理解它的结构，先要知道这一层可能产生哪几种 KV：

• 主 KV（main KV）：MLA 压缩后的主缓存，按 kv_compress_ratio 决定是否再次压缩（C1 不压缩、C2 两两压缩）。
• SWA KV（滑动窗口 KV）：用于局部窗口注意力，窗口大小为 window_size。
• 索引器 KV（index KV）：DSA 风格稀疏注意力里的 lightning indexer 键（站内 DSA 索引器与算子）。
• 局部索引器（has_local_index）：本层自带索引器，自己算 top-k；否则复用前一层索引源的 top-k 结果。

CP 下这些 KV 都必须覆盖完整序列，否则 attention 看不到前面的 K。这正是整段代码的第一句注释：

Without a local indexer, delay main KV storage to overlap CP all-gather.

1
with userbuffer.cp_session(enabled=use_cp):
2
    wait_projected_kv = _make_cp_gather_getter(projected_kv, use_cp) if not yoco_full_kv_stored else None
3

4
    q_a, kv_swa = self._project_q_and_swa(hidden_states, shared_hidden_bf16)
5
    ...
6
    kv_swa, wait_kv_swa = self._prepare_swa_kv(kv_swa, kv_cache.swa_cache, input_metadata, use_cp)
7

8
    q_b, index_q_b = self._project_q_b(q_a, input_metadata.positions)
9
    ...
10
    kv_ca_index = self._compute_kv_ca_index(
11
        wait_projected_kv, index_weights, index_q_b, kv_cache, input_metadata, cross_layer_runtime,
12
    )

1
def _make_cp_gather_getter(x, enabled):
2
    if x is None:
3
        return None
4
    if enabled:
5
        return userbuffer.cp.all_gather(x)
6
    return lambda: x

1
def all_gather(self, x: torch.Tensor) -> Callable[[], torch.Tensor]:
2
    assert self._session is not None, "all_gather requires an active CP session"
3
    if not x.is_contiguous():
4
        contiguous_x = self.get_inplace_tensor(x.shape, x.dtype)
5
        contiguous_x.copy_(x)
6
        x = contiguous_x
7
    handle = self._buffer.all_gather(x)
8
    gathered_shape = (self._size, *x.shape)
9
    return lambda: handle.wait().view(gathered_shape).flatten(0, 1)

1
out = userbuffer.cp.get_inplace_tensor(kv_swa.shape, kv_swa.dtype) if use_cp else None
2
if input_metadata.is_prompt:
3
    cast_args = CastArgs("e4m3", self.swa_cache_num_per_channels, round_sf=True, sf_format="ue8m0", cast_back=True)
4
    out_args = out
5
...
6
result = fused_norm_rope_hadamard_cast_store_kvcache(
7
    kv_swa,
8
    norm_args=NormArgs(...),
9
    rope_args=RopeArgs(input_metadata.positions, ...),
10
    cast_args=cast_args,
11
    out_args=out_args,
12
)
13
kv_swa = result.unflatten(-1, (1, self.qk_head_dim))
14
wait_kv_swa = userbuffer.cp.all_gather(kv_swa) if use_cp else None

1
total_k_swa_cur = (input_metadata.mirror_cp_reverse_index.shape[0] if use_cp else kv_swa.shape[0]) if self.use_swa else 0
2
...
3
if use_cp:
4
    kv_swa = gather(wait_kv_swa(), input_metadata.mirror_cp_reverse_index, out=kv_swa_cur)
5
else:
6
    kv_swa_cur.copy_(kv_swa)

1
use_cp = mirror_cp_reverse_index is not None
2
with userbuffer.cp_session(enabled=use_cp):
3
    if use_cp:
4
        main_k = self.kv_norm(
5
            main_k,
6
            out=userbuffer.cp.get_inplace_tensor(main_k.shape, main_k.dtype),
7
        )
8
        wait_main_k = userbuffer.cp.all_gather(main_k.unsqueeze(1))
9
        main_k = gather(wait_main_k(), mirror_cp_reverse_index)
10
fused_norm_rope_hadamard_cast_store_kvcache(
11
    main_k,
12
    norm_args=None if use_cp else NormArgs(self.kv_norm.weight, ...),
13
    ...
14
)

1
if yoco_full_kv_stored:
2
    use_cp = PG.context_parallel_group().size() > 1
3
    expected_rows = input_metadata.cp_mirror_len if use_cp else input_metadata.total_q
4
    assert kv_ca_index.shape[0] == expected_rows, (...)

1
if self.use_swa:
2
    kv_index = kv_swa_full_index
3
    if self.has_packed_kv:
4
        kv_index = kv_index_cat(kv_index, kv_ca_index, total_k_swa_cached, total_k_swa_cur)

压缩器：为什么必须先收集再压缩

V4.1 的 KV 压缩（C2）把每两个相邻 token 的 KV 合成一个压缩条目。Compressor.forward_pre_norm 的 CP 处理只有一句：

1
if self.kv_compress_ratio == 1:
2
    ...
3
    if input_metadata.mirror_cp_reverse_index is not None:
4
        projected_kv = gather(projected_kv.contiguous(), input_metadata.mirror_cp_reverse_index)
5
    return projected_kv

C1（不压缩）路径直接返回；C2 路径则把 mirror_cp_reverse_index 传给压缩内核：

1
kv_compressed = compress_and_update_state_prefill(
2
    score, latent, state_cache,
3
    input_metadata.state_block_idx,
4
    input_metadata.compressed_batch_mapping[self.ratio_key],
5
    input_metadata.cu_seqlens_q,
6
    input_metadata.cu_seqlens_k,
7
    input_metadata.cu_compressed_seqlens_q[self.ratio_key],
8
    self.kv_compress_ratio,
9
    input_metadata.mirror_cp_reverse_index,
10
)

为什么要收集？压缩的配对是按原始序列的相邻 token 进行的：token 0 与 1 配成一条、2 与 3 配成一条。zigzag 切分后，rank 0 的 mirror 张量里第 6 行紧邻第 7 行，但它们在原始序列里可能是第 2 段尾和第 3 段头，并不相邻；而真正相邻的第 3 段尾与第 4 段头分别在不同 rank 手里。若各 rank 各压各的，配对的边界会错、跨 rank 的配对根本看不到。

所以正确做法是：收集完整序列 → 按原始顺序压缩 → 再写回缓存。这解释了为什么 ks / ke 在完整序列坐标系里（前文提到），也解释了为什么 positions_c（压缩 token 的位置）不需要重排——压缩结果本身是按原始顺序排列的。

这里有一个容易误读的地方，值得单独说清楚：C2 路径并没有在上一层先把张量收集并重排好再交给压缩器，真正的重排发生在压缩内核内部。看 compress_and_update_state_prefill 的文档字符串，两个输入张量的形状标为：

1
score:  [total_q | cp_size * cp_mirror_len, dim], fp32
2
latent: [total_q | cp_size * cp_mirror_len, dim], fp32
3
mirror_cp_reverse_index: [total_q], int32,
4
    本 token 在 permuted（按 rank 拼接）score/latent 里的行号

也就是说，CP 打开时进入压缩内核的是每个 rank 的 mirror 张量按 rank 拼接后的完整序列（$p \cdot \texttt{cp\_mirror\_len}$ 行），行序是「rank0 的头块 → rank0 的尾块 → rank1 的头块 → …」，并不是原始序列顺序。参考实现 compress_and_update_state_prefill_ref 的取数逻辑把这个关系写得很直白：

1
else:
2
    # 本行在原始 query 序列中的位置
3
    orig_idx = pos - qstart + qs
4
    # 映射到按 rank 拼接后的行号
5
    idx = mirror_cp_reverse_index[orig_idx] if has_index else orig_idx
6
    score_val = score[idx, :]
7
    latent_slice[j, :] = latent[idx, :]

它在按原始序列顺序遍历每一个待压缩的位置 pos，但并不要求输入张量本身按原始顺序排列——而是用 mirror_cp_reverse_index[orig_idx] 现查现取。state_cache 的写入循环同理。这样压缩器的配对语义（原始相邻 token 配成一条）保持不变，跨 rank 的配对也能被正确拼起来，代价是内核里多一次索引间接寻址。

回到 forward 链路上看：projected_kv 的 all-gather 发生在 _compute_kv_ca_index → _store_main_kv → PackedWkv.forward_store_kv_cache → Compressor.forward 这条链中，但 gather 结果并没有被重排成原始顺序，而是以按 rank 拼接的布局直接送进压缩内核，由内核用逆索引完成「逻辑上的重排」。_store_main_kv 里 wait_projected_kv() 取到的正是这个按 rank 拼接的张量。

这样设计的原因很实际：压缩内核本来就要逐行读取 score/latent 并做 softmax 加权，多一个索引查表几乎不增加成本；而在上层再插一次全量按索引重排（又是一个只做数据搬运的算子），读写一遍 total_q × dim 的 fp32 张量，对一段很长的 prompt 来说是可观的显存带宽开销。把重排「折进」压缩器的访存里，相当于把它消掉了。

C1（不压缩）路径则没有这么幸运：它没有可以借用的「逐行读取」环节，于是只能显式 gather(projected_kv, mirror_cp_reverse_index) 一次，把按 rank 拼接的张量重排成原始顺序再返回。

顺带一提，这里能看到 CP 与 YoCo 的交汇。PackedWkv 的输入在 _store_yoco_full_kv 场景下来自 consumer_full_hidden——一个在 YOCO 截断边界之前捕获的「完整长度」隐藏状态。断言写着：

1
expected_rows = full_metadata.cp_mirror_len if use_cp else full_metadata.total_q
2
assert full_hidden.shape[0] == expected_rows, "consumer_full_hidden must be in the full (per-rank mirror) layout"

也就是说，即使是「完整长度」的隐藏状态，在 CP 下也是本 rank mirror 布局的，行数是 cp_mirror_len，不是 total_q。它需要在各 rank 上各算各自那段 KV，再把收集回来的 KV 拼成完整序列——这是下一节要展开的内容。

当 YoCo 截断遇到 CP：compact mirror

YoCo 的 prefill 截断（站内 两本块账与 YoCo）只对最新的 keep_tokens 个 query 做完整注意力计算，前缀部分只用来生成 KV 和做压缩。CP 下这两件事要同时成立：

• 前缀部分仍然要覆盖完整序列的 KV——
• 新 query 部分仍然要按 zigzag 切分、每个 rank 只算自己的 query。

结果就是两级镜像：full mirror（完整镜像，行数 cp_mirror_len）用来生成 KV；compact mirror（截断后的镜像，行数 yoco_consumer_cp_mirror_len）用来算 query。yoco.py 里有明确的实现：

1
mirror_cp_index_c, mirror_cp_reverse_index_c, ks_all, ke_all = prepare_cp_mirror(
2
    ...,
3
    self.yoco_consumer_cp_mirror_len,
4
)
5
view.mirror_cp_index = mirror_cp_index_c
6
view.mirror_cp_reverse_index = mirror_cp_reverse_index_c
7
view.cp_mirror_len = self.yoco_consumer_cp_mirror_len

先做一次 prepare_cp_mirror，对截断后的 query 长度重新算镜像索引；再对每个 rank 的完整镜像只取 compact 部分，gather 回来之后重新分配成 compact mirror 的顺序。这里用到 yoco_cp_boundary_fused：

1
contrib_index, gathered_reverse_index = yoco_cp_boundary_fused(
2
    cu_seqlens_q, cu_seqlens_q_c, cp_size, cp_rank, total_q_c, contrib_pad_len,
3
)
4
...
5
view.yoco_cp_boundary_index = gather(gathered_reverse_index, mirror_cp_index_c)

yoco_cp_boundary_fused 的参考实现从截断点 clip = sq - keep 开始遍历：

1
for rank in range(cp_size):
2
    for half, block in enumerate((rank, 2 * cp_size - 1 - rank)):
3
        for j in range(ps):
4
            local_idx = ps * block + j
5
            if clip <= local_idx < sq:
6
                reverse_entries.append((cu_c + local_idx - clip, rank, len(contribs[rank])))
7
                contribs[rank].append(cu_elem + half * ps + j)

只有 [clip, sq) 区间内的行才算「贡献」。contrib_pad_len 取各 rank 贡献行数的最大值，所以 gather 回来的结果在每个 rank 上是等长的（短的 rank 补 padding），逆索引以 rank * contrib_pad_len + ordinal 编码。

为什么截断让事情变复杂？ 因为截断前，每个 rank 的 mirror 行数是 cp_mirror_len；截断后，每个 rank 只保留 query 部分，但每条请求保留下来的 query 数量相同（keep_tokens），而 zigzag 切分下每个 rank 从截断区间里拿到的数量是不等的。所以需要「收集后重新按 compact 的 zigzag 切分」，而不是直接沿用 full mirror 的行顺序。contrib_pad_len 就是用来吸收这个不等的。

attn.py 的 _store_yoco_full_kv 与 forward 里各有一个 session：

1
# update_main_cache / _store_yoco_full_kv
2
with userbuffer.cp_session(enabled=use_cp):
3
    wait_projected_kv = _make_cp_gather_getter(projected_kv, use_cp)
4
    pre_rope_kv = self._store_main_kv(wait_projected_kv, kv_cache, full_metadata)
5

6
# forward
7
yoco_full_kv_stored = cross_layer_runtime.consumer_full_metadata is not None
8
if yoco_full_kv_stored:
9
    # With CP > 1 an inner cp session is opened here (before the main session below, with no nesting).
10
    self._store_yoco_full_kv(cross_layer_runtime, kv_cache)

注释里的 「with no nesting」 是硬约束。CPManager.session 里有 assert self._session is None，session 不可嵌套；因此 full-KV 那一段（在 _store_yoco_full_kv 内）必须先于主分支的 session 打开、并在主 session 打开之前关闭。代码的写法是先把 _store_yoco_full_kv 放在 with 之外，再进入主 session。

而 yoco.py 的 get_consumer_view 也有断言：

1
assert self.mirror_cp_index is not None, "prepare_for_mirror_cp must run before the compaction boundary"

也就是说，两种镜像有严格的先后：先 full mirror（入口处 prepare_for_mirror_cp），后 compact mirror（截断边界处）。顺序反了，full mirror 的索引就会建立在 compact 的坐标系上。

yoco_cp_boundary_index 与 DFlash 的特殊路径

DFlash 草稿头（站内 DeepInfer 的投机解码路径）在 prefill 时采集目标层特征，供 decode 做草稿。dflash.py 里有：

1
# prefill CP > 1 时 dflash_input 是 per-rank mirror 布局（cp_mirror_len 行），各层 all-gather 后按原 varlen 顺序写入。
2
mirror_cp_reverse_index = input_metadata.mirror_cp_reverse_index
3
if mirror_cp_reverse_index is None:
4
    ...
5
    assert dflash_input.shape[0] == input_metadata.cp_mirror_len

DFlash 的特征张量来自多个目标层，每个目标层都只算本 rank 的 query 行，所以它天然是 mirror 布局。写入草稿缓存时需要先 all-gather、再按逆索引还原到原始顺序——这与主模型出口的做法一致。如果目标层全在 YoCo 的 compact 半区，行数则取 consumer_metadata.cp_mirror_len（即 yoco_consumer_cp_mirror_len），在 deepseek.py 里能看到这段：

1
rows_c = consumer_metadata.cp_mirror_len if PG.context_parallel_group().size() > 1 else consumer_metadata.total_q
2
dflash_input = hidden_states.new_empty((rows_c, self.dflash_input_size))

这段代码也说明另一件事：CP 与投机解码、YoCo、结构化输出等特性不是正交的，每一处新特性都要问一句「这时的行数是 total_q 还是 cp_mirror_len」。

两种 CP 实现：按序列切分与按 head 切分

同一个「让每个注意力头看到完整上下文」的问题，在 DeepSpeed Ulysses 里有一个不同的落点：Ulysses 不做 KV 收集，而是按 attention head 切分——每个 rank 持有全部 token 的若干个 head，于是「一个 head 的注意力」仍然能看到完整序列，只需要在 Q/K/V 投影后按 head 做一次 all-to-all 重排。下图把两种切法放在一起：

（图源：DeepSpeed Ulysses 论文，arXiv:2309.14509）

两者的取舍很清楚：按 head 切分每层只做一次 Q/K/V 的 all-to-all，通信次数少；交换量与序列长度乘以 hidden 成正比，不过它与 KV 收集不同——不需要在 attention 内部反复搬 K/V。它的硬约束是 head 数必须能被并行度整除，而且对 MLA 这种「所有 head 共享一份 KV latent」的结构不友好：latent 无法按 head 拆开，切了也没有意义。按序列切分对 head 数没有约束，能处理 MLA/GQA，代价是每个 KV 源层都要交换一次与当前 chunk 长度成正比的 KV。DeepInfer 走的是后者。

元数据广播：为什么只有 rank 0 算

回到 DeepInfer。broadcast_tensor_dict 的输入侧有一段 CP-only 的逻辑：

1
if PG.context_parallel_group().rank() == 0:
2
    cp_broadcast_object(non_tensor_values)
3
else:
4
    for path, value in cp_broadcast_object(None).items():
5
        *parents, last = path.split(".")
6
        ...
7
        setattr(obj, last, value)
8
if tensor_list:
9
    cp_broadcast_multi_tensor(tensor_list)

这段逻辑与第二十一篇（CP 组对象广播）讲的是同一个机制：CP 组只在 rank 0 上构造完整元数据，非零 rank 通过共享内存对象通道收非张量字段，通过 NCCL 收张量。prepare_cp_mirror 本身不需要通信——因为它输入的是 cu_seqlens_q 与 cp_rank，而 cu_seqlens_q 在广播之后各 rank 已一致，cp_rank 本地可知。唯一有 rank 依赖的部分是 mirror_cp_index 的取值，而 prepare_cp_mirror 的 cp_rank 参数正好承担了这个差异：同样的输入在 rank 0 与 rank 3 上产出的 mirror_cp_index 不同，但 mirror_cp_reverse_index 完全相同——它的定义基于原始位置，与 rank 无关。

这也解释了两个索引在用途上的差异：

既然逆索引与 rank 无关，为什么每个 rank 还要各算一遍？因为内核计算开销很低（两次二分查找），而广播一张长度为 total_q 的 int32 表（16K token 即 64 KB）的开销并不更小；更重要的是，mirror_cp_index 本来就需要各 rank 分别算，一次 prepare_cp_mirror 同时产出四张表，没有理由为了省掉逆索引再单独设计一条广播。

同构层之间复用 gather 缓冲：use_shared_kv 与 is_kv_source

上面提到 KV 收集结果可以在同构层之间复用，这就引出 CP 视角下另一个容易被忽略的细节。

DeepInfer 把结构相同的层归成一组（SharedLayerTopology），组内只有「源层」负责产生 KV，后续消费者层共用同一份结果。源层的判定在 builder.py 里写得很直接：

1
is_kv_source_layer=(shared_topology is None
2
                    or shared_topology.layer_kv_map[runtime_layer_id] == runtime_layer_id),

即某一层如果被映射到自身，它就是这一段的源层。

对 CP 来说，这意味着 all-gather 只发生在源层。消费者层读取的是 cross_layer_runtime.loaded_prefill_kv——源层收集并重排后放进来的整段 KV，形状为 (total_kv, 1, qk_head_dim)。_prepare_prefill_kv 的文档字符串写得很明确：

A KV source starts a contiguous share-KV segment. Its following consumers reuse the buffer until the next KV source starts a new segment.

两个约束随之而来：

• 源层开始新段时，会先把上一段的 loaded_prefill_kv 置空，再分配新的缓冲区——否则新段会覆盖旧段正在被下游读取的数据。
• 消费者层要对齐 loaded_prefill_kv_main_offset：它必须等于本层期望的 main_kv_offset，否则直接断言失败。这说明同一段内各层的 SWA/压缩切分方式必须一致，不能出现某一层比源层多算或少算几行 K 的情况。

从通信账上看，这条复用把「每层一次 gather」降为「每个源层一次 gather」。若 40 层里只有 8 层是源层，CP 的 KV gather 次数就从 40 次降到 8 次——这是 CP 通信成本里最不容易被注意、但影响最大的一个乘数。

缓冲区与内存账

缓冲区由 userbuffer.allocate_cp 分配：

1
max_ag_bytes_per_session = max_num_tokens * max_ag_bytes_per_token * cp_group.size()
2
max_ag_bytes_per_session += max_num_ags_per_session * deep_ep.get_num_rdma_alignment()
3
alignment = deep_ep.get_num_allocation_alignment()
4
num_bytes = ((max_ag_bytes_per_session + alignment - 1) // alignment) * alignment

• max_num_tokens 是 max_num_prefilling_tokens，而后者在 finalize_model_config 里被除以了 CP 大小——每个 rank 的 token 预算只是总预算的一部分。
• max_ag_bytes_per_token 是模型侧逐层累加的最大值：V4AttentionBase.get_max_ag_bytes_per_token 返回 qk_head_dim * 2 字节，再加上 SWA 的 qk 维、投影 KV 维等。
• max_num_ags_per_session = 8 允许一次 session 里发起多条 gather（模型首尾各一条、注意力若干条），每条需要 RDMA 对齐。
• 整体再按分配对齐向上取整。

这块缓冲区在 DeepEP 的 BucketBuffer 上分配（deep_ep.BucketBuffer(group, num_bytes, explicitly_destroy=True)），因此 CP 的 gather 与 EP 的 dispatch/combine 共用同一套 RDMA 通路与分配器；但缓冲区本身仍然是从显存里划出来的，它不会凭空消失，只是可以和 EP 的缓冲区按同一套对齐规则管理。

可以粗略估算：以 max_num_prefilling_tokens = 32768（总预算 65536、CP=2）、每 token 每 rank 字节约 1 KB 计，CP=2 时缓冲区约 32768 × 1024 × 2 ≈ 64 MiB。对于 64 GB 张的量级不算什么，但它是在 prefill 预算里预先扣掉的——这也是官方选择生产暂不启用 CP 的原因之一，见下节。

一次具体批次的账：cp_mirror_len 与 padding 浪费

抽象公式放到具体批里更容易看出成本。假设一批里有两条请求，长度分别为 4096 与 300，CP=4：

cp_mirror_len = 1024 + 76 = 1100，而 total_q = 4396。每个 rank 只处理 1100 行 query，四条 rank 合计 4400 行——只比 total_q 多出 4 行 padding。这是因为长请求恰好整除，短请求的 padding 也只有 4 行，整体浪费很小。

但如果批里全是短请求，浪费比例会上升。例如 16 条长度 300 的请求、CP=8：每条 $\mathrm{ps}_i = \lceil 300/16 \rceil = 19$，每条 mirror 长度 38，cp_mirror_len = 608，而 total_q = 4800；每个 rank 608 行，8 个 rank 共 4864 行，多出 64 行。也就是说，短请求批的 padding 开销来自「每条请求都向上取整到 $2p$ 的整数倍」，批内请求越多、CP 越大，累计的取整误差越大。

这解释了一个工程上的选择：CP 更适合「少量、超长」的请求，而不是「大量、很短」的请求。后者用连续批处理和 DP 更划算；CP 的索引内核、mirror 重排与跨 rank 收集都是为长序列摊薄的设计。

通信账：CP 下一次 prefill 到底通信几次

在第十三篇给过一次粗账，这里按本篇的实际代码路径再数一遍（假设 40 层，每层是「主 KV + SWA + 索引器」的来源层，CP=2）：

如果不考虑层间共享，按「每层 3 次」计，40 层约 120 次 all-gather；实际由于源层复用会低不少，但仍在几十到上百次的量级，加上首尾与元数据，与 TP 在 40 层的通信次数相当，甚至更多。这是 CP 真正的代价：TP 的通信在每层一次（all-reduce），而 CP 的通信在每层多次（不同 KV 来源各一次）。对通信带宽有限的节点，这种差距会直接把 CP 的收益吃掉。

值得注意的是，CP 与 TP 可以组合：TP 切的是权重，CP 切的是序列，两者解决的问题不同。DeepInfer 当前不做 TP（第十三篇解释过原因：MoE 的参数账使 EP 更划算），所以 CP 若要启用，面对的是「没有 TP 可叠加」的场景——单卡模型放得下，只是 prefill 太慢。

KV 收集之后：收集的是「当前 KV」，不是完整主缓存

读到这里容易产生一个误解：既然每层都收集完整 KV，是不是意味着每层都把完整主 KV 写进了分页缓存？答案是否定的。

DeepInfer 的 prefill 注意力并不从分页缓存里读取完整主 KV。_prepare_prefill_kv 拼出来的 kv 由三部分组成：

1
[ SWA 缓存（已落盘的历史窗口） | 当前请求新算出的 SWA KV | 压缩后的主 KV 缓存 ]
2
     total_k_swa_cached              total_k_swa_cur            total_k_compressed

其中压缩后的主 KV 通过 load_paged_kv_cache 从 ca_cache 读出，当前 token 的 KV 则由本层现算。也就是说，主 KV 以压缩形式落盘，完整 KV 是每层临时拼出来的；CP 收集的是「当前这批 token 的投影 KV + 压缩缓存」，不是一份完整的主 KV 副本。

这对 CP 的意义在于：收集缓冲区的大小由 max_num_prefilling_tokens 决定，而不是由 max_model_len 决定。每条请求的当前 chunk 只有一部分 token，收集的 KV 行数与之成正比；历史部分已经在压缩缓存里，不需要跨 rank 搬运。若主 KV 不做压缩，CP 每层要搬的就是整条历史的 latent，通信量会随上下文长度线性增长，CP 的可用范围会小得多。

也正因为如此，_prepare_prefill_kv 里 use_shared_kv 的共享逻辑才有意义：同一段 KV 源层算出的收集结果，可以被后续若干消费者层复用，直到下一个 KV 源层开启新段。没有这层复用，每层都要重新收集一次完整 KV，CP 的通信成本会再翻几倍。

官方为什么暂不启用 CP

ParallelConfig 里虽然有 context_parallel_size，但官方部署文档与生产配置都没有开 CP。原因有两层：

1. 部署层：CP 会改变一个引擎内 KV、调度与元数据的所有布局约束——例如 cp_align = cp_size * 2 要求 chunked prefill 的每个 chunk token 数是 2 * cp_size 的倍数；max_num_prefilling_tokens 被除以 CP；--cp 必须整除 ep_size。所有这些约束都会渗透到调优参数里，而生产环境更倾向于用不受这些约束的机制解决问题。
2. 替代方案已存在：prefill 侧的长上下文压力交给 Prefill Load Balancer 按长度分档处理（把长 prompt 路由到专门的处理节点），decode 侧的问题由 Wide-EP 解决。在 64 卡以上 EP 已经让 decode 不再是瓶颈的前提下，再叠一层 CP 收益不如先把 EP 调好。

但代码是完整可用的：测试覆盖了 cp_size ∈ {2,4,8}，YoCo、DFlash、多模态都写了 CP 分支，说明这不是遗留代码，而是为「单卡放得下大模型、但 prefill 太慢」的场景预留的路径。理解这条路径，对读任何一个 CP 实现都用得上。

几处容易被误读的细节

一、use_cp 只在 prefill 下为真

attn.py 里：

1
use_cp = input_metadata.is_prompt and PG.context_parallel_group().size() > 1

_prepare_swa_kv 里 not input_metadata.is_prompt 分支直接 assert not use_cp。decode 时不走 mirror 布局，因为 decode 的每个请求只有一个 query token，序列维并行的收益极低。这一约束也影响到 CUDA Graph——decode 的图捕获里 use_cp 恒为假，路径稳定；prefill 没有图，布局可以动态。

二、gather 有两种含义

同名的 gather 在不同位置指代不同东西，读代码时最容易混淆：

第二类是最常见的，凡是 gather(x, input_metadata.mirror_cp_index) 或 gather(x_gathered, mirror_cp_reverse_index) 都是单卡内的按索引重排。文档字符串里写「gather(x, mirror_cp_index) reorders a varlen tensor into this rank’s CP-mirror layout」，这里的 gather 也是重排。

三、session 不可嵌套，闭包不可跨 session

CPManager.session() 的 assert self._session is None 与「闭包必须在 session 内消费」是两条硬约束。YoCo 路径的注释专门提醒「inner cp session … with no nesting」，就是因为 full-KV 的 session 与主分支的 session 若嵌套会触发断言。类似地，all_gather 返回的闭包如果在 session 外调用，拿到的可能是被后续 session 覆盖的缓冲区——这种错误不会在出错点报错，而是在结果里表现为静默的数值错误，排查代价极高。

四、_make_cp_gather_getter 在 CP=1 时返回恒等闭包

这看似是 trivial 的实现，却是让调用方不必写双路径的关键：

1
return lambda: x

有了它，_store_main_kv(wait_projected_kv, ...) 的写法在 CP 开关两种情况下完全一致，代价是一次 Python 闭包调用——在 prefill 的 CPU 开销里完全可以忽略。

五、cp_mirror_len 的向上取整和 cp_align 不一样

cp_align = cp_size * 2 约束的是 chunk 大小（调度器每个 chunk 的 token 数），cp_mirror_len 是每条请求切块后累计的行数。前者影响调度层面的分块，后者影响张量的行数。两者都用到 2 * cp_size，但含义完全不同：chunk 是调度单位，请求是可切分的对象。读调度器代码时（站内 C++ 调度器）看到 cp_align，不要和 cp_mirror_len 混在一起。

小结

DeepInfer 的 CP 路径可以压缩为一条数据流：

1
原始 varlen 布局
2
   │  prepare_cp_mirror（rank 相关）
3
   ▼
4
本 rank mirror 布局（cp_mirror_len 行）
5
   │  逐层：Q 投影 ∥ 发起 KV all-gather → 等待 → 按逆索引重排 → RoPE → 写缓存
6
   ▼
7
本 rank mirror hidden
8
   │  all-gather + 按逆索引重排
9
   ▼
10
原始 varlen 布局（total_q 行）→ 采样

三个核心设计：

1. 镜像切分让各 rank 的 query 负载对齐，而 mirror 布局让每个 rank 的 token 在张量里保持连续——一次索引预处理内核同时产出正向与逆向映射，以及压缩 KV 的 ks/ke。
2. query 对齐的字段在 mirror 布局，key 对齐的字段在原始布局。这解释了为什么 positions 要重排而 positions_c 不重排，为什么压缩要在完整序列坐标系里做、为什么 positions_c 可以不重排，以及为什么 attention 内核只拿 kv_index 就能对两种布局的 query 与 key 正常工作。
3. 通信全部走异步接口：_make_cp_gather_getter 返回闭包，「发起即返回、等待即取值」，把 all-gather 与 Q 投影、SWA 投影重叠；session 负责缓冲区的分配与生命周期。

代价是每层多次集合通信，与 TP 的通信模式不相上下。官方因此把它保留为「单卡放得下、prefill 太慢」时的预备路径，生产上交给 Prefill Load Balancer 和 Wide-EP。对想理解 CP 实现的人来说，这套代码是一个完整样本：索引预处理、布局切换、异步通信、与截断/投机解码的特性交互，每一处都有对应的断言和测试。

参考资料

1. DeepInfer GitHub 仓库
2. DeepEP：面向 MoE 与专家并行的通信库
3. Ring Attention with Blockwise Transformers for Near-Infinite Context
4. Striped Attention: Faster Ring Attention for Causal Transformers
5. DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models