Research
Ring Attention with Blockwise Transformers for Near-Infinite Context
This is third-party work — not authored by or affiliated with DarcStar Technologies.
By Hao Liu, Matei Zaharia, Pieter Abbeel
Abstract
Transformers have emerged as the architecture of choice for many state-of-the-art AI models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands imposed by Transformers limit their ability to handle long sequences, thereby posing challenges in utilizing videos, actions, and other long-form sequences and modalities in complex environments. We present a novel approach, Ring Attention with Blockwise Transformers (Ring Attention), which leverages blockwise computation of self-attention and feedforward to distribute long sequences across multiple devices while fully overlapping the communication of key-value blocks with the computation of blockwise attention. Our approach enables training and inference of sequences that are up to device count times longer than those achievable by prior memory-efficient Transformers, without resorting to approximations or incurring additional communication and computation overheads. Extensive experiments on language modeling and reinforcement learning tasks demonstrate the effectiveness of our approach in allowing millions of tokens context size and improving performance.
DarcStar commentary
Our perspective on this work and how it informs our research — not part of the original paper.
Of all the ways people have tried to make attention scale, ring attention is the one we've spent the most time building ourselves — we wrote our own distributed implementation, cross-chunk combine and all — so we'll be plain about what it gets right and where the hard part actually is.
Ring attention goes after the wall that really caps context: memory. Shard the sequence across a ring of devices, stream the key-value blocks around the ring, and overlap that passing with the blockwise compute so the communication hides behind the math. Do it right and the context you can hold grows with the number of devices — exactly, with no approximation and no whole-sequence bottleneck living on any one card. The quiet elegance is that online softmax is associative: you can combine partial attention computed on separate shards without ever assembling the whole thing in one place. The same math runs on one GPU or a whole ring without changing a line. That's the idea worth keeping.
Here's the part the diagram makes look easy. The correctness lives in the seams — the causal mask and the softmax denominator have to be combined across chunks exactly, and a surprising number of ring implementations get that cross-chunk combine subtly wrong: the loss curve looks fine while the attention is quietly incorrect. And there's a tempting shortcut — just gather every block onto every device — that throws away the entire memory win the moment you reach for it. If your “ring” is doing an all-gather, it isn't the thing this paper is about.
The other tooth is what bites at real scale: past a point, the ring's communication, not its compute, is what stalls the machine. The bottleneck at millions of tokens isn't FLOPs — it's keeping the pipes full and the topology honest. Solve compute with one lever, memory with the ring, communication with another still; drop any one and you hit a different wall.
So this one sits on our shelf as external work we've spent real time inside — the memory lever done cleanly, and a genuinely exact one at that. Its ambition — near-infinite context, the whole video or episode or corpus in view at once — is the right target. We just think you earn it by keeping the combine provably exact, refusing the all-gather fallback, making the ring topology-aware, and — the ground we care about most — keeping it correct when the attention underneath is also sparse. That last combination is still open, and it's where the interesting work is.