Why architecture matters here
Consider the alternatives, because they clarify what ring attention buys. You could use an approximation — sliding-window attention, sparse patterns, or a linear-attention variant. These reduce the asymptotic cost genuinely, but they change the function being computed. A model trained with full attention does not behave identically when you swap in a windowed approximation at inference, and a model trained with the approximation carries its limitations permanently. You are making a quality decision disguised as an infrastructure decision, and you will not discover the cost until some evaluation depends on a long-range dependency the window dropped.
Or you could gather: all-gather every device's keys and values so each device has the full KV tensor, then compute locally. This is exact and trivially simple, and it defeats the entire purpose — after the gather, every device holds O(S) worth of KV again, which is the wall you were climbing. It also concentrates the same total bytes into a single burst, giving you peak bandwidth demand and no overlap opportunity. Ring attention moves the identical volume spread evenly across N steps, each hop overlapping real arithmetic.
The architectural payoff is that context length becomes a scaling parameter rather than a hardware constraint. Want to double the context? Double the ring — per-device memory stays flat because each device still holds S/N tokens. That is a qualitatively different position from 'wait for next year's chip', and it composes cleanly with the other parallelisms, since sequence, model, and batch sharding partition orthogonal axes.
The cost you accept is that your attention layer now has a network dependency in its inner loop. A single slow device does not slow only itself; it slows the entire ring, because every other device eventually waits on the block it should have forwarded. The blast radius of one degraded link is the whole collective, and that coupling is the price of exactness at scale.
The architecture: every piece explained
Start with the sharding. The sequence of length S is split into N contiguous blocks, one per device, so device i owns tokens [i·S/N, (i+1)·S/N). Critically, device i owns both the queries and the keys/values for its slice. This is what makes the memory arithmetic work: nothing is replicated, so per-device memory for the KV cache is O(S/N · d) where d is the head dimension. The queries stay pinned to their device for the entire operation — only the KV blocks move. That asymmetry is deliberate, and it matters, because it means the output shard also stays pinned: device i produces the output rows for its own tokens and never needs to send them anywhere.
The ring topology is a logical construct, not necessarily a physical one. Device i sends to device (i+1) mod N and receives from (i-1) mod N. On hardware with a physical ring or mesh — NVLink within a node, a torus across nodes — the logical ring maps onto physical links so every hop is a single-hop transfer at full bandwidth. Without that structure some hops traverse multiple links, and because the ring is synchronous your effective hop latency becomes the worst hop rather than the average. Topology-aware rank assignment is not a micro-optimization; it is the difference between overlapping and stalling.
The online softmax accumulator is the mathematical heart. A standard softmax needs the global maximum and the global sum of exponentials before it can normalize anything, which appears to require seeing all the scores at once. The online formulation sidesteps this by carrying three running values per query row: the running maximum m, the running sum of exponentials l, and the running unnormalized output o. When a new block arrives with its own local maximum and local sum, you compute the new global maximum, rescale the existing accumulator by exp(m_old - m_new), rescale the incoming block by exp(m_local - m_new), and add. The rescaling factor corrects for the fact that the earlier partial results were normalized against a maximum that has now been superseded. The result after absorbing all N blocks is bit-for-bit the same as a single-shot softmax, up to floating-point associativity.
The communication layer is where the performance lives. Each step issues an asynchronous send of the current KV block to the next device and an asynchronous receive of the next block from the previous device, then immediately begins computing attention against the block already in hand. The send and receive complete in the background while the arithmetic runs. This requires double-buffering — you need a buffer holding the block you are computing on and a separate buffer receiving the incoming block, because you cannot overwrite data you are still reading. Two KV block buffers per device is the standard cost, which nudges per-device memory to O(2 · S/N · d) but keeps the asymptotic win intact.
Causal masking adds a wrinkle, because it is where naive implementations leave half their throughput on the floor. When device 2 receives a KV block from device 3 — tokens strictly later than device 2's own queries — every score is masked out and the block-step is wasted work. A causal-aware ring skips the roughly half of (query block, KV block) pairs that are fully masked, but that creates load imbalance: device 0 has almost nothing to do while device N-1 does nearly all the work. The standard fix is zigzag assignment — give each device two non-adjacent slices, one from the front of the sequence and one from the back, so every device carries a balanced mix of cheap early and expensive late tokens.
End-to-end flow
Walk a single forward pass through a 4-device ring on a 400,000-token sequence. Each device holds 100,000 tokens. At initialization, device i holds Q_i, K_i, V_i, and its accumulator (m, l, o) is initialized to negative infinity, zero, and zero respectively — the identity element for the online softmax merge.
Step 0: every device computes attention between its local Q_i and its local K_i, V_i — the diagonal block, which needs no communication but uses a triangular mask, since device i's query at local position p attends to local keys 0..p. Simultaneously every device fires off an async send of its KV block to its right-hand neighbour and an async receive from its left. The partial output merges into the accumulator, which simply adopts its values since the accumulator was empty.
Step 1: the async transfers from step 0 have completed during the step-0 arithmetic. Device i now holds KV block (i-1). Under causal masking, device 1 computing against block 0 is fully useful — block 0's tokens all precede device 1's tokens. But device 0 computing against block 3 is entirely masked, since block 3's tokens are all later. Device 0 skips the arithmetic entirely and just forwards the block. Meanwhile the next round of async sends and receives is already in flight. The devices that did real work merge their partials: compute the new running max across the accumulator's m and the block's local max, rescale both sides by the appropriate exponential factors, and accumulate.
Steps 2 and 3 proceed identically, each device computing against a KV block one further around the ring, skipping the fully masked pairs, and merging what it computes. After step N-1 = 3, every KV block has visited every device. Each device's accumulator has absorbed every key its queries were permitted to see. The final normalization divides the running output o by the running sum l, and device i writes out the attention output rows for its own 100,000 tokens. No output ever crosses the interconnect — the output shard is already where it needs to be for the next layer, which is sharded the same way.
The backward pass mirrors this with one important addition. Gradients need the same all-to-all access, so the ring rotates again — this time carrying both the KV blocks and their gradient accumulators. Because the gradient with respect to K_j and V_j receives contributions from every device's queries, those buffers travel with the block, accumulate at each stop, and arrive home fully summed after N steps. The backward ring therefore moves roughly twice the bytes of the forward ring, which is why backward, not forward, is usually what exposes an under-provisioned interconnect.