第 8 课：上下文并行 CP

• 对应官方章节：Context Parallelism / Ring Attention / Zig-Zag Ring Attention
• 前置：第 3 课（all-gather / all-to-all 概念）、第 7 课（SP）
• 状态：未开始

本课目标

完成后你需要能够：

• 说明 CP 与 SP 的区别：CP 沿序列切分整个模型（含 attention），独立于 TP 使用。
• 推演 Ring Attention 的三步循环，说明为什么注意力是唯一需要跨卡通信的模块。
• 解释朴素 Ring Attention 在因果掩码下的负载不均，以及 Zig-Zag 如何均衡。
• 对比 all-gather 与 all-to-all（ring）两种 K/V 交换实现的取舍。

核心概念

1. 上下文并行（CP）的提出背景与第一性原理

随着大语言模型应用从短文本对话演进到长文档分析、代码库解析、超长推理（如 128k、256k 乃至 1M 上下文），系统显存面临着序列长度 $s$ 引发的非线性膨胀危机：

• 自注意力机制二次方显存：原生 Attention 的分数矩阵维度为 $(s, s)$，显存复杂度为 $O(s^2)$；即便采用 FlashAttention 避免物化分数矩阵，前向与反向中的输入/输出激活、残差连接以及为反向保存的 K/V 激活仍然随 $s$ 呈线性增长。在 128k 序列下，仅单样本激活显存就高达数百吉字节。
• TP 扩展性的物理壁障：第 6 课已指出，张量并行（TP）受限于机内 NVLink 超高带宽，通常最大只能取 $\text{TP}=8$。仅靠机内 TP 无法跨机扩展以消化百万级别的超长序列。

上下文并行（Context Parallelism, CP）应运而生。如图 8-1 与图 8-2 所示，CP 是一种全局序列分块并行架构：

Ring Attention 与因果掩码的负载均衡

环形 KV 轮转

• 序列切成 $P$ 块，每卡的 Q 块常驻
• K/V 块沿环传递 $P-1$ 次
• 每步用 online softmax 合并 $(m, d, O)$，结果与完整注意力相同

Zig-Zag 切分

• 朴素切分下 rank $i$ 要算 $i+1$ 块，负载不均
• 切成 $2P$ 块，rank $i$ 拿第 $i$ 块和第 $2P-1-i$ 块
• 每卡有效块数都是 $2P + 1$

每步传 K、V 共 $2 \cdot \frac{S}{P} \cdot d$ 个元素，计算约 $4 (\frac{S}{P})^2 d$ FLOPs：计算随块长平方增长、通信线性增长，块足够长时通信可被计算掩盖

图 8-1：Ring Attention 的 KV 轮转与 Zig-Zag 负载均衡

上下文并行长序列切分

超长序列 S (如 128k Tokens)

分块 Chunk 0 -> GPU 0

分块 Chunk 1 -> GPU 1

分块 Chunk 2 -> GPU 2

分块 Chunk 3 -> GPU 3

图 8-2：上下文并行按序列切块并分配至独立 GPU 拓扑

• 它将长达 $s$ 的输入序列在最外层直接沿序列维度均匀切分为 $N_{\text{cp}}$ 个子序列片段，分配给集群中的 $N_{\text{cp}}$ 张 GPU（通常跨物理节点，结合 RDMA 网络）。
• 每张 GPU 分配到的局部序列长度仅为 $s_{\text{local}} = \frac{s}{N_{\text{cp}}}$。
• 与逐 Token 算子的天然正交性：Transformer 中的 MLP（Feed-Forward）、LayerNorm/RMSNorm、Embedding 以及各类残差连接完全是逐 Token 独立计算（Token-wise）的。在 CP 模式下，这些模块在各卡上进行完全本地的张量运算，全程产生 0 跨卡通信开销。
• 自注意力的唯一性挑战：整个网络中，唯独多头自注意力（Self-Attention）模块打破了 Token 间的独立性——根据 Attention 公式 $\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q K^T}{\sqrt{d}}\right) V$，每个 Token 的 Query 必须与全序列所有历史 Token 的 Key/Value 进行点积计算。因此，如何在各 GPU 仅持有局部切片的前提下，低显存、高带宽地计算出全量注意力，是 CP 的核心难题。

2. Ring Attention 的环形流转与异步重叠机制

为了解决注意力跨卡计算的问题，UC Berkeley 提出了经典的 Ring Attention（环形注意力） 算法。其核心思想是：让局部 Query 始终固定在本地显存，而让 Key/Value 块沿着逻辑通信环异步流转轮巡。

如图 8-3 所示，整个通信与计算流水线由双缓冲（Double Buffering）机制驱动：

GPU 0 (Q_0 固定)

GPU 1 (Q_1 固定)

环形 P2P 异步通信 (KV 轮转)

Step 0: 计算本地 Q_i 与 K_i, V_i 的注意力

1发送 K_0, V_0 -> GPU 1，接收来自前驱的 KV

2发送 K_1, V_1 -> GPU 2，接收来自前驱的 KV

Step 1: 异步接收下一块 KV，同时计算当前重叠块注意力

利用 Online Softmax 动态更新分母与局部输出，无需中间显存

图 8-3：Ring-Attention 环形 KV 轮转与双缓冲通信计算重叠流程

1. 初始状态：每个 Rank $r \in [0, N_{\text{cp}}-1]$ 均持有本地的子序列块 $Q_r, K_r, V_r$，形状均为 $(b, s_{\text{local}}, h)$。

2. $N_{\text{cp}}$ 步环形流水循环：在每一个时间步 $t \in [0, N_{\text{cp}}-1]$：

   • 异步通信启动：通信流启动非阻塞 P2P 传输（dist.isend 与 dist.irecv），将当前卡持有的 $K, V$ 块异步投递给环形拓扑的下一个邻居 $(r+1) \pmod{N_{\text{cp}}}$，同时异步接收来自上一个邻居 $(r-1) \pmod{N_{\text{cp}}}$ 的新 $K, V$ 块。
   • 计算与通信重叠：在后台网络传输的同时，GPU 计算流在本地执行矩阵乘：将固定的本地 $Q_r$ 与当前持有的 $K, V$ 计算局部注意力，并通过 Online Softmax 算法就地合并、缩放并累加到本地的上下文输出张量 $O_r$ 中。
   • 同步与双缓冲轮换：计算完毕后，等待后台异步 P2P 传输收包完成，交换当前缓冲区与接收缓冲区（Buffer Swap），进入下一轮迭代。

3. 显存恒定特性：在整个环形轮转过程中，每张卡只需要额外保留用于接收下一个块的临时 KV 双缓冲空间。每卡显存仅为 $O(s_{\text{local}}) = O(\frac{s}{N_{\text{cp}}})$，与全网序列总长度 $s$ 的膨胀完全解耦。

3. 与 FlashAttention 的理论同源性：Online Softmax

为什么可以将原本需要全量点积的 Softmax 拆分成小块分别计算？其数学基石正是 Online Softmax（亦是 FlashAttention 的核心，第 13 课将详解其 CUDA 实现）。

在朴素计算中，Softmax 的分母需要全序列求和：$\ell = \sum_{j=1}^s \exp(S_j - m)$。若简单地对每个块独立做 Softmax 再拼接，会导致各块的分母归一化基准不一致，计算结果必然完全错误。

Ring Attention 在本地维护三个统计状态量：

• $m^{(t)}$：当前已见所有 KV 块的最大点积分数（用于防止指数溢出）；
• $\ell^{(t)}$：当前已见所有 KV 块的归一化指数和；
• $O^{(t)}$：当前累积的局部加权输出向量。

当第 $t$ 个新的 KV 块到达，产生局部点积分数 $S_{\text{new}}$ 时： $$\tilde{m} = \max(m^{(t-1)}, \max(S_{\text{new}}))$$ $$\ell^{(t)} = \ell^{(t-1)} \cdot \exp(m^{(t-1)} - \tilde{m}) + \sum \exp(S_{\text{new}} - \tilde{m})$$ $$O^{(t)} = O^{(t-1)} \cdot \frac{\exp(m^{(t-1)} - \tilde{m}) \cdot \ell^{(t-1)}}{\ell^{(t)}} + \frac{\exp(S_{\text{new}} - \tilde{m})}{\ell^{(t)}} \cdot V_{\text{new}}$$ 这保证了在 $N_{\text{cp}}$ 轮迭代结束时，本地输出 $O$ 与全量非分块注意力算出的结果在数值上完全等价。

4. 因果掩码下的严重负载失衡与 Zig-Zag 拓扑平衡

在自回归因果语言模型（Autoregressive / Causal LLM）中，注意力矩阵呈下三角掩码结构：位置为 $i$ 的 Query 只能 Attend 历史位置 $j \le i$ 的 Key/Value，未来的 Token 全被掩码置零（Masked Out）。

A. 朴素连续切分的负载失衡

若将长序列简单地连续切片（例如 4 卡场景，Rank 0 分得 $[0, \frac{s}{4}]$，Rank 3 分得 $[\frac{3s}{4}, s]$）：

• Rank 0：持有的 Query 全部位于序列最前端。它只需要与本地的 KV 块计算，其余 3 个远端发来的 KV 块全部处于未来时间步，计算量全部为 0。Rank 0 在第 1 步后就陷入空转饥饿。
• Rank 3：持有的 Query 全部位于序列末尾。它必须与集群中所有 4 个 KV 块全部进行点积运算，工作量高达 Rank 0 的 4 倍。
• 气泡率（Bubble Rate）：在下三角注意力矩阵下，纯连续切分的总有效计算量仅为方阵的 $50\%$，但由于木桶效应，每轮必须等待最慢的 Rank 3，系统整体算力浪费高达近 $50\%$。

B. Zig-Zag Ring Attention 的数学自洽解

为了消除这种空转，Zig-Zag Ring Attention 改变了分块映射策略：让每张 GPU 同时持有一块“序列开头的早 Token”和一块“序列末尾的晚 Token”

设每个 Rank 分配 $K=2$ 个 Chunk，全序列共 $2 N_{\text{cp}}$ 个块。Zig-Zag 按照如下规律交错分配： $$\text{Rank } r \text{ 持有 Chunk } \{ r, \ 2 N_{\text{cp}} - 1 - r \}$$ 以 $N_{\text{cp}} = 4$ 为例：

• Rank 0: Chunk $\{0, 7\}$
• Rank 1: Chunk $\{1, 6\}$
• Rank 2: Chunk $\{2, 5\}$
• Rank 3: Chunk $\{3, 4\}$

因果掩码下，第 $q$ 个 Query Chunk 需要计算的有效 KV 块数量恰为 $q + 1$。因此，Rank $r$ 的总计算工作量为： $$W_r = (r + 1) + (2 N_{\text{cp}} - 1 - r + 1) = (r + 1) + (2 N_{\text{cp}} - r) = 2 N_{\text{cp}} + 1$$ 变量 $r$ 被完全抵消 对于任意 Rank $r$，其总工作量恒等于常数 $2 N_{\text{cp}} + 1$（例如 4 卡下每张卡分得 $2 \times 4 + 1 = 9$ 个块的计算任务），因果掩码导致的等待被消除。

5. 两种工程通信实现：All-to-All / Ring vs All-Gather

在主流大模型框架中，上下文并行存在两种流派：

特性维度	Ring Attention (逐块 P2P 轮转)	DeepSpeed Ulysses (All-to-All 转置)
通信拓扑	环形单向 P2P (isend/irecv)	全交换 All-to-All (类似数据转置)
显存占用	$O(\frac{s}{N_{\text{cp}}})$，另加一块 KV 的接收缓冲	注意力阶段每卡持有全序列、$\frac{H}{N_{\text{cp}}}$ 个头的 K/V，总大小 $s \cdot \frac{h}{N_{\text{cp}}}$ 与 Ring 的 $\frac{s}{N_{\text{cp}}} \cdot h$ 相同，只是切分维度从序列换成了头；另需 All-to-All 收发缓冲
通信-计算重叠	极佳：P2P 传输完全被单块计算掩盖	难以重叠：进入注意力前必须等待 All-to-All 同步完成
对注意力头数依赖	完全无约束，头数少于 $N_{\text{cp}}$ 亦可运行	强约束：注意力头数必须能被 $N_{\text{cp}}$ 整除
长文本极限	适合百万级（1M+）超长序列	适合中长（32k~64k）且节点间全互联场景

具体演示与数值推演

设集群规模为 $N_{\text{cp}} = 4$ 张 GPU，全序列划分为 8 个 Chunk（编号 0 至 7，因果掩码下 Chunk $q$ 合法 Attend 的 KV 范围是 $[0, q]$）：

• 朴素连续分配方案（Sequential Assignment）：
  • 各卡持有的 Chunk：Rank 0: $\{0, 1\}$，Rank 1: $\{2, 3\}$，Rank 2: $\{4, 5\}$，Rank 3: $\{6, 7\}$。
  • 各卡有效计算对数：
  • Rank 0: $(0+1) + (1+1) = 1 + 2 = 3$ 对
  • Rank 1: $(2+1) + (3+1) = 3 + 4 = 7$ 对
  • Rank 2: $(4+1) + (5+1) = 5 + 6 = 11$ 对
  • Rank 3: $(6+1) + (7+1) = 7 + 8 = 15$ 对
  • 失衡分析：最大负载（15）是最小负载（3）的 5.0 倍。Rank 0 大部分时间处于等待阻塞状态。
• Zig-Zag 镜像分配方案（Zig-Zag Assignment）：
  • 各卡持有的 Chunk：Rank 0: $\{0, 7\}$，Rank 1: $\{1, 6\}$，Rank 2: $\{2, 5\}$，Rank 3: $\{3, 4\}$。
  • 各卡有效计算对数：
  • Rank 0: $(0+1) + (7+1) = 1 + 8 = 9$ 对
  • Rank 1: $(1+1) + (6+1) = 2 + 7 = 9$ 对
  • Rank 2: $(2+1) + (5+1) = 3 + 6 = 9$ 对
  • Rank 3: $(3+1) + (4+1) = 4 + 5 = 9$ 对
  • 均衡分析：所有卡的工作量都等于 9 对，负载完全均衡。虽然每张卡都需要与环上其他卡发生数据交换，但单卡发送与接收的总通信量恒定为 $2 \times \frac{s}{N_{\text{cp}}}$，没有产生任何额外的数据膨胀。

代码填空题

模拟 chunk 分配与 Ring Attention 轮转，验证 Zig-Zag 的负载均衡。

def sequential_assignment(num_ranks: int, chunks_per_rank: int):
    """
    朴素连续分配：rank r 拿连续的 chunks_per_rank 个 chunk。
    返回 dict: rank -> chunk 序号列表（chunk 从 0 开始编号）。
    """
    assign = {}
    for r in range(num_ranks):
        assign[r] = ______   # 填空：第 r 个 rank 拿到的 chunk 列表
    return assign

def zigzag_assignment(num_ranks: int, chunks_per_rank: int):
    """
    Zig-Zag：第 r 个 rank 交替从序列两头取。
    序号：rank r 第 k 个 chunk 的全局编号。
    规则：编号 < num_ranks 的 chunk 正序发（chunk r 给 rank r）；
          超过后反向交错：chunk (2*num_ranks - 1 - r) 给 rank r。
    对 chunks_per_rank = 2（最常见的 2 chunk/卡）：rank r 拿 {r, 2p-1-r}。
    """
    total = num_ranks * chunks_per_rank
    assign = {}
    for r in range(num_ranks):
        chunks = []
        for k in range(chunks_per_rank):
            # 填空：交替取数公式（奇数块从尾部取）
            if k % 2 == 0:
                chunks.append(______)      # 正序：第 k/2 轮正序编号
            else:
                chunks.append(______)      # 反序：从尾端回推
        assign[r] = chunks
    return assign

def causal_workload(assign: dict[int, list[int]], num_chunks: int) -> dict[int, int]:
    """
    因果掩码下每个 rank 需要计算的 (q_chunk, kv_chunk) 对数。
    一个 rank 有多个 query chunk 时，工作量 = Σ_q (q 号 + 1)，
    因为 q 可以 attend kv_chunk 0..q（因果）。
    """
    workload = {}
    for r, q_chunks in assign.items():
        total = 0
        for q in q_chunks:
            total += ______                # 填空：该 query chunk 的合法 kv 数量
        workload[r] = total
    return workload

def ring_rotation(num_ranks: int, rank: int, start_chunk: int) -> list[int]:
    """
    Ring Attention 轮转：返回该 rank 在每一步（0..num_ranks-1）持有的
    K/V chunk 编号（K/V 每步从 rank r-1 传到 rank r）。
    第 t 步：本 rank 持有的是 t 步前从上游 rank (r - t) 出发的 chunk，
    即编号 (start_chunk - t) mod num_ranks。
    """
    steps = []
    for t in range(num_ranks):
        steps.append(______)               # 填空：(start - t) 沿环回绕
    return steps

if __name__ == "__main__":
    p = 4
    seq_assign = sequential_assignment(p, 2)
    zz_assign = zigzag_assignment(p, 2)
    print("sequential:", seq_assign)
    print("zig-zag   :", zz_assign)

    w_seq = causal_workload(seq_assign, p * 2)
    w_zz = causal_workload(zz_assign, p * 2)
    print("sequential workload:", w_seq, " max/min =", max(w_seq.values()) / min(w_seq.values()))
    print("zig-zag workload   :", w_zz, " max/min =", max(w_zz.values()) / min(w_zz.values()))
    assert len(set(w_zz.values())) == 1, "zig-zag 应完全均衡"
    print("Zig-Zag 负载完全均衡 ✓")

    # Ring 轮转示例：rank 2 的 K/V 持有序列
    print("rank 2 轮转:", ring_rotation(p, 2, 2))

三个问答题

Q1

CP 中为什么只有 attention 模块需要跨卡通信？为什么交换的是 K/V 而不是 Q？如果把整条序列的 K/V 一次性复制到每卡（all-gather 式），与 ring 式相比显存和延迟各付出什么代价？

Q2

用因果掩码解释为什么朴素连续切分下”GPU 1 立即算完、GPU 4 要等三轮”，并说明 Zig-Zag 为什么能均衡负载。Zig-Zag 下每个 rank 需要从其他所有 rank 获得 K/V，这是否增加总通信量？

Q3

教材指出 Ring Attention 与 FlashAttention 都依赖 online softmax。请解释：分块计算注意力时，为什么不能简单地对每个块分别 softmax 再拼接？还需要维护哪些额外统计量？（可以只给概念，第 13 课会用代码实现。）

检查与通过标准

总分 10 分：代码正确 4 分（分配公式、工作量统计、轮转公式、均衡验证）、三题各 2 分、通过线 8 分。

一票否决项：

• 分不清 CP 与 SP（本书定义）。
• 认为 CP 中 MLP/LayerNorm 也需要跨卡交换激活。
• 说不出因果掩码导致负载不均的具体机制。
• 认为”分别 softmax 再拼接”是正确做法。

Tip💡 点击展开参考答案与详细代码实现

参考答案（仅 answer 分支）

先完成题目再核对；核对后必须解释关键公式，并改一个规模重新计算。

def sequential_assignment(num_ranks: int, chunks_per_rank: int):
    """
    朴素连续分配：rank r 拿连续的 chunks_per_rank 个 chunk。
    返回 dict: rank -> chunk 序号列表（chunk 从 0 开始编号）。
    """
    assign = {}
    for r in range(num_ranks):
        assign[r] = list(range(r * chunks_per_rank, (r + 1) * chunks_per_rank))   # 填空：第 r 个 rank 拿到的 chunk 列表
    return assign

def zigzag_assignment(num_ranks: int, chunks_per_rank: int):
    """
    Zig-Zag：第 r 个 rank 交替从序列两头取。
    序号：rank r 第 k 个 chunk 的全局编号。
    规则：编号 < num_ranks 的 chunk 正序发（chunk r 给 rank r）；
          超过后反向交错：chunk (2*num_ranks - 1 - r) 给 rank r。
    对 chunks_per_rank = 2（最常见的 2 chunk/卡）：rank r 拿 {r, 2p-1-r}。
    """
    total = num_ranks * chunks_per_rank
    assign = {}
    for r in range(num_ranks):
        chunks = []
        for k in range(chunks_per_rank):
            # 填空：交替取数公式（奇数块从尾部取）
            if k % 2 == 0:
                chunks.append((k // 2) * num_ranks + r)      # 正序：第 k/2 轮正序编号
            else:
                chunks.append(total - 1 - (k // 2) * num_ranks - r)      # 反序：从尾端回推
        assign[r] = chunks
    return assign

def causal_workload(assign: dict[int, list[int]], num_chunks: int) -> dict[int, int]:
    """
    因果掩码下每个 rank 需要计算的 (q_chunk, kv_chunk) 对数。
    一个 rank 有多个 query chunk 时，工作量 = Σ_q (q 号 + 1)，
    因为 q 可以 attend kv_chunk 0..q（因果）。
    """
    workload = {}
    for r, q_chunks in assign.items():
        total = 0
        for q in q_chunks:
            total += q + 1                # 填空：该 query chunk 的合法 kv 数量
        workload[r] = total
    return workload

def ring_rotation(num_ranks: int, rank: int, start_chunk: int) -> list[int]:
    """
    Ring Attention 轮转：返回该 rank 在每一步（0..num_ranks-1）持有的
    K/V chunk 编号（K/V 每步从 rank r-1 传到 rank r）。
    第 t 步：本 rank 持有的是 t 步前从上游 rank (r - t) 出发的 chunk，
    即编号 (start_chunk - t) mod num_ranks。
    """
    steps = []
    for t in range(num_ranks):
        steps.append((start_chunk - t) % num_ranks)               # 填空：(start - t) 沿环回绕
    return steps

if __name__ == "__main__":
    p = 4
    seq_assign = sequential_assignment(p, 2)
    zz_assign = zigzag_assignment(p, 2)
    print("sequential:", seq_assign)
    print("zig-zag   :", zz_assign)

    w_seq = causal_workload(seq_assign, p * 2)
    w_zz = causal_workload(zz_assign, p * 2)
    print("sequential workload:", w_seq, " max/min =", max(w_seq.values()) / min(w_seq.values()))
    print("zig-zag workload   :", w_zz, " max/min =", max(w_zz.values()) / min(w_zz.values()))
    assert len(set(w_zz.values())) == 1, "zig-zag 应完全均衡"
    print("Zig-Zag 负载完全均衡 ✓")

    # Ring 轮转示例：rank 2 的 K/V 持有序列
    rot = ring_rotation(p, 2, 2)
    print("rank 2 轮转:", rot)
    assert rot == [2, 1, 0, 3], "K/V 从 r-1 流向 r，第 t 步应持有 chunk (r - t) mod N"

补全后的代码（关键填空处）

def sequential_assignment(num_ranks, chunks_per_rank):
    assign[r] = list(range(r * chunks_per_rank, (r + 1) * chunks_per_rank))

def zigzag_assignment(num_ranks, chunks_per_rank):
    total = num_ranks * chunks_per_rank
    for k in range(chunks_per_rank):
        if k % 2 == 0:
            chunks.append((k // 2) * num_ranks + r)          # 正序轮
        else:
            chunks.append(total - 1 - (k // 2) * num_ranks - r)  # 反序轮

def causal_workload(assign, num_chunks):
    total += q + 1        # 因果：query chunk q 可 attend kv chunk 0..q，共 q+1 个

def ring_rotation(num_ranks, rank, start_chunk):
    steps.append((start_chunk - t) % num_ranks)   # K/V 从 r-1 流向 r，第 t 步持有 chunk (r-t) mod N

运行结果（p=4，每卡 2 chunk）：

sequential: {0:[0,1], 1:[2,3], 2:[4,5], 3:[6,7]}   工作量 {3,7,11,15}，失衡 5 倍
zig-zag   : {0:[0,7], 1:[1,6], 2:[2,5], 3:[3,4]}   工作量 {9,9,9,9}，完全均衡 ✓
rank 2 轮转: [2, 1, 0, 3]

三个问答题要点

Q1

• 只有 attention 需要全序列信息（token 要 attend 所有 K/V）；MLP/LN 逐 token 独立，序列切分免费；
• 交换 K/V 而非 Q：每个 token 的 Q 只需自己那份；K/V 是全量共享的”被 attend 对象”；
• all-gather 式：一次通信、简单，但每卡临时显存 = 全部 K/V（O(seq·h) 额外）；ring 式：分块轮转，临时显存只多一块，通信与计算重叠，代价是多步的基延迟。

Q2

• 因果掩码下 softmax 按行（query）计算：rank 持有序列开头的 query 时，合法 K/V 都在本地或很近；持有末尾 query 的 rank 需要所有前序 K/V，且要等其他 rank 算完才能汇总（online softmax 的 block 间依赖）；
• 朴素分配：rank 0 工作量 1 对、rank 3 工作量 4 对 → 等待链；
• Zig-Zag：每卡都有”早 + 晚”混合的 query → 工作量均衡（{9,9,9,9}）；每个 rank 需要来自其他所有 rank 的 K/V（通信”全连接”），但总通信量不变（每块 K/V 仍只传一次/一圈），只是不再能”早退”。

Q3

• 不能分别 softmax 再拼接：softmax 分母是整行的 exp 和，只看到部分块时不知道全局归一化因子；
• 需要维护：running max m（数值稳定 + 已算块的重新缩放）、running sum l（exp(score−m) 之和）；每来一块：m_new = max(m, rowmax)，旧输出按 exp(m−m_new) 缩放，新块按 exp(score−m_new) 累加；
• 这就是 FlashAttention 的 online softmax（第 13 课代码实现）。

通过要点

• CP 沿序列切整个模型（可独立使用）；SP 只切 TP 配套的 LN/dropout 区域。
• Ring Attention 三步：async send K/V → 本地计算 → recv；共 cp 步。

官方主参考

• Ultra-Scale Playbook
• PyTorch distributed documentation

← 上一课
第 7 课：序列并行 SP 下一课 →
第 9 课：流水线并行 PP