The memory bill for a million tokens
Start with why a single accelerator loses. Take a mid-sized model: d_model = 4096, 32 layers, bf16. One layer’s K and V for a context of N tokens occupy 2 · N · d · 2 bytes. At N = 1,048,576 that is 2 × 1.05e6 × 4096 × 2 ≈ 17 GB per layer — over 500 GB across 32 layers before you count queries, the residual stream, or a single weight.
Grouped-query attention softens this: with 8 KV heads out of 32, the KV width drops from 4096 to 1024 and the bill falls 4×, to roughly 135 GB. Better, still hopeless on one 80 GB device. And note what did not save you: FlashAttention. It removes the [N, N] score matrix, which is the N^2 term, but the N · d activations are untouched. The N^2 problem is a kernel problem; the N · d problem is a capacity problem, and capacity problems are solved by adding devices.
Where context parallelism sits
Every large-model deployment already splits work along several axes, and it helps to see which dimension each one cuts:
| Axis | Splits | Communication |
|---|---|---|
| Data (DP/FSDP) | the batch | all-reduce of gradients |
| Tensor (TP) | hidden dim / heads | all-reduce twice per layer |
| Pipeline (PP) | layers | point-to-point activations |
| Context (CP) | the token axis | K/V blocks around a ring |
Ring Attention is the leading implementation of that last row. The reason it needs its own axis is that attention is the only layer that mixes tokens. Split the sequence and every other layer — embeddings, LayerNorm, the MLP, the residual add — keeps working untouched, because each acts on one token at a time. Attention alone requires every query to meet every key, and that single coupling is what the ring exists to service — everything else about sequence sharding is embarrassingly parallel.
The mechanism, compressed
So that this piece stands alone: cut the sequence into P contiguous blocks of b = N/P tokens, one per device. Device i holds Q_i, K_i, V_i, each [b, d]. Queries never move. Over P steps, the K/V blocks hop one seat forward around the ring; at each step a device computes the [b, b] partial attention between its own Q_i and whichever K/V block is currently visiting, then merges it into a running (m, l, o) accumulator — running max, running denominator, running unnormalized output.
The merge is exact because softmax is invariant to the constant you subtract, so an arriving block with a larger max can retroactively rescale everything already accumulated by α = exp(m_old − m_new). That is the whole trick, and it is the same accumulator FlashAttention uses on-chip — lifted from one GPU’s memory hierarchy to a pod’s interconnect. Ring Attention is exact: no approximation, no quality loss, only a different cost structure.
Ring versus all-to-all: counting bytes
The ring is not the only way to shard the token axis. The main alternative, Ulysses-style head parallelism, keeps the sequence split for every layer except attention, then does an all-to-all that re-shards from “by token” to “by head” just before attention and back again after, so each device computes a few complete heads over the whole sequence with an unmodified kernel.
The byte counts differ sharply. Per device, a ring ships one K/V block per step for P-1 steps: 2 · d_kv · (N − b) elements, which grows with total context. Ulysses moves four tensors (Q, K, V, O) once each through an all-to-all, roughly 4 · b · d elements — a factor of roughly P/2 less traffic. The catch is that Ulysses cannot scale past the head count (P ≤ H), it needs a genuine all-to-all rather than nearest-neighbour links, and it collides with tensor parallelism, which already wants the head axis. The ring moves more bytes but scales to arbitrary P, uses only neighbour links, and overlaps cleanly. Production stacks often do both: Ulysses inside a node, a ring across nodes.
A bandwidth budget that decides everything
Whether the ring is free or fatal is one arithmetic comparison per step: transfer time versus matmul time. Take N = 1M, P = 64, so b = 16,384, with d = 4096 and bf16.
Compute per step = 4 · b^2 · d (QK^T and PV)
= 4 · 16384^2 · 4096 ≈ 4.4e12 FLOP
÷ 400 TFLOP/s ≈ 11 ms
Transfer per step = 2 · b · d_kv · 2 bytes
MHA (d_kv=4096): 268 MB → NVLink 900 GB/s = 0.30 ms (hidden, 37× margin)
→ IB 50 GB/s = 5.4 ms (marginal, 2×)
GQA (d_kv=1024): 67 MB → IB 50 GB/s = 1.3 ms (hidden, 8×)
Three lessons fall straight out. Inside a node the ring is essentially free. Across a slow fabric it is borderline — and once a causal mask halves the useful compute per step while leaving the transfer unchanged, a 2× margin becomes a stall. And GQA is a communication optimization as much as a memory one: shrinking d_kv shrinks exactly the tensor the ring puts on the wire. The general ratio is compute/comm = O(b), so when the ring stalls, the fix is a bigger block — fewer, larger devices’ worth of tokens each — not a faster network.
Causal masking wastes half the ring
The clean picture assumes every query block attends to every key block. With an autoregressive mask, query block i only needs key blocks j ≤ i, so device i does i+1 block-pairs of work while device 0 does one. Total useful work is P(P+1)/2 out of P^2, but the step is gated by the busiest device, so utilization tends to (P+1)/2P → 50%. Half the pod idles while device P-1 grinds.
Zigzag (or striped) ring attention fixes this by splitting into 2P half-blocks and giving device i both chunk i and chunk 2P−1−i. Chunk c carries roughly c+1 units of work, so device i carries (i+1) + (2P−i) = 2P+1 — identical for every i. Pairing an early chunk with a late one makes each device’s load the same by construction. The cost is that a device’s tokens are no longer contiguous, which you must remember when assigning positions.
The ring is a prefill algorithm
A point that catches teams out: Ring Attention pays for itself during prefill, not decode. Prefill has b queries meeting b keys, so compute is O(b^2 d) against O(b d) of traffic — the O(b) ratio that lets transfers hide. During autoregressive decode there is one query token per sequence. Compute collapses to O(b d), the same order as the transfer, and the ratio becomes O(1). Rotating gigabytes of KV cache to serve one token is pure waste.
The right decode pattern inverts the movement: keep the KV cache sharded and stationary, broadcast the single query vector to all P devices, let each compute a partial (m, l, o) against its own shard, and combine the P triples with one small reduction — the same rescale-and-add merge, applied once at the end. You move a query vector and P partial states instead of the entire cache. Same math, opposite dataflow.