1. 背景：为什么长上下文需要上下文并行

大模型推理的瓶颈正在从「算力」转向「显存 + 带宽」。当上下文长度从 4K 涨到 128K、1M 时，两个问题会同时爆发：

• KV Cache 显存爆炸：以 LLaMA-3-70B 为例，80 层、8 个 KV Head、head_dim=128、fp16 精度，单 token 的 KV Cache 约为 2 × 80 × 8 × 128 × 2 Byte ≈ 320 KB。128K 上下文就是约 40 GB，单卡根本放不下。
• Attention 计算量平方增长：O(L²) 的注意力分数矩阵，即使 FlashAttention 做了分块，单卡在超长序列下也会被显存和算力同时卡住。

常见的并行手段各有短板：

并行方式	切分维度	对长上下文的作用
数据并行 DP	Batch	不减少单条序列的显存
张量并行 TP	Hidden / Head	每卡仍需完整序列的 KV
流水并行 PP	Layer	不减少单层激活与 KV
上下文并行 CP	Sequence	直接切分序列，KV 与激活按卡分摊

上下文并行（Context Parallel, CP）的核心思路很直接：把一条长序列沿 token 维度切成 N 段，每张卡只持有自己那一段的 Q/K/V。但 Attention 是全局的——第 0 段的 Query 需要和所有段的 Key 做点积。于是问题就变成了：如何在 N 张卡之间高效地交换 KV，让每张卡都能算完自己那一段的注意力输出。

Ring Attention 就是这个问题最经典的答案。

2. 核心思想：把 KV 沿着环形传递

2.1 朴素做法与它的代价

最直观的实现是 All-Gather：每张卡把自己的 K/V 广播出去，所有卡都拿到完整的 K/V，然后各自算自己那段 Q 的注意力。

# 朴素 All-Gather 方案（伪代码）
def naive_cp_attention(q_local, k_local, v_local, cp_group):
    # 每张卡都拿到完整的 K/V
    k_full = all_gather(k_local, group=cp_group)  # [L, H, D]
    v_full = all_gather(v_local, group=cp_group)  # [L, H, D]
    # 本地 Q 与全局 KV 做注意力
    out = flash_attention(q_local, k_full, v_full)
    return out

问题在于：All-Gather 的通信量是 O(L × N)，每张卡都要接收 (N-1)/N 的完整 KV，显存峰值也回到单卡水平。序列越长、卡越多，通信越不可接受。

2.2 Ring Attention 的关键洞察

Ring Attention 的洞察是：不需要一次性拿到全部 KV，可以分 N 步、每步只传一段。

把 N 张卡排成一个环，每张卡持有自己那一段的 K/V。计算分 N 轮进行：

1. 第 i 轮，每张卡用当前手里的 K/V 块计算自己 Q 的部分注意力（配合在线 Softmax 累积）；
2. 计算的同时，把 K/V 块异步发送给环上的下一张卡，并从上一张卡接收新的 K/V 块；
3. N 轮之后，每张卡都遍历过全部 K/V 块，得到完整的注意力输出。

这样做的两个好处：

• 通信量从 O(L×N) 降到 O(L)：每张卡总共只发送自己那一段 KV 一次（分 N 次小块发出），总量与序列长度线性相关，与卡数无关。
• 通信与计算重叠：第 i 轮的 KV 传输可以和该轮的注意力计算并行，只要计算时间 ≥ 通信时间，Ring Attention 就近似「零通信开销」。这也是它名字里「Ring」的由来——像接力一样转圈。

2.3 在线 Softmax：分块累积的正确性保证

分块计算注意力的前提是 Online Softmax（FlashAttention 的核心技巧）。对每个 Q 块，维护三个累积量：

• m：当前见过的最大 logit（数值稳定用）
• l：累积的 softmax 分母 Σ exp(x - m)
• acc：累积的加权 Value

当新来一个 KV 块时，用新的最大值 m_new 修正旧结果：

def online_softmax_update(m_old, l_old, acc_old, s_block, v_block):
    """s_block: [Bq, Bk] 当前 KV 块的注意力分数; v_block: [Bk, D]"""
    m_block = s_block.max(dim=-1, keepdim=True).values
    m_new = torch.maximum(m_old, m_block)

    # 用新的最大值重新缩放旧的累积量
    scale_old = torch.exp(m_old - m_new)
    p_block = torch.exp(s_block - m_new)

    l_new = l_old * scale_old + p_block.sum(dim=-1, keepdim=True)
    acc_new = acc_old * scale_old + p_block @ v_block
    return m_new, l_new, acc_new

Ring Attention 每一轮就是调用一次这个更新，N 轮后 acc / l 即为最终输出。数学上与一次性算完整注意力完全等价，只是把「一次大矩阵乘」拆成了 N 次小块乘。

3. Prefill 阶段的通信策略

Prefill 阶段处理的是整条 prompt，序列长度 L 很大，但只有一次前向。此时 CP 的切分方式是：沿序列维把 Q/K/V 都切成 N 段，每卡持有一段。

3.1 通信模式：P2P 环形传递 KV

Prefill 的 Ring Attention 是纯 P2P（点对点）通信，用 send/recv 在环上传递 KV 块，而不是集合通信：

import torch
import torch.distributed as dist

def ring_attention_prefill(q, k, v, cp_rank, cp_size, cp_group):
    """
    q/k/v: [B, H, L_local, D]，每卡持有自己那一段
    返回: [B, H, L_local, D]
    """
    B, H, L_local, D = q.shape
    scale = D ** -0.5

    # 在线 softmax 的累积量
    m = torch.full((B, H, L_local, 1), float("-inf"), device=q.device)
    l = torch.zeros((B, H, L_local, 1), device=q.device)
    acc = torch.zeros((B, H, L_local, D), device=q.device)

    k_cur, v_cur = k, v
    next_rank = (cp_rank + 1) % cp_size
    prev_rank = (cp_rank - 1 + cp_size) % cp_size

    for step in range(cp_size):
        # 1) 异步把当前 KV 块发给下一张卡（非阻塞）
        send_k = dist.isend(k_cur, dst=next_rank, group=cp_group)
        send_v = dist.isend(v_cur, dst=next_rank, group=cp_group)

        # 2) 用当前 KV 块计算部分注意力（与上面的发送重叠）
        s = (q @ k_cur.transpose(-1, -2)) * scale  # [B, H, L_local, L_local]
        m_new = torch.maximum(m, s.max(dim=-1, keepdim=True).values)
        scale_old = torch.exp(m - m_new)
        p = torch.exp(s - m_new)
        l = l * scale_old + p.sum(dim=-1, keepdim=True)
        acc = acc * scale_old + p @ v_cur
        m = m_new

        # 3) 等待发送完成，再从上一张卡接收下一块 KV
        send_k.wait()
        send_v.wait()
        if step < cp_size - 1:
            k_next = torch.empty_like(k_cur)
            v_next = torch.empty_like(v_cur)
            dist.irecv(k_next, src=prev_rank, group=cp_group).wait()
            dist.irecv(v_next, src=prev_rank, group=cp_group).wait()
            k_cur, v_cur = k_next, v_next

    return acc / l

几个工程要点：

• isend 先发、irecv 后收：先发起发送再计算，让通信和计算重叠；计算完再 wait 发送、发起接收，避免死锁。
• 因果掩码（Causal Mask）：如果序列是因果的，环上第 i 步的 KV 块对某些 Q 位置是「未来」的，可以整块跳过。Ring Attention 论文里提到，因果场景下平均能省掉约一半的 KV 传输。
• 负载均衡：因果注意力下，靠前的 Q 段要看的 KV 少、靠后的多。常见做法是把序列切成 2N 段，第 i 张卡拿第 i 段和第 2N-1-i 段，让每张卡的总计算量均衡（类似 Zig-Zag 切分）。

3.2 为什么 Prefill 适合 Ring

Prefill 的特点是计算密集、通信可被完全掩盖：每轮要算 L_local × L_local × D 的矩阵乘，而传输的只是一个 KV 块。在 A100/H100 上，只要 L_local 不太小，计算时间通常远大于 P2P 传输时间，Ring Attention 的通信开销可以压到 5% 以内。

4. Decode 阶段的通信策略

Decode 阶段完全不同：每步只生成 1 个 token，Q 的长度是 1，但 KV Cache 会随着生成长度不断增长。

4.1 问题：Decode 下 Ring 不再划算

如果 Decode 也照搬 Ring 的逐块传递，会非常低效：

• 每步只有 1 个 Query，计算量极小（1 × L × D），通信时间远大于计算时间，无法重叠；
• 要遍历 N 个 KV 块，就要 N 次 P2P 往返，延迟被放大 N 倍；
• 而 Decode 是逐 token 串行的，每一步的延迟直接决定吞吐（TPOT）。

所以 Decode 阶段的主流策略是：放弃环形逐块传递，改用 All-Gather 一次性聚合 KV。

4.2 通信模式：All-Gather KV Cache

def decode_attention_cp(q, k_local, v_local, cp_group):
    """
    Decode 阶段：q 只有 1 个 token
    q: [B, H, 1, D]
    k_local/v_local: [B, H, L_local, D] 本卡持有的 KV Cache 分片
    """
    # 一次性把所有卡的 KV 聚合起来（All-Gather）
    k_full = all_gather_along_seq(k_local, group=cp_group)  # [B, H, L, D]
    v_full = all_gather_along_seq(v_local, group=cp_group)  # [B, H, L, D]

    # 单 token 的注意力，直接算
    scale = q.shape[-1] ** -0.5
    s = (q @ k_full.transpose(-1, -2)) * scale   # [B, H, 1, L]
    p = torch.softmax(s, dim=-1)
    out = p @ v_full                              # [B, H, 1, D]
    return out

为什么这里 All-Gather 反而更好？

• 只有一次集合通信：All-Gather 的延迟是 O(log N) 级别的 ring/tree 实现，比 N 次 P2P 往返少得多；
• 通信量可控：Decode 每步只聚合一次 KV，虽然总量是 O(L)，但相比 N 次往返的延迟开销，一次 All-Gather 更划算；
• 可与计算重叠：现代推理框架（vLLM、SGLang）会把 All-Gather 与上一层的计算、或与 MoE 的通信重叠。

4.3 更进一步的优化：KV Cache 分片 + 本地计算

在超长上下文 Decode 中，还有一种更激进的策略：每张卡只保留自己那一段 KV，Q 广播到所有卡，各卡算局部注意力，最后 Reduce。

def decode_attention_reduce(q, k_local, v_local, cp_group):
    """每卡算局部注意力，最后 All-Reduce 求和"""
    scale = q.shape[-1] ** -0.5
    s = (q @ k_local.transpose(-1, -2)) * scale  # [B, H, 1, L_local]
    p = torch.softmax(s, dim=-1)
    out_local = p @ v_local                       # [B, H, 1, D]

    # 局部 softmax 不能直接相加，需要带 max 的归约
    # 简化写法：先 All-Reduce 分子分母
    dist.all_reduce(out_local, op=dist.ReduceOp.SUM, group=cp_group)
    return out_local

注意：局部 softmax 直接 All-Reduce 是错的，因为各卡的 softmax 分母不同。正确做法是各卡同时归约 (max, sum_exp, weighted_sum) 三元组，或者用 Flash-Decoding 的 split-K 思路：各卡算局部 (m_i, l_i, acc_i)，再做一次带 max 修正的合并。

def merge_partial_softmax(partials):
    """partials: List[(m_i, l_i, acc_i)]，合并为全局结果"""
    m_global = torch.stack([p[0] for p in partials]).max(dim=0).values
    l_global = torch.zeros_like(partials[0][1])
    acc_global = torch.zeros_like(partials[0][2])
    for m_i, l_i, acc_i in partials:
        scale = torch.exp(m_i - m_global)
        l_global += l_i * scale
        acc_global += acc_i * scale
    return acc_global / l_global

5. Prefill 与 Decode 的策略对比

维度	Prefill	Decode
Q 长度	长（整条 prompt）	1（单 token）
计算/通信比	高，通信易掩盖	低，通信是瓶颈
推荐通信模式	P2P 环形传递 KV	All-Gather KV 或 局部计算 + Reduce
通信次数	N 轮 P2P	1 次集合通信
关键优化	因果跳过、Zig-Zag 负载均衡、通信计算重叠	通信与层计算重叠、split-K 归约
典型实现	Ring Attention、Striped Attention	Flash-Decoding、KV 分片 All-Gather

一句话总结：Prefill 用「时间换空间」——分 N 轮传，把通信藏进计算里；Decode 用「空间换时间」——一次聚合，把延迟压到最低。

6. 工程实践中的注意事项

• 序列长度必须能被 cp_size 整除，否则需要 padding 或变长切分，后者实现复杂度高很多。
• CP 与 TP 的组合：通常 CP 沿序列切、TP 沿 head 切，两者正交。但 CP 的 All-Gather 会与 TP 的 All-Reduce 竞争带宽，需要规划通信组。
• 因果掩码的正确性：环上传递的 KV 块要带上位置信息，判断哪些 Q 位置需要跳过，否则会算进「未来」token。
• 数值稳定性：在线 softmax 的 m 初始化为 -inf，第一轮要特殊处理，避免 exp(-inf - (-inf)) 产生 NaN。
• 框架支持：目前 vLLM、SGLang、Megatron-LM 都已支持 CP；Ring Attention 的原始实现在 Google 的 ringattention 仓库，FlashAttention-3 也集成了 CP 支持。

7. 小结

Ring Attention 用「环形接力」把长序列注意力的 KV 通信从 O(L×N) 压到 O(L)，配合在线 Softmax 保证数学等价。它在 Prefill 阶段大放异彩，因为计算密集、通信可被完全掩盖；而 Decode 阶段由于单 token 计算量太小，反而更适合一次性的 All-Gather 或 split-K 归约。理解这两种阶段的差异，是做好长上下文推理系统设计的关键。