2023年10月，伯克利大学的研究团队提交了一篇论文，声称可以让Transformer处理"接近无限"长度的上下文。实验数据令人咋舌：在512块TPU v4上，7B模型可以训练超过800万Token的序列；在1024块TPU v4上，这个数字飙升到1600万。

三个月后，同一团队发布了World Model论文，展示了如何用这项技术训练能理解一小时视频的多模态模型。2024年，主流大模型的上下文窗口从4K一路膨胀——Llama-4宣称支持1000万Token，Gemini 1.5 Pro达到200万。这些数字背后，都指向同一个技术：Ring Attention。

一个100M Token的序列需要多少内存

要理解Ring Attention为何如此重要，需要先直面一个残酷的现实。

假设一个相对"modest"的模型：隐藏维度 $d=1024$，批大小 $b=1$。处理一个1亿Token的序列，光是存储每一层的输出就需要多少内存？

自注意力的核心计算是 $\text{Attention}(Q,K,V) = \text{softmax}(QK^T/\sqrt{d})V$。每层需要存储的激活值大小是 $2bsh$ 字节（bfloat16精度），其中 $s$ 是序列长度，$h$ 是隐藏维度。代入数值：

$$2 \times 1 \times 100,000,000 \times 1024 \times 2 \text{ bytes} \approx 400 \text{ GB}$$

这还只是一层。一个典型的Transformer有几十层，加上前馈网络、梯度、优化器状态，单层激活值就已经超过了任何GPU的容量——NVIDIA H200的HBM是141GB，AMD MI300X是192GB，最新的GB200也不过288GB。

更残酷的是注意力的内存复杂度。标准的注意力矩阵 $S=QK^T$ 大小是 $s \times s$，对于1亿Token，这个矩阵需要：

$$100,000,000 \times 100,000,000 \times 2 \text{ bytes} = 20,000 \text{ TB}$$

显然，没有任何单机能处理这个规模。长上下文的本质挑战，不是参数量，而是激活值的内存墙。

Flash Attention解决了什么，又没解决什么

Flash Attention是这个领域的里程碑。它通过分块计算（tiling）避免显式存储完整的注意力矩阵，将内存复杂度从 $O(s^2)$ 降到 $O(s)$。

核心思想是：不需要一次性计算完整的 $QK^T$，可以将Q、K、V分成小块，逐块计算注意力，同时使用Online Softmax技巧在块间传递归一化统计量。具体而言，每个块计算完成后，只保留三个标量：当前最大值 $m$、分子累加器 $N$、分母累加器 $D$。

但Flash Attention有一个根本性的局限：它仍然要求所有数据在同一设备上。

分块计算可以减少激活值内存，但每一层的输出仍然需要存储，供下一层使用。这意味着，无论Flash Attention多么高效，序列长度 $s$ 仍然受限于单设备的HBM容量。

伯克利团队在BPT（Blockwise Parallel Transformer）论文中做过一个实验：在80GB的A100上，7B模型使用Flash Attention + 梯度检查点，最大上下文大约是16K Token。要突破这个瓶颈，唯一的出路是跨设备分布。

Ring Attention的核心洞察

Ring Attention的出发点是一个简单但深刻的观察：

分块注意力的计算顺序是可交换的——只要正确合并各块的归一化统计量，最终结果与计算顺序无关。

这意味着，如果设备A持有查询块 $Q_1$，设备B持有键值块 $K_2, V_2$，设备A不需要等待从B获取数据才能开始计算——它可以先用本地的 $K_1, V_1$ 计算，同时异步接收B的数据。

Ring Attention将这个观察转化为一个精巧的分布式算法：

  1. 将序列切成 $N$ 块，分配给 $N$ 个设备
  2. 每个设备持有一个查询块 $Q_i$ 和对应的 $K_i, V_i$
  3. 设备按环形拓扑组织：设备 $i$ 向设备 $(i+1) \mod N$ 发送数据，从设备 $(i-1) \mod N$ 接收数据
  4. 每一轮，设备计算本地 $Q_i$ 与当前 $K_j, V_j$ 的注意力，同时异步传递/接收下一批 $K, V$
  5. 经过 $N$ 轮，每个设备都看到了所有 $K, V$，完成了完整注意力的计算

关键在于通信与计算的重叠。如果计算时间大于通信时间，网络延迟就被完全隐藏了。

sequenceDiagram
    participant GPU0
    participant GPU1
    participant GPU2
    participant GPU3
    
    Note over GPU0,GPU3: 第一轮：计算本地K/V，同时传递
    GPU0->>GPU1: K0, V0
    GPU1->>GPU2: K1, V1
    GPU2->>GPU3: K2, V2
    GPU3->>GPU0: K3, V3
    
    Note over GPU0,GPU3: 第二轮：计算接收的K/V，继续传递
    GPU0->>GPU1: K3, V3
    GPU1->>GPU2: K0, V0
    GPU2->>GPU3: K1, V1
    GPU3->>GPU0: K2, V2
    
    Note over GPU0,GPU3: ...重复直到所有设备都见过所有K/V

Online Softmax：分块计算的前提

要实现分块注意力，必须解决一个数学问题：Softmax如何分块计算？

标准Softmax定义是：

$$\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_{j=1}^{n} e^{x_j}}$$

分母需要所有元素的指数和，这看起来无法分块。但有一个技巧：

设当前块的最大值是 $m_1$，已经累积的归一化因子是 $l_1 = \sum_{k \in \text{processed}} e^{x_k - m_1}$。当处理新块时，新块的最大值是 $m_2$，新块的指数和是 $l_2 = \sum_{k \in \text{new\_block}} e^{x_k - m_2}$。

新的全局最大值是 $m = \max(m_1, m_2)$，新的全局归一化因子是：

$$l = l_1 \cdot e^{m_1 - m} + l_2 \cdot e^{m_2 - m}$$

这个公式允许我们增量地更新归一化统计量，而不需要一次性访问所有元素。

更稳定的实现使用对数空间：

$$\log l = \log(e^{\log l_1 + m_1 - m} + e^{\log l_2 + m_2 - m})$$

这被称为LogSumExp技巧，是Flash Attention和Ring Attention共同依赖的数学基础。

内存复杂度的革命性改进

Ring Attention的内存分析是其最令人印象深刻的成果之一。

对于bfloat16精度，每个设备需要存储：

  • 1个查询块：$bc d$ 字节
  • 1个键块 + 1个值块：$2 \cdot bc d$ 字节
  • 接收缓冲区（下一批K, V）：$2 \cdot bc d$ 字节
  • 输出缓冲区：$bc d$ 字节

总计 $6bcd$ 字节，其中 $c$ 是块大小。注意这个式子不包含序列长度 $s$——块大小 $c$ 是由硬件带宽决定的常数，与实际处理的序列长度无关。

方法 单层激活值内存
标准Transformer $2bns^2$
Memory Efficient Attention $8bsh$
Blockwise Parallel Transformer $2bsh$
Ring Attention $6bch$

$b$=批大小，$s$=序列长度，$h$=隐藏维度，$n$=注意力头数，$c$=块大小（常数）

这个改进意味着：上下文长度可以随设备数量线性扩展。如果有 $N$ 个设备，Ring Attention可以处理 $N$ 倍长的序列。

最小序列长度：带宽与算力的博弈

通信与计算的重叠不是免费的午餐。它要求：

$$\text{计算时间} \geq \text{通信时间}$$

设设备算力为 $F$ FLOPS，互连带宽为 $B$ bytes/s，块大小为 $c$，隐藏维度为 $d$。

计算一个注意力块需要：

  • $QK^T$：$2dc^2$ FLOPs
  • 注意力分数 $\times V$：$2dc^2$ FLOPs
  • 总计：$4dc^2$ FLOPs

传递一个K/V块需要：

  • $K + V$：$4dc$ bytes

要重叠，需要：

$$\frac{4dc^2}{F} \geq \frac{4dc}{B}$$

简化得：

$$c \geq \frac{F}{B}$$

这就是最小块大小的公式。而最小序列长度是这个值的6倍（因为需要存储多个块）。

硬件配置 算力 (TF) 带宽 (GB/s) 最小块大小 最小序列长度
A100 + NVLink 312 300 ~1K ~6K
A100 + InfiniBand 312 12.5 ~25K ~150K
TPU v4 275 268 ~1K ~6K
TPU v5e 196 186 ~1K ~6K

这解释了为什么Ring Attention需要高速互连。使用普通InfiniBand连接的A100，最小序列长度高达15万Token，在实用中可能得不偿失。

因果掩码的负载均衡挑战

自回归模型（如GPT）使用因果掩码，每个Token只能attend到它之前的Token。这给Ring Attention带来了一个微妙但严重的问题：负载不均衡。

考虑4个设备，每个处理序列的1/4：

  • 设备0处理Token 0-24，计算 $25 \times 25$ 的注意力
  • 设备1处理Token 25-49，计算 $25 \times 50$ 的注意力
  • 设备2处理Token 50-74，计算 $25 \times 75$ 的注意力
  • 设备3处理Token 75-99，计算 $25 \times 100$ 的注意力

设备3的工作量是设备0的4倍！在环形传递中，后面的设备会成为瓶颈，前面的设备在等待时空转。

Striped Attention（条纹注意力）通过重新分配Token解决了这个问题。核心思想是：不按顺序切分，而是交错分配。

假设4个设备，Striped Attention的分配方式是：

  • 设备0：Token [0, 4, 8, 12, …]
  • 设备1：Token [1, 5, 9, 13, …]
  • 设备2：Token [2, 6, 10, 14, …]
  • 设备3：Token [3, 7, 11, 15, …]

这样每个设备都持有"早期"和"晚期"Token的混合，因果掩码下的计算量更加均衡。

实验表明，在256K序列长度上，Striped Attention比标准Ring Attention快1.45倍；在786K Token上，加速比达到1.65倍。

实战：从8K到118K的内存演变

AKASA团队的一篇技术博客详细记录了使用Ring Attention扩展上下文的实际过程。他们从一个基线开始：在单块H100上微调Llama-8B，上下文被限制在1K Token左右。

第一步是FSDP分片。将模型权重、优化器状态、梯度分布到4块GPU上，每块GPU的内存占用从70GB降到约12GB，上下文窗口扩展到8K Token。但随着序列变长，激活值开始主导内存占用。

第二步是引入Ring Attention。激活值被分布到多块GPU上，峰值内存从60GB降到20GB。更长的序列成为可能。

第三步是梯度检查点。通过在反向传播时重计算激活值，进一步压缩内存。最终配置（FSDP + Ring Attention + 梯度检查点）可以在4块GPU上处理118K Token，峰值内存约30GB。

代价是吞吐量下降约58%。但这是值得的——否则根本无法训练这个规模的上下文。

与其他方案的对比

长上下文训练有多种技术路径，Ring Attention不是唯一选择。

DeepSpeed Ulysses是另一种序列并行方案。它使用All-to-All集合通信，在注意力计算前将序列分块，计算后再聚合。优点是集合通信通常比P2P更稳定；缺点是并行度受限于注意力头数，不适用于GQA/MQA场景。

All-Gather方案通过预先收集所有K/V块来计算注意力。实现简单，但需要在每个设备上临时存储完整序列，内存效率低于Ring Attention。

**USP（Unified Sequence Parallelism）**是一个有趣的混合方案，结合了Ulysses和Ring Attention的优点。在小规模场景使用Ulysses获得高吞吐，在大规模场景切换到Ring Attention扩展上下文。

方案 通信方式 并行度限制 内存效率 适用场景
Ring Attention P2P 无限制 最高 超长序列
DeepSpeed Ulysses All-to-All 注意力头数 中等 短序列、高吞吐
All-Gather All-Gather 无限制 低 原型验证
USP 混合 灵活 高 通用

实际工程中，USP正在成为主流选择。NVIDIA的TransformerEngine已经集成了这一方案。

World Model：百万Token的实际应用

Ring Attention最引人注目的应用是World Model论文。研究团队训练了一个能够处理100万Token的多模态模型，可以理解一小时的完整视频。

关键技术栈包括：

  • 渐进式上下文扩展：从4K开始，逐步扩展到1M，降低训练成本
  • 混合模态训练：文本、图像、视频混合，避免模态坍缩
  • 掩码序列打包：高效处理不同长度的序列

模型在视频问答任务上展现出惊人的能力。例如，给定一段一小时的视频，其中某个瞬间显示车里有柠檬，模型能够正确回答相关问题——这需要模型"记住"视频中每个细节。

工程实践指南

要在生产环境中使用Ring Attention，需要考虑几个关键因素。

硬件要求：高速互连是必须的。NVLink（300GB/s+）或TPU ICI（200GB/s+）可以获得良好的通信计算重叠。普通InfiniBand（12.5GB/s）会导致最小序列长度过高。

最小序列长度：通常需要至少6K-10K Token才能有效重叠。对于更短的序列，传统方法可能更快。

与FSDP结合：Ring Attention处理序列并行，FSDP处理模型并行。两者结合可以同时扩展上下文和模型大小。

因果掩码：务必使用Striped Attention或Zigzag变体，否则会遭遇严重的负载不均衡。

实现选择：

  • PyTorch: zhuzilin/ring-flash-attention（基于Flash Attention）
  • JAX: 作者官方实现（论文附录）
  • 生产环境: NVIDIA TransformerEngine的 attn_forward_func_with_cp
# 使用示例（PyTorch + yunchang库）
from yunchang import LongContextAttention, set_seq_parallel_pg

# 设置进程组：Ulysses=2, Ring=4
set_seq_parallel_pg(sp_ulysses_degree=2, sp_ring_degree=4, rank, world_size)

# 创建注意力层
attn = LongContextAttention(ring_impl_type="zigzag")

# 前向传播
output = attn(q, k, v, causal=True)

局限与未来

Ring Attention并非完美。它的主要局限包括：

吞吐量损失：即使通信完全重叠，分块计算的效率仍低于连续计算。实测中通常有20-60%的吞吐量下降。

推理挑战：自回归生成时每次只产生一个Token，计算量极小，难以与通信重叠。Flash Decoding等变体正在解决这个问题。

大规模部署风险：P2P通信在大规模部署中可能出现死锁，需要仔细设计通信模式。

新兴竞争者：线性注意力（如Mamba）和状态空间模型提供了另一种长上下文解决方案，计算复杂度是真正的 $O(n)$，而非依赖于分布式。

但就目前而言，Ring Attention是唯一一个可以在不牺牲模型质量的前提下，将上下文扩展到百万级Token的实用技术。随着模型规模和上下文需求的持续增长，它将继续在长上下文LLM的训练中扮演关键角色。

参考资料

  1. Liu, H., Zaharia, M., & Abbeel, P. (2024). Ring Attention with Blockwise Transformers for Near-Infinite Context. ICLR 2024.
  2. Liu, H., & Abbeel, P. (2023). Blockwise Parallel Transformer for Large Context Models. NeurIPS 2023.
  3. Nrusimha, A., et al. (2023). Striped Attention: Faster Ring Attention for Causal Transformers.
  4. Liu, H., et al. (2025). World Model on Million-Length Video And Language With Blockwise RingAttention. ICLR 2025.
  5. Fang, J., & Zhao, S. (2024). USP: A Unified Sequence Parallelism Approach for Long Context Generative AI.
  6. Dao, T., et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022.
  7. Coconut Mode - Ring Attention Explained
  8. GPU MODE Lecture 13: Ring Attention
  9. AKASA Blog - Ring Attention: Shedding Light on the Dark Art of Attention Sharding
  10. GitHub: zhuzilin/ring-flash-attention