2310.01889-ring-attention-blockwise-transformers-near-infinite-context
Ring Attention with Blockwise Transformers for Near Infinite Context
Ring Attention 的核心贡献是把 exact Transformer 的长序列瓶颈从“单卡必须驻留整段序列输出”改成“每张设备只驻留本地 query block,并让 key/value block 沿 ring 轮转”:它复用 blockwise exact attention 的在线 softmax 统计量,在不近似 attention 的前提下把最大上下文长度扩展到接近设备数倍;代价边界在于 exact attention 的二次计算仍然存在,通信是否可隐藏取决于 block size、计算吞吐和互联带宽。
Source
- Title: Ring Attention with Blockwise Transformers for Near-Infinite Context
- arXiv: https://arxiv.org/abs/2310.01889
- HTML: https://arxiv.org/html/2310.01889v4
- PDF: https://arxiv.org/pdf/2310.01889
- Code/Project: https://github.com/haoliuhl/ringattention
- Authors: Hao Liu, Matei Zaharia, Pieter Abbeel
- Submitted: 2023-10-03
- Current version read: v4, 2023-11-27
- Subjects: Computation and Language (cs.CL)
- Venue/status: ICLR 2024, according to the author homepage and OpenReview profile.
作者与关系
- Hao Liu: UC Berkeley.
- Matei Zaharia: UC Berkeley EECS / Sky Lab and Databricks CTO. In this paper, he brings the distributed systems and large-scale data/AI systems line into a long-context Transformer training problem.
- Pieter Abbeel: UC Berkeley.
阅读目标与判断边界
本笔记关注:
- Ring Attention 在数学语义上如何保持 exact softmax attention。
- 它如何把长上下文训练的 activation memory 从全局序列长度
依赖改成本地 block size 依赖。 - 它和 FlashAttention、sequence parallelism、ZeRO/FSDP、linear/sparse attention 路线的关系。
判断边界:
- 本文主要分析 v4 论文和作者公开资料;不把后续 Large World Model、Gemini long-context 或 Ring Attention 衍生实现的结果归入本文原始证据。
- 论文的 MFU 表格部分以图片呈现,TeX 源里只保留 context size 行;本笔记只记录可从正文和表格文本确认的数值。
- 论文强调 training,更少给出现代 serving stack 下的真实 inference latency / throughput benchmark;推理部分主要来自 Appendix C 的条件分析。
论文脉络
1. 研究问题、背景和价值
Transformer 长上下文有两类成本。
第一类是 attention score/probability matrix 的
第二类是每一层输出 activation 的驻留成本。即使 attention matrix 不落盘,下一层 self-attention 仍需要上一层所有 token 的输出;如果单张设备必须保存整段序列的上一层输出,activation memory 仍然随全局序列长度
因此,本文要解决的问题是:在保持 exact Transformer 语义的条件下,把超长序列的 layer output 和 attention computation 分摊到多设备上,并尽量让通信隐藏在计算后面。
这个问题重要的原因在于,长上下文应用覆盖文本检索、书籍、视频、代码库、科学数据、trial-and-error RL 经验和长反馈学习。这些场景里,模型需要同时访问很长历史,简单截断会改变任务定义,近似 attention 可能带来不可控质量损失。
2. 已有解决方案与不足
已有方向大致分为四类。
- Memory-efficient / FlashAttention 类方法避免
attention matrix 的 HBM materialization。它们保留 exact attention 语义,但主要处理单设备或单 kernel 内部的 IO 和 activation 问题,全局 layer output 仍需要被下一层访问。 - Blockwise Parallel Transformer (BPT) 将 attention 和 feedforward 都 blockwise 计算,把单层 activation 降到
量级。这里的 仍是当前设备需要持有的序列长度,因此超大 下仍受单设备 HBM 限制。 - FSDP / ZeRO 类方法分片参数、梯度和 optimizer state。它们解决模型状态冗余,但长上下文的 activation residency 是另一个维度。
- Sequence parallelism / all-to-all 类方法沿序列维分片。它们能分摊部分 activation,但通常需要 gather 全序列或引入难以完全隐藏的通信。论文特别指出,已有 ring topology self-attention 方案仍存在非重叠通信开销,在大 context 下会成为瓶颈。
Ring Attention 的切入点是把第 1 类的 blockwise exact attention、第 2 类的 blockwise FFN、以及第 4 类的跨设备序列分片合并起来,并用 ring communication 的顺序性降低同时驻留的 KV blocks。
3. 作者可能的思考路径
可以从一个已有事实出发:softmax attention 的每个 query block 对所有 key/value blocks 的聚合可以分块计算。只要维护每行的最大值、归一化分母和 numerator,key/value block 的访问顺序不会改变最终 exact softmax 结果。
如果顺序可以交换,就可以问一个系统问题:每张设备是否必须一次性拿到所有远端 KV?答案是否定的。设备只需要在某个时刻拿到一个 KV block,并把它并入当前 query block 的在线 softmax 累积状态。
然后出现通信问题:所有设备都需要遍历所有 KV blocks。直接 all-gather 会让每台设备临时存很多远端 KV,破坏内存目标;逐个 fetch 会产生等待。ring 的思路是让每个设备每轮只和相邻设备交换当前 KV block,同时计算自己手上 query block 与当前 KV block 的 attention。只要这一轮计算时间足够长,KV 传输就被隐藏。
这条思路的关键 intuition 是:长上下文让单次 blockwise attention 计算变重,重计算反而提供了隐藏通信的窗口。因此,超长序列中的
4. 核心假设或切入点
Ring Attention 依赖三个假设:
- Attention block 的合并具有顺序不变性:不同 KV blocks 的贡献可以按任意顺序并入在线 softmax 统计量。
- 每台设备的本地 block size
足够大,使得 attention block computation 时间大于 KV block transfer 时间。 - 互联拓扑支持相邻设备高效交换,例如 TPU torus、GPU NVLink 或足够高带宽的 GPU 集群互联。
5. 方法 / 系统 / 理论框架
5.1 基础 attention 语义
论文仍使用标准 scaled dot-product attention:
方法改变执行顺序和数据布局,保留原始 attention 语义。
5.2 跨设备 sequence block 分片
设有
对每一层:
- 每个 host 固定自己的 query block
。 - 每轮拿当前持有的
和 做 blockwise attention。 - 用 online softmax 统计量更新当前 query block 的 numerator、denominator 和 max score。
- 同时把当前
发送给 ring 上下一个 host,并从上一个 host 接收新的 。 - 所有 KV blocks 轮转完成后,每个 host 得到自己 query block 的 attention output。
- Feedforward 只作用于本地 block,可以在本地 blockwise 执行。
实现上,论文的 JAX 代码用 jax.custom_vjp 定义 forward/backward,用 jax.lax.ppermute 做相邻 host 的 KV block 交换。forward 中保存 output、denominator、max_score 等 online softmax 状态;backward 也沿 ring 扫描 KV blocks 来累积
5.3 内存上界
BPT 仍需要每层
Ring Attention 把每台 host 的驻留对象限制为常数个 blocks:
- 当前 query block:
个 block。 - 当前 key/value blocks:
个 blocks。 - 正在接收的 key/value blocks:
个 blocks。 - 当前 output block:
个 block。
因此一台 host 的核心 activation 上界为:
其中
5.4 通信隐藏条件
论文用一个简化 roofline 条件估算最小 block size。设每台 host 计算吞吐为
一个 query block 和一个 KV block 的 attention score 计算需要
发送和接收 key/value blocks 的通信量约为:
要让通信被计算覆盖,需要:
化简得到:
这解释了论文表 2 的硬件差异:A100 NVLink、TPU v3/v4/v5e 的最小 block size 约 1K,本地最小 sequence length 约 6K;A100 InfiniBand 因带宽较低,最小 block size 约 24.5K,本地最小 sequence length 约 149.5K。
5.5 与 parallelism 的组合
论文建议大规模训练中把 Ring Attention 和 FSDP / tensor parallelism 组合使用。FSDP 负责模型参数分片;Ring Attention 负责 sequence/context 维度扩展;tensor parallelism 可在 global token batch 过大时减少 batch 压力。
Appendix A 给出一个示例:512 张 A100 上,如果 30B 模型需要 8 张卡做模型分片,剩余设备可形成 32-way Ring Attention,把 context size 扩大 32 倍;如果 3B/7B 模型不需要 FSDP,则可以把更多设备用于 context 扩展。
5.6 推理适用性
论文主要讨论训练,因为训练有更重的 activation memory。Appendix C 指出 Ring Attention 也可用于 autoregressive inference:长上下文 KV cache 可以沿 ring 流动,让每张设备不必持有完整 KV cache。
推理时 query token 数通常为 1,通信隐藏条件变成另一个形态。论文以 32 张 TPUv5e 服务 LLaMA-7B 为例,传统按 attention heads 维切分时 batch size 1 约支持 256K context;Ring Attention 通过循环 KV cache 可把上下文扩展 32 倍。论文给出的隐藏条件是
5.7 长上下文并不免费
Appendix D 给出训练 FLOPs 随 context length 的增长方式。单 sequence FLOPs 近似为:
其中
在总 token 数固定的数据集上,从旧 context
这说明长上下文训练的 per-dataset 成本低于 naive 的 context length 平方倍数,但仍会显著增长。论文给出例子:170B 模型从 4K 扩到 10M context,长度增加 3072 倍,per-dataset FLOPs 增加 162.6 倍。
6. 结论链条
论文的结论链条可以压缩为:
- Exact attention 可以按 KV block 顺序无关地增量计算。
- 因此每张设备可以固定本地 query block,让 KV blocks 沿 ring 流动。
- 每轮计算
对当前 的贡献,同时传输下一轮 KV。 - 当
时,通信可以被 blockwise attention 计算隐藏。 - 每台设备只需要常数个 blocks 的 activation,单设备 memory 上界从
变为 。 - 在高带宽互联上,最大 context size 可以接近按设备数线性扩展。
- 这种扩展仍保留 exact attention 语义,但训练总 FLOPs 会随 context length 增长。
关键实验/定理
结果 1:最大训练 context size
- 设置:LLaMA-style 3B/7B/13B/30B 模型;8x A100 NVLink、32x A100 InfiniBand、TPUv3-512、TPUv4-1024、TPUv5e-256;对比 vanilla Transformer、memory efficient attention、memory efficient attention + feedforward / BPT、Ring Attention;使用 FSDP 保持总 token batch size 可比。
- 指标:end-to-end training 能支持的最大 context size,单位为
tokens。 - 结果:
| Hardware / model | Best prior | Ring Attention | Gain |
|---|---|---|---|
| 8x A100 NVLink, 3B | 64K | 512K | 8x |
| 8x A100 NVLink, 7B | 32K | 256K | 8x |
| 8x A100 NVLink, 13B | 16K | 128K | 8x |
| 32x A100 InfiniBand, 7B | 128K | 4096K | 32x |
| 32x A100 InfiniBand, 13B | 64K | 2048K | 32x |
| TPUv4-1024, 3B | 32K | 16384K | 512x |
| TPUv4-1024, 7B | 16K | 8192K | 512x |
| TPUv4-1024, 13B | 16K | 4096K | 256x |
| TPUv4-1024, 30B | 8K | 2048K | 256x |
- 解读:context size 的提升基本跟可用于 ring 的设备数成比例。这个实验强力支撑“单设备 activation memory 不再是主限制”的结论。
结果 2:通信隐藏所需 block size
- 设置:用
分析多种硬件互联。 - 指标:最小 block size
与本地最小 sequence length 。 - 结果:
| Hardware | FLOPS | Interconnect bandwidth | Minimal block size | Minimal local sequence length |
|---|---|---|---|---|
| A100 NVLink | 312 TF | 300 GB/s | 1.0K | 6.2K |
| A100 InfiniBand | 312 TF | 12.5 GB/s | 24.5K | 149.5K |
| TPU v3 | 123 TF | 112 GB/s | 1.1K | 6.6K |
| TPU v4 | 275 TF | 268 GB/s | 1.0K | 6.2K |
| TPU v5e | 196 TF | 186 GB/s | 1.1K | 6.3K |
- 解读:Ring Attention 对高带宽互联非常友好;跨节点 InfiniBand 要求大很多的本地 block,实际效率更依赖 batch、tensor parallelism 和网络拓扑。
结果 3:MFU 与大 context 训练
- 设置:BPT 与 Ring Attention 使用相同 BPT implementation;GPU batch size 2M tokens,TPU batch size 4M tokens;FSDP + JAX SPMD,部分设置叠加 tensor parallelism。
- 指标:model FLOPs utilization / throughput;论文表 4 给出 context size 对照和 MFU 曲线图。
- 结果:Ring Attention 在以下设置扩展 context:
| Model / hardware | BPT context | Ring Attention context |
|---|---|---|
| 7B, 8x A100 | 32K | 256K |
| 13B, 8x A100 | 16K | 128K |
| 13B, 32x A100 | 64K | 2048K |
| 30B, TPUv4-1024 | 16K | 2048K |
| 65B, TPUv4-1024 | 8K | 1024K |
- 解读:论文主张 Ring Attention 在显著扩长 context 时没有明显损害 MFU 或 throughput。由于具体 MFU 数值在图中,当前笔记只保留可复核的 context size 对照和结论强度。
结果 4:In-context RL / ExoRL
- 设置:把 Ring Attention 用于 Agentic Transformer / ExoRL,比较 BC、DT、AT + memory efficient attention、AT + BPT、AT + Ring Attention。
- 指标:6 个 ExoRL 任务累计 return。
- 结果:AT + BPT 在 128 trajectories 下 OOM;AT + Ring Attention 能训练 128 trajectories,total average 为 113.66,高于 32 trajectories 下 AT + BPT 的 111.13。单项任务上 Walker Run 从 105.88 到 110.45,Cheetah Run 从 178.75 到 181.34。
- 解读:这里的性能提升幅度不大,但关键价值是可训练更长的 trajectory context;它更多支撑可行性和应用方向,质量增益仍需要更多任务验证。
结果 5:LLaMA-13B 512K finetune 与 line retrieval
- 设置:32x A100 80GB,ShareGPT 125K cleaned conversations,LLaMA-13B finetune 到 512K context;评测 long-range line retrieval。
- 指标:不同 context length 下的 retrieval accuracy。
- 结果:Ring Attention-13B-512K 在长 context 下保持较高 retrieval accuracy;GPT-3.5-turbo-16K、Vicuna-16B-16K、Claude-2-100K 受各自最大 context 限制。
- 解读:该实验说明 Ring Attention 可用于长 context LLM finetune,并非只停留在 synthetic max-context benchmark;但数据和训练预算有限,不能视为完整长上下文能力评测。
证据链强度评估
强证据
- 内存分析清晰:把每台设备驻留对象列成 6 个 blocks,得到
上界,直接解释为何单设备 memory 与全局 解耦。 - 通信隐藏条件
给出了可复核的硬件约束,并解释 NVLink/TPU 与 InfiniBand 场景差异。 - 最大 context size 表覆盖多种硬件和模型规模,结果与“按设备数扩展”主张一致。
中等强度证据
- MFU/throughput 结论合理,但关键数值主要在图片曲线中,文档化复核粒度低于表格数字。
- ExoRL 证明更长 trajectory context 可训练,性能提升幅度较小,任务范围有限。
- LLaMA-13B 512K line retrieval 支撑长上下文可用性,但没有覆盖更广泛真实长文档、多跳推理、代码库和 agent 任务。
需要谨慎的推论
- “near-infinite context” 应理解为 activation memory 可以随设备数扩展,而非计算成本消失。Exact softmax attention 的 quadratic arithmetic 仍然存在。
- 训练中的高带宽通信重叠条件不能直接外推到所有推理服务场景;decode query length、KV cache layout、batching、tensor parallelism 和 serving scheduler 都会改变实际瓶颈。
- 2023 年实验基于 JAX/SPMD 和当时的 baseline,后续 FlashAttention-2/3、context parallelism、Ulysses、RingAttention 变体、FlexAttention/FlashInfer 等会改变工程对比。
OpenReview / 审稿意见吸收
- Venue status: 当前档案未记录公开 peer-review 状态。
- Public reviews: 当前档案未记录可可靠匹配的 OpenReview / ARR / 会议 reviewer comments。
- Ratings / confidence: 无公开评分可用于校准。
- Reviewer consensus: 暂无。
- Main criticisms: 暂无公开 reviewer 质疑可引用;可信度主要由论文、技术报告、项目证据和本地一致性检查决定。
- Author response: 暂无公开 rebuttal 记录。
- 对本文可信度的影响: 按未完成公开审稿吸收处理,结论需要依赖实验设置、baseline 强度、复现证据和跨论文一致性校准。
本地讨论补充
1. 讨论收敛点
- 初版归档将 Ring Attention 定位为 distributed exact attention / context parallelism 基础节点:它接在 FlashAttention/BPT 的 blockwise exact attention 之后,把同一块计算顺序扩展到多设备 ring 上。
- 它解决的主问题是长序列 activation residency 和 KV block 通信编排;它没有改变 attention family,也没有降低 exact attention 的二次计算阶。
- 对系统实践来说,Ring Attention 的可用性首先取决于本地 block 的 arithmetic intensity 是否足以隐藏互联通信,其次才是实现细节。
2. 修正后的理解
- Ring Attention 与 FlashAttention 的关系可以表述为:FlashAttention 关注单设备/单 kernel 内的 HBM-SRAM IO 和在线 softmax;Ring Attention 关注跨设备序列分片时如何让 KV blocks 流动并隐藏通信。两者可以组合。
- Ring Attention 与 ZeRO/FSDP 的关系可以表述为:ZeRO/FSDP 分片模型状态;Ring Attention 分片 sequence/context activation。大模型训练常需要两者叠加。
- Ring Attention 与 Lightning/DLA 的关系可以表述为:Ring Attention 保持 exact softmax 语义并分布式执行;Lightning/DLA 通过 linear attention / state memory 改变 architecture,以换取更好的长序列复杂度。
3. 后续复验指标
- 在现代 FA2/FA3/FlashInfer/FlexAttention 和 PyTorch/XLA/DeepSpeed context-parallel stack 下复测 training MFU、communication overlap rate、peak activation 和 end-to-end throughput。
- 分离 intra-node NVLink、inter-node InfiniBand、TPU ICI/torus 上的性能边界。
- 对 inference separately 记录 prefill/decode latency、KV cache residency、batch size、query length、context parallel degree 和 tensor parallel degree。
主要启发
- 长上下文系统优化至少有四个正交层:kernel IO、跨设备 context parallelism、模型结构压缩/linear attention、serving scheduler。Ring Attention 属于第二层。
- Exact attention 的 blockwise 顺序可交换性是系统设计入口:只要在线 softmax 统计量能正确合并,KV blocks 可以按 ring 顺序流动。
- 通信重叠条件比“是否用了 ring”更关键。高带宽互联下 1K block 就可能够用;低带宽互联下需要非常大的 local sequence 才能隐藏通信。
- 长 context 的训练成本应按 per-dataset FLOPs ratio 看。Ring Attention 让 memory feasible,但不会让 10M/100M context 变成低成本训练。
- 对 RL 和 agentic workflow,Ring Attention 的价值在于把更长 trial-and-error history 放进单个 context;是否带来能力提升仍由数据、任务结构和训练目标决定。
局限
- Exact attention 的
arithmetic 仍然存在,超长 context 下 FLOPs 成本很高。 - 对硬件互联敏感,A100 InfiniBand 的最小 block size 显著大于 NVLink/TPU 场景。
- Inference 评估主要是理论条件和示例分析,缺少现代 serving engine 下的系统 benchmark。
- 质量实验有限:ExoRL 提升幅度较小,line retrieval 是长上下文检索能力的局部测试。
- 论文 baseline 处在 2023 年软件栈,后续 context parallelism、Ulysses、RingAttention 变体、FA2/FA3、FlexAttention、FlashInfer 和 vendor kernels 需要重新对比。
- 方法要求每层都能以 blockwise 方式组织 attention 与 FFN;复杂 attention mask、packed sequences、多模态 token layout 和 MoE routing 可能引入额外系统约束。
跨论文关系
- 与 2205.14135:FlashAttention 用 tiling + online softmax 避免
attention matrix materialization;Ring Attention 依赖同类 blockwise exact attention 语义,并把 KV block 访问顺序放到跨设备 ring 中。 - 与 2307.08691:FA2 解决单 GPU attention kernel 的 work partitioning、sequence parallelism 和 non-matmul FLOPs;Ring Attention 解决多设备 context parallelism 中的 KV communication overlap。两者位于不同系统层。
- 与 1910.02054:ZeRO/FSDP 减少 model states 冗余,Ring Attention 减少单设备 sequence activation residency。两者共同解释现代超长上下文大模型训练为什么需要多维 parallelism。
- 与 2309.14509 DeepSpeed Ulysses:两者都保持 exact attention 语义并处理跨设备 context/sequence 分片。Ulysses 通过 all-to-all 在 sequence-partitioned 和 head-partitioned layout 间转换;Ring Attention 通过 KV block ring rotation 与 online softmax state 累积远端 blocks。两者提供了 sequence/head transpose 与 KV-stream 两种 context parallelism 语言。
- 与 2308.16369:Sarathi 处理 inference serving 中 prefill/decode batching 和 pipeline bubbles;Ring Attention 处理训练/推理中的超长 context activation/KV cache 分布。二者都使用 chunk/block 视角,但作用层级不同。
- 与 2405.17381 和 2606.10650:Ring Attention 保留 exact softmax attention;Lightning/DLA 改为 linear attention / state memory。前者是系统并行路线,后者是 architecture 路线。
- 与 2026-04-24:DeepSeek-V4 的 million-token context 依赖 compressed attention、heterogeneous KV cache 和 deterministic kernels;Ring Attention 是更早的 exact attention context-parallel 基础节点,可作为对照理解“保持 exact semantics”和“引入压缩”的取舍。
- 跨论文关系定位:记录 Ring Attention / distributed exact attention / context parallelism 节点,连接 FlashAttention family、ZeRO/FSDP、Sarathi 和 long-context architecture/serving 材料。
Reference Intake Brief
Target
- Intended target system: 新增论文笔记和作者页,更新索引行与 long-context systems 关系章节。
- Existing related assets: 2205.14135、2307.08691、1910.02054、2308.16369、2405.17381、2606.10650、2026-04-24。
- Proposed form: 新建独立 Markdown 文档;更新作者档案、索引行和对应论文的关系章节。
Reusable Elements
- Ring Attention = blockwise exact attention + ring KV rotation + communication/computation overlap。
- 单设备 activation 上界从
型依赖变为 。 - 通信隐藏条件
是评估部署可行性的核心指标。
Risks
- Copyright/over-copying: 只保留必要公式、表格数字和本地分析,不复制长段论文原文。
- Unsourced or unverifiable claims: 作者主页、OpenReview、GitHub、DBLP 和 arXiv 信息均有来源;不记录未验证中文姓名。
- Tone/brand mismatch: 保持技术分析风格,避免夸张化理解 “near-infinite”。
- Safety/compliance issues: 无直接安全双用途操作细节。
- Overlap with existing assets: 与 FlashAttention/FA2 和 Lightning/DLA 关系已区分系统层与 architecture 层。
Skipped
| Material | Reason |
|---|---|
Blockwise Parallel Transformer (2305.19370) 完整分析 |
属于 Ring Attention 的直接上游论文,当前只在关系中记录,后续可独立归档。 |
Large World Model (2402.08268) |
属于 RingAttention 后续应用线,当前任务只分析 2310.01889。 |
| 现代 RingAttention / context parallelism 实现 benchmark | 超出本文版本范围,需要单独工程调研。 |
Recommendation
Decision: merge
Why: Ring Attention 是 exact long-context Transformer training 的关键 distributed attention 节点,补齐当前档案中 FlashAttention kernel、ZeRO/FSDP training memory、Sarathi serving scheduling、Lightning/DLA architecture 路线之间的缺口。它提供了一个清晰判断框架:长上下文 memory feasibility 可以靠 blockwise context parallelism 解决,但计算成本和互联条件仍然决定真实可用性。