Context Parallel

在 Megatron 的术语中，SP 主要让 TP 之间的部分操作按序列分片保存和计算，例如 LayerNorm 与 Dropout。进入需要完整输入的 TP 计算前，布局会再次转换。

CP 则把一条序列分到不同 rank，并让 attention 处理跨分片的依赖。每个 rank 可以负责自己的 token 的其他逐 token 计算。阅读资料时还要看其定义：Ulysses 也使用 Sequence Parallelism 这个名称，但处理的是跨完整 attention 的序列并行。

每个 KV 分块都能得到局部 attention 输出，以及该块的 LogSumExp。合并时，用这些归一化信息重新加权局部输出。直接把局部输出取平均会丢失不同分块的权重。

对同一条 query，设第 $j$ 块的输出为 $O_j$，LogSumExp 为 $\ell_j$。下式用最大值 $a$ 稳定指数计算。因果掩码下，只合并这条 query 实际可见的分块。

Ring 通常保留本地 Q，让 KV 分块依次流经各 rank。它可以逐块计算，并尝试用计算覆盖下一块的传输。每块过小时，通信启动和同步会更显著。

Ulysses 通过 All-to-All 把序列分片换成头分片。每张卡在本地处理所分配头的完整序列，再交换回原布局。可用并行度与头的切分方式有关，GQA、MQA 的 KV 头布局尤其需要检查。

USP 和 LoongTrain 进一步组合头维与上下文维的切分。选择时需要一起看序列长度、头数和互连，不能只比较通信算子的名字。

设总序列长 $S$、batch 为 $B$、KV 头数为 $h_{kv}$、头维为 $d$，每元素 $b$ 字节，共 $p$ 个 CP rank。KV 均匀分片，每个 rank 初始持有 $2BSh_{kv}db/p$ 字节。

在传完其余 $p-1$ 块的基本环形实现中，每个 rank 的发送量为下面的式子，接收量相同。这只计算一层 attention 前向的 KV 载荷。反向还会涉及梯度，因果调度、KV 复制和其他切分也会改变实际计数。

例如 $B=1,S=32768,h_{kv}=2,d=128,p=8$，BF16 下每卡初始 KV 为 4 MiB，前向发送 28 MiB。

序列前部的 query 只能看到较少的历史 token，后部的 query 能看到更多。连续地均分序列，虽然均分了 token 数，却没有均分可见的 query-key 配对数。

Zigzag 一类布局会把前部与后部的片段配在一起，让不同 rank 的计算量更接近。实现还需要同步调整位置、掩码及输出还原。

新 query 很少，但历史 KV 可以很长。CP 可以分片保存历史 KV，让多个 rank 各自计算对当前 query 的局部贡献，再按 Softmax 的归一化信息合并。

因此长上下文推理会考虑传 Q、传 KV 或传局部输出。Pass-Q 与 Pass-KV 的取舍还受到新 token 数、已有 KV 长度和头数影响，训练时的通信路径不能直接照搬。

如果 CP rank 保存相同的模型权重，每个 rank 得到的是这些相同参数在不同 token 上的梯度贡献。它们仍需合并，才能对应整条序列的损失。

KV 的梯度也要回到它所属的序列分片。检查实现时，需要同时追踪 attention 内的数据交换，以及模型参数梯度所在的归约组。

先固定模型、序列长度和 batch，再比较各组合的 attention、MLP、通信与显存峰值。GQA 的 KV 头数和跨节点带宽分别会怎样改变交点？