Ring Attention: what the near-infinite context paper actually says

Reading Liu et al. (UC Berkeley, 2023) after a 128K-context training job that wouldn't fit on any number of GPUs we had.

The request came from the model team: they needed to train on 128K-token documents. Legal contracts, full codebases, multi-day conversation transcripts. The context had to be real — truncating to 32K changed the model's behavior on the tasks they cared about.

The math killed every approach we tried first. FlashAttention gets you from O(N²) memory to O(N) — but O(N) at 128K tokens, 80 hidden layers, BF16, is still 7.5 GB per layer per device just for the activations. At 80 layers, a single model replica won't fit in 640 GB of aggregate device memory without a lot of engineering. And that's before optimizer state, gradients, or the KV cache.

Tensor parallelism (Megatron-style) solves the weight matrix problem but doesn't help with sequence-length-scaled activations. ZeRO partitions optimizer state but the sequence dimension stays local. The problem is that attention, by design, requires every query to see every key-value pair — and those KV pairs scale with the sequence you're trying to attend to.

Ring Attention — "Ring Attention with Blockwise Transformers for Near-Infinite Context," Liu, Yan, Stone, Moritz, Liang, and Abbeel, UC Berkeley 2023 — addresses this by distributing the sequence dimension itself across devices. Each device owns a chunk of tokens and computes attention over those queries while rotating the keys and values through a ring of devices until every query has seen every KV pair. The core claim: you can train transformers on sequences that scale with the number of devices, not the memory of any single device.

What actually limits sequence length before Ring Attention

To understand the paper's contribution, you need to start with where FlashAttention leaves off.

FlashAttention (Dao et al., 2022) eliminates the O(N²) memory requirement for the attention score matrix by tiling the computation. Rather than materializing the full N×N attention score matrix in HBM, it processes attention in blocks that fit in SRAM, using online softmax to accumulate results without storing intermediates. Memory goes from O(N²) to O(N) for attention. This is the key result — but O(N) still grows linearly with sequence length, and that linear growth eventually hits the device memory ceiling.

At 128K tokens, with hidden dimension 4096 and 32 attention heads:

The Q, K, V projections are each [128K, 4096] — about 1 GB per matrix at BF16
Intermediate activations through each transformer block: another 4-6 GB
With 40 layers: you're looking at 200-250 GB just for activations with activation checkpointing disabled

FlashAttention helps enormously relative to standard attention but doesn't change the fundamental scaling: one device still needs to hold all the KV pairs for one sequence, because the tiling is within a device. Each query still attends to all N keys, and all N keys need to be somewhere the device can read from. At 128K tokens, N = 128,000. On an 80 GB A100, that's the constraint.

The solution space before Ring Attention:

Sequence truncation: changes the model's behavior on long-context tasks, often unacceptably
Gradient checkpointing: trades compute for memory by recomputing activations; helps with hidden layer activations but doesn't reduce KV memory proportionally
Offloading to CPU: latency is too high for training (PCIe bandwidth ~32 GB/s vs NVLink ~600 GB/s)
Approximate attention (Longformer, BigBird): sparse patterns reduce memory but change what the model can attend to

None of these preserve exact full-sequence attention at 128K while fitting a training run on realistic hardware. That's the problem Ring Attention is solving.

The ring mechanism, step by step

Ring Attention distributes the sequence across P devices arranged in a logical ring. Device 0 holds tokens 0 to N/P−1, device 1 holds tokens N/P to 2N/P−1, and so on. Each device has the Q, K, and V vectors for its chunk of the sequence — computed from its local embeddings — and the model's attention weights.

A single forward pass through one attention layer requires N rounds:

Round 0: Each device computes the local attention contribution — how much its local queries attend to its local keys and values. This is standard FlashAttention within the device, producing a partial attention output for each local query position.

Round 1: Each device sends its K and V blocks to the next device in the ring (device i → device i+1, wrapping around). While that communication is in flight, each device has already started processing its local Q against the incoming K/V from the previous round. When the K/V arrives, computation resumes. The device updates its running partial output using the online softmax accumulation from FlashAttention — the same running maximum m and running sum l trick, now applied across rounds instead of across blocks within a device.

Rounds 1 through P−1: This pattern repeats. Each device sends its K/V to the next while computing with the K/V it just received. After P rounds, every device's queries have attended to every other device's keys and values. The running partial outputs are correct: mathematically identical to what you'd get from running full-sequence attention on a single device.

The key: communication and computation overlap. While device i computes attention using K/V from device i−1, the K/V from device i−1 is simultaneously being sent to device i−2. The ring structure ensures every device is either computing or sending (usually both) at every step. If the communication latency is less than the compute time per round, the communication cost is fully hidden.

The paper shows the communication complexity: each device sends N/P × d_k × 2 (keys and values) per round, for P rounds, totaling N × d_k × 2 bytes per head per device over the full pass. This is equivalent to one all-gather of the full KV sequence — but structured so it can be overlapped with computation.

What the online softmax accumulation actually does across rounds

The online softmax trick from FlashAttention-1 is load-bearing here, and it's worth understanding what's happening.

Standard softmax for one query row:

attention_weights = softmax(QKᵀ / sqrt(d_k))
output = attention_weights × V

You need all N keys to compute the denominator of softmax. For a single query position attending to 128K keys split across 8 devices, you're receiving those keys in 8 batches of 16K.

After processing the first device's K/V (16K keys), you have:

m_1 = max of the first 16K attention logits
l_1 = sum of exp(logit - m_1) over first 16K keys
o_1 = weighted sum of first 16K values (unnormalized)

After receiving the second device's K/V, you update:

m_2 = max(m_1, new_max)
l_2 = l_1 × exp(m_1 - m_2) + l_new × exp(new_max - m_2)
o_2 = (o_1 × l_1 × exp(m_1 - m_2) + o_new) / l_2

This rescaling is exact — no approximation. After all P rounds, o_P / l_P gives the exact attention output for that query position, identical to computing full-sequence attention at once. The running statistics are small (two scalars per query head), so they add negligible memory.

This is the mathematical core that makes Ring Attention work: FlashAttention's online softmax is associative and correct across any partition of the key-value sequence.

The causal masking complication

Causal attention (autoregressive generation, decoder-only models) adds a load imbalance problem that the paper addresses but is worth understanding clearly.

In causal attention, query position i can only attend to key positions ≤ i. This means:

Device 0 (tokens 0 to N/P−1): these early tokens attend to a small prefix of the sequence. Device 0 does O(N/P)² work on its local block, then very little work in subsequent rounds (its later-round K/V blocks are all masked out for its queries).
Device P−1 (tokens N−N/P to N−1): these late tokens attend to the full sequence. Device P−1 does O(N/P)² work on each local block it receives, every round.

Without correction, device 0 finishes much earlier than device P−1 and sits idle. Pipeline efficiency drops.

The solution is to reorder the ring schedule. Rather than each device always sending to device i+1, you pair devices so that each pair of devices always has comparable amounts of useful work to do. Device 0 and device P−1 are paired so that when device 0 is processing device P−1's K/V (almost entirely masked for device 0's queries), device P−1 is processing device 0's K/V (mostly unmasked for device P−1's queries). The work within each pair is complementary: one device's batch is mostly masked exactly when the other's is mostly unmasked.

The paper demonstrates that this reordering eliminates the load imbalance for causal attention, achieving roughly uniform device utilization throughout training. Without this, you'd see idle devices that reduce effective throughput proportionally to the imbalance.

Competing approaches: DeepSpeed Ulysses and Megatron sequence parallelism

Ring Attention isn't the only way to distribute attention across the sequence dimension. The comparison with other approaches matters for choosing the right one.

Megatron-style sequence parallelism (Korthikanti et al., 2023) adds a sequence-parallel region outside the attention computation itself. Activations like LayerNorm outputs are partitioned across devices by sequence position, using all-gather before attention and reduce-scatter after. Attention itself still runs on the full local sequence — it doesn't distribute the KV pairs. This helps with activation memory between transformer blocks but doesn't help when the attention KV for a single sequence is too large for one device. The overhead is two extra all-gathers per transformer block.

DeepSpeed Ulysses (Brandon et al., 2023) takes a different approach: before computing attention, it transposes the (sequence, head) dimensions using an all-to-all operation, so each device works on all sequence positions for a subset of heads. After attention, another all-to-all transposes back. Each device computes full-sequence attention for its heads, so no KV rotation is needed — but each device still needs to hold all N positions for its heads, which is the memory constraint that Ring Attention eliminates. Ulysses reduces communication volume relative to ring approaches (two all-to-alls vs ring communication), but doesn't handle the case where N × d_k per head exceeds device memory.

Ring Flash Attention (Li et al., 2023) is an optimized re-implementation of the ring concept that addresses some practical inefficiencies in the original paper's implementation — better overlap scheduling, more careful handling of causal masking, lower memory overhead for the intermediate statistics. For production use, this is the implementation to consider over the original paper's reference code.

The practical selection:

Megatron sequence parallelism: when activations between blocks are the bottleneck, but per-head KV fits on one device
DeepSpeed Ulysses: when you have many attention heads and the all-to-all communication is faster than ring communication for your topology; head count must be divisible by device count
Ring Attention: when K/V for a single sequence at a single head doesn't fit on one device — true long-context training

What the paper's numbers actually show

The paper trains a transformer on 1 million token sequences using 8 TPUv4 pods (256 TPUs total), demonstrating that the approach scales with device count rather than device memory. At 1M tokens, standard FlashAttention on a single device would require roughly 75 GB just for the Q, K, V activations — Ring Attention distributes this so each device holds 1/256 of the sequence.

The more relevant numbers for GPU training:

The communication volume per device is proportional to sequence length × head dimension, divided by device count. For 128K sequences on 8 A100s with 32 heads and d_k=128: each device sends about 128K × 128 × 2 bytes / 8 = 4 MB per round, for 8 rounds = 32 MB total per layer.
NVLink bandwidth (within node): ~600 GB/s bidirectional. Sending 4 MB over NVLink takes about 6.7 microseconds.
Time to compute local attention over N/P = 16K tokens per round on one A100: depends on batch size, but easily 2-10ms for a single sequence at BF16.
Communication is fully hidden behind computation in most configurations when using intra-node NVLink.

The paper reports near-linear scaling efficiency (>90%) when devices are within-node and using high-bandwidth interconnect. Efficiency drops when crossing node boundaries over InfiniBand.

Production tradeoffs no one mentions in the benchmark post

Communication-computation balance is fragile. Ring Attention's efficiency relies on the compute time per round being longer than the communication time. Increase your hardware's compute throughput (H100 vs A100: ~2× MFU), and the compute-communication ratio shifts. What hid communication cleanly on A100s may not on H100s for the same model and sequence length. Before assuming you get the paper's scaling efficiency numbers, profile your specific hardware.

The implementation complexity is substantial. Ring Attention requires custom attention kernels that know about the ring structure, handle the running softmax accumulation across rounds, implement the causal reordering for decoder models, and correctly schedule the overlapped communication. The original paper's code is a research implementation. Production use requires either Ring Flash Attention, a framework integration (EasyContext, JAX-based implementations for TPU), or significant engineering work. Teams that try to bolt ring communication onto standard PyTorch attention kernels get correctness bugs in the softmax accumulation that are very hard to detect — the outputs look close to correct but aren't exact.

Sequence parallel degree constraints. Number of devices must divide evenly into sequence length, and sequence length per device must be large enough to keep the local FlashAttention computation efficient. If N/P drops below ~512 tokens per device, the per-round compute time shrinks to the point where communication overhead is no longer hidden, and effective throughput degrades despite having more hardware. For short sequences, ring attention is worse than no ring.

It doesn't help with model parallelism integration. Ring Attention partitions the sequence dimension. Tensor parallelism (Megatron) partitions the model dimension. These interact: if you're doing both, you need to manage device groups carefully. Each tensor-parallel group typically needs to independently do its ring communication, which means the ring and tensor parallel groups must be orthogonal. Combined 3D parallelism (tensor × pipeline × data) already coordinates multiple device groups; adding sequence parallelism as a fourth dimension requires framework support that isn't universally available.

Gradient correctness requires explicit verification. The forward pass correctness follows from the online softmax math. The backward pass through ring attention is more complex: gradients with respect to Q, K, and V flow back through the ring in the reverse direction. Bugs in the backward pass don't always surface as obvious divergence — they can cause subtle gradient estimates that produce training that looks stable but converges to a worse solution. If you're using an unfamiliar implementation, validate that the gradients match a single-device reference implementation before training at scale.

Your model must be designed for long context, not just trained on it. Rotary position embeddings (RoPE) work reasonably well at modest context extensions, but standard RoPE trained at 4K doesn't extrapolate cleanly to 128K. YaRN, ALiBi, or extended RoPE training are typically needed. Ring Attention solves the hardware constraint of fitting long contexts; it doesn't solve the model's ability to use that context. Teams that enable ring attention and train at 128K without adjusting the position encoding scheme find their models don't actually use the long context window they paid for.

Failure modes in practice

Silent precision drift in the softmax accumulation. The online softmax update involves multiplying accumulated outputs by exp(m_old - m_new) to rescale when a new maximum is found. In BF16, this rescaling is imprecise when m_old - m_new is large — which happens when attention logits have high variance, common in models without QK normalization. The error is small per round but accumulates across P rounds. For 32 devices, 32 rounds of rescaling with small BF16 errors can produce outputs that diverge from the FP32 reference in the last few bits. Usually invisible in loss curves; occasionally catastrophic if those positions are critical to the training signal. Add a periodic check that compares ring attention output against single-device reference on a representative input before committing to a long training run.

Load imbalance on heterogeneous hardware. The ring is only as fast as the slowest device. In cloud environments with spot instances or heterogeneous GPU generations (one P4d.24xlarge with an A100-80GB alongside an older A100-40GB node), one slow device stalls the ring every round. Unlike pipeline parallelism where micro-batching absorbs some imbalance, ring communication is synchronous at every attention layer. If hardware is heterogeneous, profile first.

Checkpoint format incompatibility. Standard model checkpoints store full-sequence activations. Ring Attention training produces checkpoints where attention state is distributed across devices. Resuming a ring attention training job with a different number of devices — common when a training run fails mid-way and you restart with different GPU allocation — requires resharding the checkpoint. If you don't have tooling for this, a training interruption means starting from the last checkpoint that predates ring attention configuration changes, which may be far earlier than the failure point.

When NOT to use Ring Attention

When your sequence fits on one device with FlashAttention. If a single A100 can hold the K/V for your maximum sequence with FlashAttention-2 enabled, don't add ring complexity. FlashAttention is already efficient. Ring Attention adds communication overhead (even when hidden), implementation complexity, and a new class of bugs. The simpler system wins unless you actually need the distributed memory.

For inference serving. Ring Attention is designed for training, where you process full sequences forward and backward in batches. For autoregressive inference, you're generating one token at a time, and the KV cache grows incrementally. Ring Attention's communication pattern doesn't map cleanly to incremental KV cache growth — you'd need to rotate the entire KV cache through the ring on each token generation step, which is impractical. Use PagedAttention or chunked prefill approaches for serving long-context models.

When crossing node boundaries without InfiniBand. The ring communication between devices must be fast enough to hide behind computation. Within a node over NVLink (~600 GB/s), this is achievable. Across nodes over 100 GbE Ethernet (~12 GB/s), you get one-month of communication overhead per layer that you can't hide. Ring Attention across node boundaries requires at least 200 GbE or InfiniBand (HDR/NDR) to be viable. Before provisioning a multi-node ring attention training job, calculate whether your network bandwidth can support it for your sequence length.

When the overhead of causal reordering exceeds its benefit. For encoder-only models (classification, embeddings), bidirectional attention means the ring has symmetric load and no reordering is needed. For decoder-only models, the causal reordering in the paper is necessary but adds scheduling complexity and some implementation overhead. For short sequences at low device counts (4 or fewer), the simpler approach of just letting some devices idle during masked blocks might have higher actual throughput than a carefully load-balanced ring — measure it.

When you're using a custom attention variant. Sliding window attention, cross-attention between sequences of different lengths, attention with learned sparsity patterns — these require custom kernels. Ring Attention is already complex to implement correctly for standard causal attention; custom attention variants on top of ring communication is a significant engineering project. If your architecture uses non-standard attention, assume that ring attention support requires building it from scratch.

What the paper actually gives you

Ring Attention is the training-time solution to a problem that FlashAttention solves at the intra-device level. FlashAttention eliminated the N² activation memory for attention by recomputing on-chip; Ring Attention eliminates the N × d_k KV memory constraint by distributing across devices. Together, they make exact full-sequence attention at very long contexts feasible.

The IO-awareness insight from FlashAttention generalizes here: the bottleneck isn't the computation, it's data movement. Ring Attention hides inter-device data movement by overlapping it with computation, just as FlashAttention hides HBM reads by computing in SRAM. The pattern is the same; the level of the memory hierarchy is different.

For the 128K training job that started this: we combined Ring Flash Attention (Li et al., 2023) with per-device FlashAttention-2, tensor parallelism degree 4 within each 4-GPU node, and ring size 8 across 2 nodes. Sequence length per device was 16K tokens — enough to keep local FlashAttention efficient — with NVLink within nodes and InfiniBand across. Communication was fully hidden behind computation in profiling. Training stability required extending RoPE using YaRN; without it, the model's perplexity at token positions beyond 32K was flat despite being trained on 128K sequences.

The model now uses the full 128K context on the evaluation tasks. But that last part — actually using the context — required more than just fitting the hardware. The hardware constraint was Ring Attention. The modeling constraint was position encoding. Both had to be solved.

Ring Attention with Blockwise Transformers for Near-Infinite Context — Liu, Yan, Stone, Moritz, Liang, Abbeel. ArXiv 2023.
Ring Flash Attention — Li et al. GitHub, 2023.
Reducing Activation Recomputation in Large Transformer Models — Korthikanti et al., NVIDIA. MLSys 2023.
DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models — Brandon et al., Microsoft. ArXiv 2023.