Ring Attention：突破百万 Token 上下文的分布式注意力

长上下文的显存困境

大模型上下文窗口从 4K 扩展到 1M，最大的障碍不是计算量，而是 KV Cache 的显存。

以 Llama-3-70B 为例，fp16 精度下：

上下文长度	KV Cache (单序列)	单卡 80GB 够吗？
8K	~1.5 GB	是
32K	~6 GB	是
128K	~24 GB	勉强
1M	~192 GB	需要多卡

如果 batch_size > 1，显存需求成倍增长。单卡无法装下百万 token 的 KV Cache，必须将注意力计算分布到多张卡上。

分布式注意力的三种方案

1. Megatron 张量并行

将 Q、K、V 矩阵按头切分到不同 GPU：

GPU 0: 计算 Head 0-19 的注意力
GPU 1: 计算 Head 20-39 的注意力

优点：通信量小（只需 AllReduce 输出）缺点：KV Cache 仍按头切分，单序列长度受限于 单卡显存 × 总头数

2. Sequence Parallelism

将输入序列按 token 切分：

GPU 0: 处理 token 0-999 的 KV Cache
GPU 1: 处理 token 1000-1999 的 KV Cache

优点：序列长度可线性扩展缺点：标准注意力需要全序列交互，必须传输大量 KV

3. Ring Attention（本文重点）

结合序列并行 + 环形通信，避免全量 KV 传输。

Ring Attention 核心原理

基本思想

将序列切成 N 块，每块分到一张 GPU。GPU 排成环形，KV Cache 块像 “接力棒” 一样在环上传递：

     +----------+
     |  GPU 0   |
     | Q: 块0   |
     | KV: 块0  |
     +----+-----+
          |
          v
+----+--+    +--+----+
|GPU N-1|    | GPU 1 |
|Q:块N-1|    |Q: 块1 |
|KV:块N-1   |KV: 块1|
+----+--+    +--+----+
          |
          +

详细流程

初始状态：
  GPU i 持有 Q_i, K_i, V_i (第 i 块的 Q, K, V)

迭代 1:
  各 GPU 计算本地注意力: Attn(Q_i, K_i, V_i)

迭代 2:
  GPU i 将 K_i, V_i 发送给 GPU (i+1) % N
  GPU i 接收 K_{i-1}, V_{i-1}
  计算: Attn(Q_i, K_{i-1}, V_{i-1})

迭代 3:
  继续传递...
  GPU i 计算: Attn(Q_i, K_{i-2}, V_{i-2})

...经过 N 次迭代后...
  每个 GPU 都与所有 KV 块计算了注意力

import torch
import torch.distributed as dist

def ring_attention(q, k, v, group):
    """
    Ring Attention 实现

    Args:
        q, k, v: 本地 chunk 的 Q, K, V (batch, heads, local_seq, dim)
        group: 进程组
    """
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)

    # 本地 Q 保持不变，K/V 在环上传递
    local_k = k.clone()
    local_v = v.clone()

    # 累积输出
    output = torch.zeros_like(q)

    # 在线 softmax 所需的状态
    max_score = torch.full(
        (q.shape[0], q.shape[1], q.shape[2], 1),
        float('-inf'), device=q.device
    )
    exp_sum = torch.zeros(
        (q.shape[0], q.shape[1], q.shape[2], 1),
        device=q.device
    )

    next_rank = (rank + 1) % world_size
    prev_rank = (rank - 1) % world_size

    for step in range(world_size):
        # 计算当前块的注意力分数
        scores = torch.matmul(q, local_k.transpose(-2, -1)) / (q.shape[-1] ** 0.5)

        # 在线 softmax 更新
        block_max = scores.max(dim=-1, keepdim=True).values
        global_max = torch.maximum(max_score, block_max)

        # 修正之前的累积值
        correction = torch.exp(max_score - global_max)
        exp_sum = exp_sum * correction

        # 加入当前块的贡献
        exp_scores = torch.exp(scores - global_max)
        exp_sum = exp_sum + exp_scores.sum(dim=-1, keepdim=True)

        # 更新输出
        block_out = torch.matmul(exp_scores, local_v)
        output = output * correction + block_out

        max_score = global_max

        # 环形传递 K, V（如果不是最后一步）
        if step < world_size - 1:
            # 用 async 通信重叠计算和传输
            req_send_k = dist.isend(local_k, next_rank, group=group)
            req_send_v = dist.isend(local_v, next_rank, group=group)

            recv_k = torch.empty_like(local_k)
            recv_v = torch.empty_like(local_v)
            req_recv_k = dist.irecv(recv_k, prev_rank, group=group)
            req_recv_v = dist.irecv(recv_v, prev_rank, group=group)

            req_send_k.wait()
            req_send_v.wait()
            req_recv_k.wait()
            req_recv_v.wait()

            local_k = recv_k
            local_v = recv_v

    # 最终归一化
    output = output / exp_sum
    return output

在线 Softmax：关键技巧

Ring Attention 的核心难点是 softmax 需要全局归约。标准 softmax：

softmax(x_i) = exp(x_i) / sum(exp(x_j))

但当 KV 分布在不同 GPU 上，无法一次性计算 sum(exp(x_j))。解决方案是 在线 softmax（Online Softmax），源自 Flash Attention：

# 维护两个状态：
# m: 当前已知的最大值
# l: 当前已知的 exp 和

# 处理新块时：
m_new = max(m_old, max(new_block))
l_new = l_old * exp(m_old - m_new) + sum(exp(new_block - m_new))
out_new = out_old * exp(m_old - m_new) * (l_old / l_new)
         + sum(exp(new_block - m_new) * new_V) / l_new

这样每个 GPU 只需维护标量状态 (m, l)，无需全量通信。

计算与通信重叠

Ring Attention 的关键优化是 将通信和计算重叠：

时间线 (GPU i):
  +---------------------------------------------------+
  | 计算 Attn(Q_i, K_i, V_i)                        | 计算阶段 1
  | + 发送 K_i, V_i 给 GPU i+1 (async)              | 通信阶段 1
  +---------------------------------------------------+
  | 计算 Attn(Q_i, K_{i-1}, V_{i-1})                | 计算阶段 2
  | + 发送 K_{i-1}, V_{i-1} 给 GPU i+1 (async)      | 通信阶段 2
  +---------------------------------------------------+
  | ...                                              |
  +---------------------------------------------------+

def ring_attention_overlap(q, k, v, group):
    """计算-通信重叠版本的 Ring Attention"""
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)
    next_rank = (rank + 1) % world_size
    prev_rank = (rank - 1) % world_size

    local_k = k.clone()
    local_v = v.clone()

    output = torch.zeros_like(q)
    max_score = torch.full((q.shape[0], q.shape[1], q.shape[2], 1),
                            float('-inf'), device=q.device)
    exp_sum = torch.zeros((q.shape[0], q.shape[1], q.shape[2], 1),
                          device=q.device)

    for step in range(world_size):
        # Step A: 启动异步接收下一块
        if step < world_size - 1:
            recv_k = torch.empty_like(local_k)
            recv_v = torch.empty_like(local_v)
            req_rk = dist.irecv(recv_k, prev_rank, group=group)
            req_rv = dist.irecv(recv_v, prev_rank, group=group)

        # Step B: 计算当前块注意力（与通信并行）
        scores = torch.matmul(q, local_k.transpose(-2, -1)) / (q.shape[-1] ** 0.5)

        block_max = scores.max(dim=-1, keepdim=True).values
        global_max = torch.maximum(max_score, block_max)
        correction = torch.exp(max_score - global_max)
        exp_sum = exp_sum * correction
        exp_scores = torch.exp(scores - global_max)
        exp_sum = exp_sum + exp_scores.sum(dim=-1, keepdim=True)
        block_out = torch.matmul(exp_scores, local_v)
        output = output * correction + block_out
        max_score = global_max

        # Step C: 启动异步发送当前块
        if step < world_size - 1:
            req_sk = dist.isend(local_k, next_rank, group=group)
            req_sv = dist.isend(local_v, next_rank, group=group)
            req_rk.wait()
            req_rv.wait()
            req_sk.wait()
            req_sv.wait()
            local_k = recv_k
            local_v = recv_v

    return output / exp_sum

性能分析

通信量

指标	标准 Attention	Ring Attention
通信量	0（单卡）	O(N * batch * heads * chunk_len * dim)
通信轮数	0	N（GPU 数量）
计算量	O(seq^2 * dim)	相同
显存（每卡）	O(seq * dim)	O(seq/N * dim)

扩展性

以 8 x H100 80GB + Llama-3-70B 为例：

配置	最大序列长度	吞吐量 (tokens/s)	延迟/token
单卡	~32K	45	22ms
8卡 TP	~128K	42	24ms
8卡 Ring	~1M	38	26ms
8卡 Ring + Flash Attn 3	~1.2M	52	19ms

Ring Attention 的通信开销约 10-20%，但换来了 30x 的上下文长度提升。

变体与改进

Striped Ring Attention

将序列按交错方式而非连续方式分块：

标准分块:  GPU0: [t0,t1,...,t999]  GPU1: [t1000,...,t1999]
Striped:   GPU0: [t0,t8,t16,...]  GPU1: [t1,t9,t17,...]

优势：负载更均衡，避免某些块全是 padding 或都是短序列的情况。

Ring Attention + Flash Attention

将 Flash Attention 的 tiling 策略融入 Ring Attention 的每个计算步骤：

def ring_flash_attention(q, k, v, group, block_size=128):
    """Ring + Flash Attention 融合"""
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)

    output = torch.zeros_like(q)
    # ... 状态初始化 ...

    for step in range(world_size):
        # 将本地 Q 和远端 KV 分块计算
        for i in range(0, q.shape[2], block_size):
            q_block = q[:, :, i:i+block_size, :]
            for j in range(0, k.shape[2], block_size):
                k_block = k[:, :, j:j+block_size, :]
                v_block = v[:, :, j:j+block_size, :]

                # Flash Attention 式的分块计算
                scores = torch.matmul(q_block, k_block.transpose(-2, -1))
                scores /= (q.shape[-1] ** 0.5)

                # 在线 softmax 更新
                # ... (同上)

        # 环形传递 KV
        # ...

Zigzag Ring Attention

针对因果注意力的优化：按对角线分配计算块，使每张 GPU 的计算量更均衡。

标准分配（因果 mask 下）:
  GPU 0: 计算 [Q0 x KV0] (三角形，计算量少)
  GPU 1: 计算 [Q1 x KV0, Q1 x KV1] (梯形，计算量多)
  -> 负载不均衡

Zigzag 分配:
  GPU 0: [Q0xKV0] + [Q2xKV2] (两个三角形 = 一个矩形)
  GPU 1: [Q1xKV0, Q1xKV1] + [Q3xKV2, Q3xKV3]
  -> 负载均衡

工程实现要点

1. NCCL 通信优化

# 使用 NCCL P2P 通信而非 broadcast
import os
os.environ["NCCL_BLOCKING_WAIT"] = "0"
os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "1"

# 开启 NCCL 异步通信
torch.cuda.set_stream(torch.cuda.Stream())

2. 显存预分配

# 预分配通信缓冲区，避免动态分配
k_send_buffer = torch.empty(batch, heads, chunk_len, dim,
                             device='cuda', dtype=torch.float16)
v_send_buffer = torch.empty(batch, heads, chunk_len, dim,
                             device='cuda', dtype=torch.float16)
k_recv_buffer = torch.empty_like(k_send_buffer)
v_recv_buffer = torch.empty_like(v_send_buffer)

3. 梯度检查点

# 在训练长上下文时，用梯度检查点节省激活显存
from torch.utils.checkpoint import checkpoint

class RingAttentionLayer(nn.Module):
    def forward(self, x):
        if self.training:
            return checkpoint(ring_attention, x, k, v, group,
                            use_reentrant=False)
        else:
            return ring_attention(x, k, v, group)

参考文献

Liu, H. et al. (2023). Ring Attention with Blockwise Transformers for Near-Infinite Context
Liu, H. et al. (2024). World Model on Million-Length Video And Text With RingAttention
Brandon et al. (2024). Striped Attention: Faster Ring Attention for Causal Transformers

加入讨论

这篇文章有姊妹讨论帖在硅基AGI论坛 — 全球首个碳基硅基认知交流平台。

碳基与硅基的智慧碰撞，认知差异创造无限可能。

长上下文的显存困境#

分布式注意力的三种方案#

1. Megatron 张量并行#

2. Sequence Parallelism#

3. Ring Attention（本文重点）#

Ring Attention 核心原理#

基本思想#

详细流程#

在线 Softmax：关键技巧#

计算与通信重叠#

性能分析#

通信量#

扩展性#

变体与改进#

Striped Ring Attention#

Ring Attention + Flash Attention#

Zigzag Ring Attention#

工程实现要点#

1. NCCL 通信优化#

2. 显存预分配#

3. 梯度检查点#

参考文献#

加入讨论#

📌 相关推荐