The Income Tax Appellate Tribunal (ITAT) sometimes hears cases where a single dispute has been litigated across a decade — the original assessment order, three rounds of appeals, cross-references to a dozen earlier judgments, and thousands of pages of financial exhibits. Imagine a legal-AI system asked to read the entire case file in one shot, so it never has to guess which exhibit a paragraph on page 3,000 is referring to. That file, tokenized, can run past two million tokens. No transformer keeps a sequence that long inside one GPU's memory. This chapter is about the trick that makes it possible anyway: instead of asking one device to hold the whole sequence, split the sequence across many devices arranged in a ring, and let the attention computation itself travel around that ring.
The wall you hit at a million tokens
Self-attention has two separate costs, and it matters which one is actually the problem here. The first is compute: for a sequence of length n, every query attends to every key, so the attention score matrix has n² entries, and computing it costs O(n²·d) multiply-adds where d is the head dimension. The second is memory: to generate those scores you need the key and value vectors for every token — the "KV cache" — resident somewhere in fast memory, and if you compute attention the naive way, you also need to materialize the full n×n score matrix before you can take a softmax over each row.
Both costs turn punishing at long context. Take n = 1,000,000 tokens in bf16 (2 bytes per number): the naive score matrix alone is n² = 10¹² entries, i.e. 2 TB — orders of magnitude larger than any single GPU's memory. FlashAttention-style algorithms already fix this half of the problem: they never materialize the full matrix, computing attention in small blocks and combining the results with a running softmax (the technique this chapter builds on in the next section). But even with that fixed, the KV cache itself is still a problem, because it scales linearly in n and that linear term is enormous at this scale. Using Llama-2-7B's published dimensions purely to make this concrete — 32 layers, hidden size 4096, bf16 — each token's K and V vectors across every layer cost 2 (K and V) × 4096 × 2 bytes × 32 layers = 524,288 bytes, about 524 KB per token, just for the cache, before weights or activations. At n = 1,000,000 tokens that is 524,288 bytes × 1,000,000 ≈ 524 GB of KV cache alone. An H100 with 80 GB of HBM can hold roughly 80,000,000,000 ÷ 524,288 ≈ 152,600 tokens of KV cache before it runs out of room for anything else — which is not a coincidence: it is close to the ceiling of context windows (128K–200K tokens) that single-GPU-served models typically advertise. Push past that ceiling and there is no algorithmic cleverness that keeps the sequence on one chip; it has to be split across several.
Splitting the sequence across devices sounds simple — give device 0 tokens 1 through n/p, device 1 the next n/p, and so on, for p devices. The complication is that attention is not local: query 700,000 might need to look at key 12, which now lives on a different device entirely. Every device's queries need to see every other device's keys and values at some point. Ring Attention (Liu, Zaharia, and Abbeel, 2023) is the scheme that makes that exchange happen with the communication overlapped by computation, so that splitting the sequence doesn't just move the memory bottleneck into a communication bottleneck instead.
Building block: the online-softmax trick
Before the ring, one piece of the machinery needs to be nailed down, because it is what makes splitting a softmax across separate compute steps mathematically safe rather than just an approximation. Softmax normally needs the whole row of scores at once — you take the max for numerical stability, exponentiate every score, and divide by the sum. If you only have one block of keys in front of you at a time, you cannot compute that final division yet. What you can do is keep three running numbers per query — a running max m, a running normalizer l, and a running weighted-value accumulator acc — and update them every time a new block of keys and values arrives:
m_new = max(m, max(scores in this block))
correction = exp(m - m_new)
l_new = l * correction + sum(exp(scores - m_new))
acc_new = acc * correction + sum(exp(scores - m_new) * V_block)
The correction factor is the whole trick: whenever a new block reveals a larger max than anything seen so far, the previously accumulated exponentials (computed relative to the old, smaller max) get rescaled down to be consistent with the new max, instead of being thrown away and recomputed. Because of that rescaling, running this update over any partition of a row's keys, in any order, produces exactly the same final acc / l as computing softmax over the whole row in one pass — not an approximation, an exact algebraic reformulation (the same log-sum-exp identity that FlashAttention uses to stream blocks from one GPU's own memory). Ring Attention reuses this identity, but instead of streaming blocks from one device's memory, it streams them from other devices' memory, over the network.
The ring: keep queries still, rotate keys and values
Arrange p devices in a logical ring — device i is connected to device i+1 (mod p) in one direction and device i−1 (mod p) in the other. Split the sequence into p contiguous blocks and give device i the i-th block's queries, keys, and values. The queries never move for the rest of the computation — each device is permanently responsible for producing the output for its own block of queries. What moves is the keys and values.
At each of p steps: every device computes attention scores between its resident (fixed) query block and whichever key/value block currently sits on that device, folds the result into its per-query (m, l, acc) state using the online-softmax update above, and simultaneously — not afterward — sends its current key/value block to its "downstream" neighbour while receiving the next block to work on from its "upstream" neighbour. After p steps, a key/value block has visited every device exactly once (the ring is a Hamiltonian cycle through all p nodes), so every query has been compared against every key in the whole sequence, and each device's accumulated acc / l is the exact, fully normalized attention output for its query block — identical to what a single infinite-memory device would have computed.
The reason this beats simply gathering all the keys onto one device first is the word "simultaneously." Because the send/receive of the next block is issued as a non-blocking operation that runs during the current step's local attention computation, the communication cost is hidden behind useful compute rather than added on top of it — provided the compute takes at least as long as the transfer. Local attention compute for one block scales roughly with (block size)²·d, while transferring one block over the ring link scales with block size·d — quadratic against linear. That means the larger the per-device chunk of sequence (which is exactly the situation at million-token context, where each device might still hold tens of thousands of tokens), the more computation there is to hide each transfer behind, and the closer the whole ring gets to running at the speed of local computation alone, with communication effectively free.
Worked example: tracing the ring by hand
To see that the ring produces exactly the same numbers as ordinary full attention — not an approximation — trace a tiny case where every number can be checked by hand. Take a sequence of n = 4 tokens with 1-dimensional (scalar) query, key, and value vectors, so the "dot product" is just a multiplication:
Q = [ 1, 0, -1, 2]
K = [ 1, 2, 0, -1]
V = [10, 20, 30, 40]
Standard attention, computed the ordinary way. For query 1 (q₁ = 1), the raw scores against all four keys are q₁·K = [1, 2, 0, −1]. Subtracting the max (2) for stability and exponentiating gives exp([−1, 0, −2, −3]) = [0.36788, 1.0, 0.13534, 0.04979], which sum to Z = 1.55300. Dividing each by Z gives the attention weights, and the output is the weights dotted with V: (0.36788·10 + 1.0·20 + 0.13534·30 + 0.04979·40) / 1.55300 = 19.14379. Repeating this for all four queries (the same mechanical steps, just different numbers) gives:
q1 = 1: scores=[ 1, 2, 0,-1] -> output = 19.14379
q2 = 0: scores=[ 0, 0, 0, 0] -> output = 25.00000
q3 =-1: scores=[-1,-2, 0, 1] -> output = 34.37567
q4 = 2: scores=[ 2, 4, 0,-2] -> output = 19.03071
Now the ring version, with p = 2 devices and block size 2. Device 0 owns queries 1–2 and starts by holding keys/values 1–2 (block 0); device 1 owns queries 3–4 and starts by holding keys/values 3–4 (block 1). There are two ring steps: step 0 uses each device's own resident block, and step 1 uses the block received from the neighbour.
Follow query 1 on device 0. At step 0, the resident block is K,V block 0 = (K=[1,2], V=[10,20]). Scores = [1·1, 1·2] = [1, 2]. Nothing has been accumulated yet, so m becomes 2, l = exp(1−2) + exp(2−2) = 0.367879 + 1.0 = 1.367879, and acc = 0.367879·10 + 1.0·20 = 23.678794. At step 1, block 1 arrives (K=[0,−1], V=[30,40]). Scores = [1·0, 1·(−1)] = [0, −1]. The max across everything seen (2 from before, 0 and −1 now) is still 2, so the correction factor is exp(2−2) = 1 — no rescaling needed this time, because the earlier block already happened to contain the largest score. Adding the new terms: l = 1.367879 + exp(0−2) + exp(−1−2) = 1.367879 + 0.135335 + 0.049787 = 1.553002, and acc = 23.678794 + (0.135335·30 + 0.049787·40) = 23.678794 + 6.051541 = 29.730336. The final output is acc / l = 29.730336 / 1.553002 = 19.143787 — matching the standard-attention answer for q1 above to six decimal places.
That match held even though the result was assembled from two separate, non-overlapping key blocks arriving in two separate steps on two different simulated devices, never seeing the full row of scores at once. The same holds for all four queries — this was checked in full floating-point precision by executing the online-softmax update on all four queries, on both devices, and comparing every one of the four outputs against the standard computation:
import math
def online_softmax_update(m, l, acc, scores, V_block):
"""Fold one new K,V block into the running (m, l, acc) state.
Exactly reproduces the standard softmax if applied to all blocks
of a row, in any order."""
m_new = max(m, max(scores))
correction = math.exp(m - m_new)
exps = [math.exp(s - m_new) for s in scores]
l_new = l * correction + sum(exps)
acc_new = acc * correction + sum(e * v for e, v in zip(exps, V_block))
return m_new, l_new, acc_new
K_blocks = [[1, 2], [0, -1]]
V_blocks = [[10, 20], [30, 40]]
# device 0 owns q1; the ring visits its own block, then the neighbour's
q1 = Q[0]
m, l, acc = -math.inf, 0.0, 0.0
for kv_step in [0, 1]:
scores = [q1 * k for k in K_blocks[kv_step]]
m, l, acc = online_softmax_update(m, l, acc, scores, V_blocks[kv_step])
print(f"q1 ring-attention output = {acc / l:.5f}")
q1 ring-attention output = 19.14379
q1 standard attention output = 19.14379
Running the same check for all four queries across both devices gives outputs of 19.14379, 25.00000, 34.37567, and 19.03071 — identical to standard full attention in every case, confirming the online-softmax rescaling is exact, not an approximation, regardless of which order the key/value blocks arrive in.
The misconception to unlearn
The mistake nearly every student makes on first meeting Ring Attention is assuming it makes attention cheaper — that splitting the n² computation across p devices somehow reduces the total number of floating-point operations needed. It does not. Every query still has to be compared against every key exactly once; the total work is still O(n²·d), just distributed as roughly n²·d/p per device instead of all of it on one chip. Ring Attention is a parallelization and memory-capacity technique, not a compute-reduction technique — it belongs in the same family as data parallelism and tensor parallelism, not in the family of sparse attention or linear attention, which genuinely do reduce the O(n²) term algorithmically (usually by approximating or restricting which keys each query is allowed to see). Ring Attention computes exactly the same attention output as a hypothetical single device with infinite memory would — the worked example above demonstrates that equivalence numerically — and its entire benefit is that no single device ever needs to hold the whole sequence's keys and values, or the whole score matrix, at once, while the unavoidable data movement between devices is hidden behind computation rather than added as pure overhead.
Active recall
Attempt these before reading the answers below.
- Does Ring Attention reduce the total number of floating-point multiply-adds needed for self-attention on a sequence of length n? Why or why not?
- Why must Ring Attention use a running ("online") softmax instead of the ordinary softmax formula at each step?
- For a ring of p = 4 devices, how many steps does it take before every device's query block has attended to every key/value block in the sequence, and what determines whether the communication needed for that is actually hidden behind computation?
- A deployment needs 100,000 tokens of context, and the model's KV cache costs 400 KB per token across all layers. What is the total KV-cache memory required, and how would you size a 5-device ring so each device's cache fits inside a 24 GB budget?
- True or false: Ring Attention changes the numerical result of attention compared with computing it on one device with unlimited memory. Justify your answer.
- Why is Ring Attention specifically valuable at million-token context lengths but not at the 4K–8K context lengths common in older chat models?
Answers.
1. No. Total attention compute stays at O(n²·d): every one of the n² query-key pairs is still scored exactly once, just spread across the p devices in the ring (about n²·d/p work per device) instead of computed on a single chip. Ring Attention parallelizes and reduces per-device memory footprint; it does not shrink the total arithmetic. Reducing the O(n²) term itself is a different family of techniques (sparse or linear attention).
2. Ordinary softmax needs the sum over every key in the row before it can normalize any single weight, but at a given ring step a device only has one key/value block resident — not the full row. The online-softmax update lets each device combine partial results as blocks arrive, using the exp(m_old − m_new) correction to rescale earlier partial sums whenever a later block reveals a bigger maximum score. Because that rescaling is an exact algebraic identity, the final accumulated result after all blocks have arrived is bit-for-bit equivalent to computing softmax over the whole row at once, as the worked example confirmed to six decimal places.
3. p = 4 steps total (each device processes its own starting block, then 3 more arriving from around the ring), since the ring is a Hamiltonian cycle that visits every one of the p nodes exactly once starting from any node. Whether the necessary transfer is hidden depends on whether local compute time per block — which scales roughly with (block size)²·d, quadratic in the per-device chunk length — exceeds the transfer time per block over the ring link, which scales roughly with block size·d, linear. Larger per-device chunks (longer context per GPU) make it easier for compute to outpace and hide the communication.
4. Total KV cache = 100,000 tokens × 400 KB/token = 40,000,000 KB = 40 GB. Split across a 5-device ring, each device holds 100,000 ÷ 5 = 20,000 tokens' worth of keys and values, i.e. 40 GB ÷ 5 = 8 GB of KV cache per device — comfortably inside a 24 GB budget, leaving roughly 16 GB per device free for activations, the local query block, and model weights.
5. False. The worked example computed both the standard full-sequence attention and the 2-device ring-attention trace for the same toy sequence and got identical outputs — 19.14379, 25.00000, 34.37567, 19.03071 — for all four queries, because the online-softmax accumulation is an exact reformulation of the same softmax computation, not an approximation of it.
6. At 4K–8K tokens, the KV cache is only a few tens of megabytes to a few hundred megabytes — it fits inside any single GPU's memory with room to spare, so there is no memory problem for Ring Attention to solve, and the ring's communication overhead only adds cost with nothing to offset it. At million-token context, KV cache alone reaches hundreds of gigabytes (roughly 524 GB for a Llama-2-7B-scale model at n = 1,000,000 tokens, as computed earlier), which exceeds even a top-end single GPU's HBM (80–192 GB); sharding the sequence across a ring of devices becomes the only way to fit the sequence at all, and because each device's chunk is still large at that scale, the ring's communication has enough local compute to hide behind.
Practice Exercises
Now it is time to practice! Complete these challenges to solidify your understanding:
- Exercise 1: Write a short program that demonstrates the core concept from this chapter. Test it with at least 3 different inputs.
- Exercise 2: Find a real-world example where ring attention: distributed attention across devices is used in an Indian company (like TCS, Infosys, Flipkart, or ISRO). Write a paragraph explaining the connection.
- Exercise 3: Create a mind-map connecting ring attention: distributed attention across devices to at least 3 other topics you have studied.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind ring attention: distributed attention across devices, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.