上下文并行（CP）¶

TL;DR：PCP 通过序列切分加速预填充。DCP 消除 KV 缓存冗余。

ContextParallel

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

什么是 CP？¶

上下文并行（Context Parallel，CP） 是一种在多个设备上沿序列维度进行并行计算的策略。

预填充上下文并行（PCP） expands the world size of devices and uses dedicated communication domains. Its primary goal is to partition the sequence dimension during the prefill phase, enabling different devices to compute distinct chunks of the sequence simultaneously. The KV cache is sharded along the sequence dimension across devices. This approach impacts the computational logic of both the Prefill and Decode stages to varying degrees.

解码上下文并行（DCP） reuses the communication domain of Tensor Parallelism (TP) and does not require additional devices. Its main objective is to eliminate duplicated storage of the KV cache by sharding it along the sequence dimension across devices within the TP domain that would otherwise hold redundant copies. DCP primarily influences the Decode logic, as well as the logic for chunked prefill and cached prefill.

如何使用 CP？¶

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

工作原理是什么？¶

设备分布¶

We introduce new communication domains for PCP and reuse TP for DCP, and this is the new layout of devices for PCP2, DCP2, and TP4. device_world

块表¶

CP 对 KV 缓存存储进行序列分片。为了实现高效的存储和访问，token 以交错方式存储到各个设备上，交错粒度由 cp_kv_cache_interleave_size 决定，默认值为 cp_kv_cache_interleave_size=1，即“token 交错”。

鉴于 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。

For any token x, referencing the following figure, its (virtual) block index is x // virtual_block_size, and the offset within the virtual block is offset_within_virtual_block = x % virtual_block_size. The local block index is local_block_index = offset_within_virtual_block // cp_kv_cache_interleave_size, and the device number is target_rank = local_block_index % cp_size. The offset within the local block is (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 缓存沿序列维度分片存储。其影响体现在解码和分块预填充阶段的逻辑中。

预填充阶段：
As illustrated, during the Chunked Prefill computation, two distinct logic implementations are employed for MLA and GQA backends.

In the MLA backend, a Context KV Cache all_gather operation is performed to aggregate the full KV values. These are then used for attention computation with the Q values of the current chunk. Note that in multi-request scenarios, the directly gathered KV results are interleaved across requests. The reorg_kvcache function is used to reorganize the KV cache, ensuring that the KV cache of the same request is stored contiguously.
In the GQA backend, an all_gather is performed along the head dimension for Q. This is because DCP overlaps with the TP communication domain, and the Q heads within a DCP group differ. However, they need to exchange results with the locally computed KV cache for online Softmax updates. To ensure correctness during result updates, the Q values are synchronized across the DCP group via head-dimension all_gather. During the result update process, cp_lse_ag_out_rs is invoked to aggregate attn_output and attn_lse, update the results, and perform a reduce-scatter operation on the outputs. Alternatively, we can use an all-to-all communication to exchange the output and LSE results, followed by direct local updates. This approach aligns with the logic adapted for PCP compatibility.

DCP-Prefill

解码阶段： The logic during the decode phase is consistent with that of GQA's chunked prefill: an all-gather operation is first performed along the Q head dimension to ensure consistency within the DCP group. After computing the results with the local KV cache, the results are updated via the cp_lse_ag_out_rs function.

DCP-Decode

预填充上下文并行（PCP）¶

首尾风格的 Token 切分

PCP requires splitting the input sequence and ensuring balanced computational load across devices during the prefill phase. We employ a head-tail style for splitting and concatenation: specifically, the sequence is first padded to a length of 2*pcp_size, then divided into 2*pcp_size equal parts. The first part is merged with the last part, the second part with the second last part, and so on, thereby assigning computationally balanced chunks to each device. Additionally, since allgather aggregation of KV or Q results in interleaved chunks from different requests, we compute pcp_allgather_restore_idx to quickly restore the original order.

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

PCP-Partition

预填充阶段：

During the Prefill phase (excluding chunked prefill), we employ an all-gather KV approach to address the issue of incomplete sequences on individual GPUs. It is important to note that we only aggregate the KV values for the current layer at a time, and these are discarded immediately after use, avoiding excessive peak memory usage. This method can also be directly applied to KV cache storage (since the KV cache partitioning method differs from PCP sequence partitioning, it is inevitable that each GPU requires a complete copy of the KV values). All attention backends maintain consistency in this logic.

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

PCP-Prefill

解码阶段：

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

PCP-Decode

分块预填充：

Currently, there are three viable approaches for Chunked Prefill compatibility: AllGatherQ, AllGatherKV, and Ring-Attn. Since PCP performs sequence sharding on both the query sequence and the KV cache, we need to ensure that one side has complete information or employ a method like Ring-Attn to perform computations sequentially. The advantages and disadvantages of Ring-Attn will not be elaborated here.

We have implemented the AllGatherQ approach in the GQA attention backend and the AllGatherKV approach in the MLA attention backend. The workflow after AllGatherQ is identical to the decode phase, while the workflow after AllGatherKV is the same as the standard prefill phase. For details, please refer to the diagram below; specific steps will not be repeated.

One important note: AllGatherKV may lead to significant peak memory usage when the context length becomes excessively long. To mitigate this, we adopt a segmented processing strategy. By predefining the maximum amount of KV cache processed per round, we sequentially complete the attention computation and online softmax updates for each segment.

PCP-ChunkedPrefill

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

SFA DSA-CP mixed execution intentionally reuses the normal TP-sharded o_proj. This is part of the DSA-CP mixed data path, not a standalone user-facing o_proj TP switch. The mixed path is used when one instance may handle both decode-only and prefill/mixed batches, so o_proj must support two layouts at runtime:

Decode-only batches keep the decode TP path. SFA outputs are exchanged with an all-to-all in the TP group, then the original TP-sharded o_proj runs normally.
Prefill or mixed batches produce SFA outputs that are not directly compatible with the TP-sharded o_proj input layout. Before o_proj forward, each rank all-gathers the TP-sharded o_proj weight and all input-sharded quantization parameters into temporary full-weight buffers. The full-weight o_proj forward runs once for that batch, and the module is then restored to the TP parameter aliases.

The storage invariant is that the original TP-sharded o_proj parameter remains the only persistent source of truth. o_proj_tp_* tensors are aliases of the original parameter storage. o_proj_full_* tensors are reusable communication buffers for prefill/mixed full-gather execution only. They must not become a second persistent copy of the TP weight.

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