上下文并行 (CP)#
摘要：PCP通过序列分割加速预填充阶段。DCP消除KV缓存的冗余存储。
关于开发过程中的主要讨论，请参阅RFC以及被该RFC引用或引用该RFC的相关链接。
什么是CP？#
上下文并行 (CP) 是一种将计算沿着序列维度在多个设备上并行化的策略。
预填充上下文并行 (PCP) 扩展了设备的世界大小并使用专用的通信域。其主要目标是在预填充阶段划分序列维度，使不同的设备能够同时计算序列的不同片段。KV缓存沿着序列维度在设备间分片存储。这种方法在不同程度上影响了预填充和解码两个阶段的计算逻辑。
解码上下文并行 (DCP) 复用张量并行 (TP) 的通信域，不需要额外的设备。其主要目标是通过在TP域内沿着序列维度分片存储KV缓存，消除原本会存储冗余副本的设备上的重复存储。DCP主要影响解码阶段的逻辑，以及分块预填充和缓存预填充的逻辑。
如何使用CP？#
请参考上下文并行用户指南获取详细信息
工作原理#
块表#
CP对KV缓存存储执行序列分片。为了便于高效存储和访问，token以交错的方式跨设备存储，交错粒度由cp_kv_cache_interleave_size决定。
如图所示，块表中定义了一个虚拟块，其中同一CP设备组内的块组成一个虚拟块。虚拟块的大小为virtual_block_size = block_size * pcp_size * dcp_size。
对于任意token x，其（虚拟）块索引为x // virtual_block_size，在虚拟块内的偏移量为x % virtual_block_size。本地块索引为offset_within_virtual_block // cp_kv_cache_interleave_size，设备编号为local_block_index % (pcp_size * dcp_size)。在本地块内的偏移量为(local_block_index // (pcp_size * dcp_size)) * cp_kv_cache_interleave_size + offset_within_virtual_block % cp_kv_cache_interleave_size。
基于上述逻辑，调整了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更新。为了确保结果更新过程中的正确性，Q值通过头维度的all_gather在DCP组内同步。在结果更新过程中，调用cp_lse_ag_out_rs来聚合attn_output和attn_lse，更新结果，并对输出执行归约-分散操作。或者，我们可以使用all-to-all通信来交换输出和LSE结果，然后直接进行本地更新。这种方法与为PCP兼容性而调整的逻辑保持一致。
sgstr "解码阶段： 解码阶段的逻辑与GQA的分块预填充一致：首先沿着Q头维度执行all-gather操作，以确保DCP组内的一致性。使用本地KV缓存计算结果后，通过cp_lse_ag_out_rs函数更新结果。
预填充上下文并行 (PCP)#
sgstr "头尾风格的分词分区
PCP需要在预填充阶段分割输入序列并确保跨设备的计算负载均衡。我们采用头尾风格进行分割和连接：具体来说，序列首先填充到长度为2pcp_size，然后分成2pcp_size等份。第一部分与最后一部分合并，第二部分与倒数第二部分合并，依此类推，从而将计算平衡的块分配给每个设备。此外，由于KV或Q的allgather聚合会导致来自不同请求的交错块，我们计算pcp_allgather_restore_idx以快速恢复原始顺序。
这些逻辑在函数_update_tokens_for_pcp中实现。
预填充阶段：
在预填充阶段（不包括分块预填充），我们采用all-gather KV的方法来解决单个GPU上序列不完整的问题。需要注意的是，我们一次只聚合当前层的KV值，并且在使用后立即丢弃，避免过高的峰值内存使用量。这种方法也可以直接应用于KV缓存存储（由于KV缓存的分区方法与PCP序列分区方法不同，每个GPU需要一份完整的KV值副本是不可避免的）。所有注意力后端在此逻辑上保持一致。
注意：虽然Ring Attention方法也可以促进信息交换，具有更低的峰值内存并支持计算-通信重叠，但我们在评估后优先实施了all-gather KV方法，因为其开发复杂度高且重叠的收益有限。
解码阶段：
在解码阶段，我们只需要在DCP的all-to-all通信交换输出和LSE之后、进行输出更新之前，在PCP组内添加一个allgather操作。
sgstr "分块预填充：
目前，有三种可行的分块预填充兼容性方法：AllGatherQ、AllGatherKV和Ring-Attn。由于PCP对查询序列和KV缓存都执行序列分片，我们需要确保其中一方拥有完整的信息，或者采用Ring-Attn等方法顺序执行计算。此处不详细阐述Ring-Attn的优缺点。
我们已经在GQA注意力后端实现了AllGatherQ方法，在MLA注意力后端实现了AllGatherKV方法。AllGatherQ之后的工作流程与解码阶段相同，而AllGatherKV之后的工作流程与标准预填充阶段相同。详细信息请参考下图；具体步骤不再赘述。
一个重要注意事项：当上下文长度过长时，AllGatherKV可能导致显著的峰值内存使用量。为了缓解这个问题，我们采用了分段处理策略。通过预定义每轮处理的最大KV缓存量，我们依次完成每个段的注意力计算和在线softmax更新。
