上下文并行 (CP)¶

简而言之：PCP 通过序列拆分加速预填充，DCP 消除 KV 缓存冗余。

ContextParallel

开发过程中的主要讨论，请参考 RFC 以及该 RFC 引用或引用的相关链接。

什么是 CP？¶

上下文并行 (CP) 是一种跨多个设备沿序列维度并行化计算的策略。

预填充上下文并行 (PCP) 扩展了设备的全局大小，并使用专用的通信域。其主要目标是在预填充阶段划分序列维度，使不同设备能够同时计算序列的不同块。 KV 缓存沿序列维度跨设备分片。这种方法在不同程度上影响了预填充和解码阶段的计算逻辑。

解码上下文并行 (DCP) 复用张量并行 (TP) 的通信域，不需要额外的设备。其主要目标是通过沿序列维度在 TP 域内原本持有冗余副本的设备间分片 KV 缓存，来消除其重复存储。 DCP 主要影响解码逻辑，以及分块预填充和缓存预填充的逻辑。

如何使用 CP？¶

详细信息请参考上下文并行用户指南。

工作原理¶

设备分布¶

我们为 PCP 引入了新的通信域，并为 DCP 复用了 TP，这是 PCP2、DCP2 和 TP4 的新设备布局。设备世界

块表¶

CP 对 KV 缓存存储执行序列分片。为了便于高效存储和访问，令牌以交错方式跨设备存储，交错粒度由 cp_kv_cache_interleave_size 决定，其默认值为 cp_kv_cache_interleave_size=1，也称为“令牌交错”。

鉴于 PCP 和 DCP 在 KV 缓存分片方面行为相似，我们统称它们为 CP。具体来说，cp_size = pcp_size * dcp_size，并且 cp_rank = pcp_rank * dcp_size + dcp_rank。

如图所示，在块表中定义了一个虚拟块，其中同一 CP 设备组内的块构成一个虚拟块。虚拟块大小为 virtual_block_size = block_size * cp_size。

对于任何令牌 x，参考下图，其（虚拟）块索引为 x // virtual_block_size，虚拟块内的偏移量为 offset_within_virtual_block = x % virtual_block_size。本地块索引为 local_block_index = offset_within_virtual_block // cp_kv_cache_interleave_size，设备编号为 target_rank = local_block_index % cp_size。本地块内的偏移量为 (local_block_index // cp_size) * cp_kv_cache_interleave_size + offset_within_virtual_block % cp_kv_cache_interleave_size。

BlockTable

基于上述逻辑，调整了 slot_mapping 计算过程，并修改了每个设备上的 slot_mapping 值，以确保 KV 缓存按预期沿序列维度分片并存储在不同设备上。

当前实现要求 block_size % cp_kv_cache_interleave_size == 0。

解码上下文并行 (DCP)¶

如上所述，DCP 的主要功能是沿序列维度分片 KV 缓存以进行存储。其影响体现在解码和分块预填充阶段的逻辑中。

预填充阶段：
如图所示，在分块预填充计算期间，对 MLA 和 GQA 后端采用了两种不同的逻辑实现。

在 **MLA 后端**中，执行上下文 KV 缓存 all_gather 操作以聚合完整的 KV 值。然后，这些值用于与当前块的 Q 值进行注意力计算。请注意，在多请求场景中，直接收集的 KV 结果在请求之间是交错的。使用 reorg_kvcache 函数重新组织 KV 缓存，确保同一请求的 KV 缓存连续存储。
在 **GQA 后端**中，沿 Q 的头维度执行 all_gather。这是因为 DCP 与 TP 通信域重叠，并且 DCP 组内的 Q 头不同。但是，它们需要与本地计算的 KV 缓存交换结果以进行在线 Softmax 更新。为了确保结果更新期间的正确性，通过头维度 all_gather 在 DCP 组内同步 Q 值。在结果更新过程中，调用 cp_lse_ag_out_rs 来聚合 attn_output 和 attn_lse，更新结果，并对输出执行 reduce-scatter 操作。或者，我们可以使用全到全通信来交换输出和 LSE 结果，然后直接进行本地更新。此方法与为 PCP 兼容性调整的逻辑一致。

DCP-Prefill

解码阶段： 解码阶段的逻辑与 GQA 的分块预填充一致：首先沿 Q 头维度执行 all-gather 操作，以确保 DCP 组内的一致性。使用本地 KV 缓存计算结果后，通过 cp_lse_ag_out_rs 函数更新结果。

DCP-Decode

预填充上下文并行 (PCP)¶

头尾式令牌分区

PCP 要求在预填充阶段对输入序列进行切分，并确保各设备间的计算负载均衡。我们采用首尾拼接的方式进行切分与合并：具体来说，先将序列填充至 2*pcp_size 长度，然后等分为 2*pcp_size 份。将第一部分与最后一部分合并，第二部分与倒数第二部分合并，以此类推，从而为每个设备分配计算量均衡的数据块。此外，由于 KV 或 Q 的 allgather 聚合会导致来自不同请求的数据块交错排列，我们计算 pcp_allgather_restore_idx 以快速恢复原始顺序。

这些逻辑在函数 _update_tokens_for_pcp 中实现。

PCP-Partition

预填充阶段：

在预填充阶段（不包括分块预填充），我们采用 all-gather KV 的方法来解决单个 GPU 上序列不完整的问题。需要注意的是，我们每次仅聚合当前层的 KV 值，使用后立即丢弃，从而避免过高的峰值内存占用。该方法也可直接应用于 KV 缓存存储（由于 KV 缓存的分区方式与 PCP 序列分区不同，每个 GPU 不可避免地需要一份完整的 KV 值副本）。所有注意力后端在此逻辑上保持一致。

注意：虽然 Ring Attention 方法也能以更低的峰值内存实现信息交换并支持计算通信重叠，但经评估，其开发复杂度较高且重叠收益有限，因此我们优先实现了 all-gather KV。

PCP-Prefill

解码阶段：

在解码阶段，我们只需在 DCP all-to-all 通信交换输出和 LSE 之后、进行输出更新之前，在 PCP 组内添加一个 allgather 操作。

PCP-Decode

分块预填充：

目前，有三种可行的方法来实现分块预填充兼容性：AllGatherQ、AllGatherKV 和 Ring-Attn。由于 PCP 对查询序列和 KV 缓存都进行了序列分片，我们需要确保一侧拥有完整信息，或采用 Ring-Attn 等方法顺序执行计算。 Ring-Attn 的优缺点在此不再赘述。

我们在 GQA 注意力后端实现了 AllGatherQ 方法，在 MLA 注意力后端实现了 AllGatherKV 方法。 AllGatherQ 之后的工作流与解码阶段相同，而 AllGatherKV 之后的工作流与标准预填充阶段相同。具体步骤请参考下图，此处不再赘述。

一个重要提示：当上下文长度过长时，AllGatherKV 可能导致显著的峰值内存占用。为缓解此问题，我们采用分段处理策略。通过预定义每轮处理的最大 KV 缓存量，我们依次完成每个分段的注意力计算和在线 softmax 更新。

PCP-ChunkedPrefill

SFA DSA-CP 混合 `o_proj` 路径¶

SFA DSA-CP 混合执行有意复用了普通的 TP 分片 o_proj。这是 DSA-CP 混合数据路径的一部分，而非独立的面向用户的 o_proj TP 开关。当单个实例可能同时处理纯解码批次和预填充/混合批次时，会使用混合路径，因此 o_proj 必须在运行时支持两种布局：

纯解码批次 保持解码 TP 路径。 SFA 输出通过 TP 组内的 all-to-all 进行交换，然后原始的 TP 分片 o_proj 正常运行。
预填充或混合批次 产生的 SFA 输出与 TP 分片 o_proj 的输入布局不直接兼容。在 o_proj 前向传播之前，每个 rank 将 TP 分片的 o_proj 权重和所有输入分片的量化参数 all-gather 到临时的全权重缓冲区中。全权重 o_proj 前向传播对该批次执行一次，然后模块恢复为 TP 参数别名。

存储不变性在于，原始的 TP 分片 o_proj 参数仍然是唯一持久化的真实来源。 o_proj_tp_* 张量是原始参数存储的别名。 o_proj_full_* 张量是仅用于预填充/混合全收集执行的可重用通信缓冲区。它们绝不能成为 TP 权重的第二个持久化副本。

这种耦合方式保留了现有的解码 TP 行为，支持预填充/混合 DSA-CP 批次，并避免了添加一个其状态可能与 DSA-CP 混合执行产生偏差的额外配置路径。