Ring Attention: Blockwise Transformers for Near-Infinite Context

Transformers win on quality, then trip on context. Self-attention is an $n$-to-$n$ interaction: even after you stop materializing the full $s \times s$ score matrix, each layer still wants the prior layer’s full-length activations. Put a hundred million tokens through a modest hidden size and you need more HBM than a single GPU or TPU slice has — not because the math is approximate, but because one device is being asked to own the whole sequence.

Ring Attention keeps the original Transformer equations. It reorganizes where the work lives. The sequence is cut into blocks and spread across a ring of hosts. Each host holds one query block. Key/value blocks circulate around the ring while each host runs blockwise attention and feedforward. When the block is large enough that compute hides the transfer, you get longer context without approximate attention and without an extra communication tax.

What this paper explains

Ring Attention (Liu, Zaharia, Abbeel) is a systems recipe for exact long-context Transformers:

1. Start from blockwise parallel transformers (BPT). Prior work already computes attention and FFN block-by-block so you do not materialize the full score matrix. That cuts activation pressure, but the layer output is still length $s$ on the device that owns it.
2. Shard the sequence across hosts. Host $i$ owns query block $i$ (the outer loop of blockwise attention) and the FFN for that block — no all-to-all required for those steps.
3. Circulate KV blocks on a ring. The inner loop needs every KV block. Instead of gathering them (delay + memory), each host sends its current KV block to the next host and receives from the previous host, overlapping transfer with blockwise compute.
4. Scale context with device count. The abstract’s claim: sequences up to device-count times longer than prior memory-efficient Transformers, without approximation or added communication/compute overheads when the arithmetic-intensity condition holds.

This is not another KV-cache eviction story (StreamingLLM sinks, H₂O heavy hitters) and not a serving batching story (Orca, DistServe, Sarathi). It is sequence-parallel exact attention for training and inference.

Prior limits

• Vanilla attention stores $O(s^{2})$ scores — dead on arrival for million-token contexts.
• Memory-efficient attention (FlashAttention-style online softmax / tiling) and BPT remove the score-matrix wall and shrink FFN peaks, but Table 1 in the paper still lists prior SOTA activations as $2bsh$ — linear in full sequence length $s$.
• Why that still fails. The next layer’s attention needs every prior-layer token. The paper’s concrete illustration: ~100M tokens, batch 1, hidden 1024 → >1000GB of activation footprint for storing layer outputs — while modern HBM is typically <100GB per device.
• Naive multi-host fetch of remote KV blocks either stalls on receives or accumulates KV copies until memory savings vanish. Sequence-parallel attention in prior ring topologies still left non-overlapped communication that BPT-scale memory did not fix.

The mechanism

Blockwise outer/inner loops. Write queries, keys, values as sequences of blocks. For a fixed query block $Q_{i}$, you walk KV blocks $K_{j},V_{j}$, accumulate online-softmax statistics, then run the FFN on that query block’s attention output. Host $i$ runs the outer iteration for $Q_{i}$.

Permutation invariance of the KV walk. For a fixed query block, the order of KV blocks does not matter as long as you combine block statistics correctly when rescaling. That is what licenses a ring schedule instead of a gather.

Ring schedule. Hosts $1,\ldots,N$ form a circle. While host $i$ attends its query block against the KV block it currently holds, it concurrently:

• sends that KV block to host $(i+1)$,
• receives the next KV block from host $(i-1)$.

After $N$ steps, every query block has seen every KV block — exact attention — and no host ever needed the full sequence in HBM.

Arithmetic intensity (overlap condition). Let $F$ be device FLOPs/s and $B$ host-to-neighbor bandwidth. For block size $c$ and head/hidden width $d$, the paper’s simplified attention-only bound is

If $c$ clears that bar, KV transfer hides under compute (forward and backward). Table 2 in the paper lists device-specific $F/B$ thresholds.

Memory on one host. Store: current query block, current KV, incoming KV buffer, attention/FFN output for the query block → six blocks → $6bch$ bytes (bfloat16 accounting in Table 1), with $c$ independent of $s$. Contrast prior BPT total $2bsh$.

A one-line mental model: sequence length $s \approx N \cdot c$ across $N$ hosts, while each host’s activation ceiling tracks $c$, not $s$ — provided the ring is fed and $c \ge F/B$.

Algorithm / figure walkthrough

1. Partition the length-$s$ sequence into $N$ blocks of size $c$ (plus the usual head/batch axes).
2. Place block $i$ on host $i$; form queries/keys/values locally for that block.
3. For step $t = 0,\ldots,N-1$: each host runs blockwise attention (and eventually FFN) between its fixed query block and the KV block currently in hand, while ppermute-style neighbor exchange rotates KV one hop along the ring.
4. Combine online-softmax numerators/denominators across KV visits so the result matches full attention.
5. Stack layers the same way — the architecture is unchanged; only the execution topology is a ring.

The opening visual is the same arc: ENTER with one device drowning in length-$s$ activations; RUN with query cubes fixed and KV tokens circulating; LEAVE with context that grows when you add hosts.

What to notice when reading

• Exactness. Softmax is still global over the sequence; they are not sparsifying or compressing attention (contrast Infini-attention / sinks).
• Independence from tensor parallel. Context-length tables use FSDP-style sharding for the model; Ring Attention is an extra axis for sequence.
• Block size is a systems dial. Too small → communication shows; large enough → overlap. Memory is $O(c)$, so you still pick $c$ deliberately.
• Training-first narrative, inference-ready. Autoregressive decode is lighter (few queries), but Appendix C sketches circulating the KV cache on a ring for larger served contexts.
• Jax ppermute. Appendix A’s sketch matches the story: collective permute around axis_name, not a full all-gather of KV.

Results and evidence

Numbers below are from the paper’s tables/figures (arXiv:2310.01889) — not re-measured here.

Table 1 — activation sizes (bytes/layer, bfloat16). Ring Attention’s total is $6bch$ vs memory-efficient attn+FFN $2bsh$. The decisive difference is $c$ vs $s$.

Context-length table (paper §5.1; tokens $\times 10^{3}$). Ring Attention vs prior memory-efficient attn+FFN, same FSDP-style setup:

The “vs SOTA” column tracks device count on the GPU rows (8× / 32×) — matching the abstract’s “device count times longer” slogan when the ring is fully used for sequence.

MFU (Table 4 / §5.2). Ring Attention pushes much longer contexts (e.g. 7B on 8× A100: prior 32K vs Ring 256K; 13B on 32× A100: prior 64K vs Ring 2048K) while the authors report MFU staying in line with BPT expectations despite more attention FLOPs.

In-context RL (Table 5). On ExoRL, Action Transformers with Ring Attention train 128 trajectories where BPT hits OOM at that width; average return 113.66 vs BPT’s 111.13 at 32 trajectories (authors’ runs; BC/DT baselines lower).

LLM line retrieval (§5.4 / Figure 3). LLaMA-13B finetuned with Ring Attention to 512K context on 32× A100 (ShareGPT-derived data; compute-budget limited vs their “millions of tokens” capability claim) maintains high retrieval accuracy out to long contexts where 16K–100K chat baselines cannot continue.

Appendix C inference sketch. Example: LLaMA-7B on 32× TPUv5e — conventional head-parallel serve ~256K KV; Ring Attention can circulate KV for 32× larger context when bandwidth/FLOPs satisfy the overlap inequality (they illustrate $B/F \approx 2.4$ under an assumed 40% MFU).

Limitations

• Needs a fast ring. Overlap assumes neighbor bandwidth high enough relative to FLOPs; weak interconnects reintroduce stalls.
• Block-size engineering. Memory is $O(c)$; picking $c$ is not free — too small fails overlap, too large pressures HBM.
• FLOPs still grow with $s$. Appendix D shows per-dataset FLOPs ratios vs 4K context; small models pay larger relative attention cost when context balloons.
• Finetune context capped in the LLM experiment. 512K is a budget choice, not the method’s ceiling; do not read it as a hard max.
• Complementary, not a replacement for FlashAttention kernels, FSDP/tensor parallel, or serving schedulers — you still compose those tools.
• Venue. Treat claims as the arXiv tech report’s reported measurements; no conference camera-ready is asserted in the abs comments.

How to read the paper

1. Figure 1 — max context under end-to-end training on TPUv4-1024 (motivation).
2. §2 — why $2bsh$ still loses to HBM.
3. Figure 2 + §3 — the ring cartoon and the $c \ge F/B$ derivation.
4. Table 1 — $6bch$ vs $2bsh$.
5. §5.1 context table — device-count scaling in practice.
6. Table 4 / Table 5 / Figure 3 — MFU, RL, line retrieval.
7. Appendix A + C — Jax ppermute sketch and inference KV circulation.

Knowledge check

Keep reading / Sources

• Paper: Ring Attention with Blockwise Transformers for Near-Infinite Context (arXiv:2310.01889)
• Code named in-paper: llm_large_context
• Prior blockwise parallel transformers (cited as BPT in §2–3)
• Contrast in Daily: FlashAttention, StreamingLLM, H₂O
• Serving (different problem): DistServe, Sarathi-Serve
• Sharding cousin: ZeRO