Ring vs. Tree Attention

Why This Matters

Flash Attention shows that one device can compute exact attention without ever materializing the full $N \times N$ score matrix, by tiling $K, V$ into blocks and maintaining a running online-softmax state. The natural next question, once a context is too long or a batch too large for one device: what happens when the $K, V$ blocks themselves live on different devices?

The local attention calculation does not change. Ring Attention (Liu, Zaharia, Abbeel, ICLR 2024) and Tree Attention (Shyam et al., 2024) use different schedules to combine device-local partial results under the same exact real-arithmetic reduction.

The shared building block: the online-softmax merge

Each device holds a shard of $K, V$ and computes a local, unnormalized partial state from its own shard: a running max score $m$ and a running sum-of-exponentials $l = \sum_i e^{s_i - m}$ (plus the corresponding unnormalized output accumulator, omitted here for clarity). Two partial states $(m_1, l_1)$ and $(m_2, l_2)$ combine via

$$m = \max(m_1, m_2), \qquad l = l_1 e^{m_1 - m} + l_2 e^{m_2 - m}.$$

This normalizer merge is associative and commutative. It is the log-sum-exp combine rule for a running maximum and normalizer. That fact lets devices fold partial states in any order or tree shape while preserving the same $(m, l)$ in exact arithmetic. Full attention also merges a separately scaled output accumulator; this page uses the normalizer pair because it makes the topology visible without introducing vector-valued state.

Two topologies, same answer, different communication cost

Ring Attention rotates $K, V$ shards around a ring of $P$ devices. Each device sends its local shard to a neighbor, receives another device's shard, and updates its running online-softmax state. After $P - 1$ hops, every device has processed every shard. This is $P - 1$ sequential point-to-point exchanges, each moving a full $K, V$ shard. Liu, Zaharia, and Abbeel overlap each exchange with local computation, but the round count and payload per round still grow linearly with device count.

Tree Attention has each device compute a local $(m, l)$ state against the shared query, then combines those small states rather than moving the underlying $K, V$ data. Pairs merge in round 1, survivors pair in round 2, and the schedule finishes in $\lceil \log_2 P \rceil$ rounds, with up to $P/2$ merges per round. Shyam et al. prove the underlying per-query calculation admits an $N/P + O(\log P)$ parallel decomposition (Theorem 1, arXiv:2408.04093). Their §6.1 experiment on 16 nodes and 128 H100 GPUs at a 5.12M-token sequence length reports close to $8\times$ lower decoding latency than their Ring Attention implementation. Their §6.2 memory formulas give $\text{Mem}_{\text{ring}} = 4btd + 2bd$ versus $\text{Mem}_{\text{tree}} = 2btd + 2bd + 2bn_h$ for batch size $b$, local chunk length $t$, model dimension $d$, and head count $n_h$ (arXiv:2408.04093, §6.2, Eqs. 8-9). Tree Attention's dominant term is half of Ring Attention's. Measured on two RTX 4090s, doubling hidden size from 2048 to 4096 doubled the absolute memory gap between the two methods, from 524MB to 1040MB (§6.2).

Neither topology introduces an attention approximation. In exact arithmetic, both evaluate the same attention reduction as a single-device Flash Attention calculation. In floating point, different reduction order can produce machine-epsilon differences. The interactive diagram checks the $(m, l)$ normalizer state against a direct calculation; it does not model the output accumulator.

Main Theorem

Proposition

The Online-Softmax Normalizer Merge Is a Commutative Semigroup

Statement

For two normalizer states $(m_1, l_1)$ and $(m_2, l_2)$, define $\operatorname{merge}((m_1, l_1), (m_2, l_2)) = (m, l)$ where $m = \max(m_1, m_2)$ and $l = l_1 e^{m_1-m} + l_2 e^{m_2-m}$. This operation is associative and commutative. Folding any partition of a score sequence into partial states, then merging those states in any order or binary-tree shape, produces the same $(m, l)$ as computing the normalizer over the whole sequence at once.

Intuition

$\operatorname{merge}$ is the log-sum-exp combine rule carried alongside a running maximum for numerical stability. Addition and max are associative and commutative; the exponential rescaling puts both local sums on the same maximum scale. The result depends on the included scores, not their grouping.

Proof Sketch

Commutativity follows because swapping the inputs leaves both $m$ and $l$ unchanged. For associativity, expand $\operatorname{merge}(\operatorname{merge}(s_1, s_2), s_3)$ and $\operatorname{merge}(s_1, \operatorname{merge}(s_2, s_3))$. Both reduce to $m = \max(m_1, m_2, m_3)$ and $l = \sum_i l_i e^{m_i-m}$. Induction extends the identity to any finite fold order or tree shape. Flash Attention uses the same identity across sequential blocks on one device.

Why It Matters

This is the entire reason Ring Attention and Tree Attention can both be exact. Neither paper is proposing a new numerical method; both propose a different schedule for combining the same associative operation, so the choice of topology is a communication-engineering decision, not a sacrifice of exactness.

Failure Mode

The result requires disjoint score subsets from the same query row. Merging states from different queries, or double-counting a score across shards, produces the wrong $(m, l)$ without an error signal. IEEE 754 arithmetic also changes the result: ring and tree parenthesize rescaling and addition in different orders, so each can differ from a direct calculation by machine epsilon. The full attention output needs its separately scaled output accumulator in addition to this normalizer pair.

report a correction →

Common Confusions

Watch Out

Ring and Tree Attention do not approximate attention

Both compute the exact same attention output as single-device Flash Attention, up to floating-point rounding. They are not sparse or linear attention methods (which do trade exactness for speed); they only change how many communication rounds combining distributed partial results takes.

Watch Out

This does not replace Flash Attention

Ring Attention and Tree Attention both use Flash Attention's tiled, online-softmax computation within each device; they extend it across devices, they do not compete with it. A cluster running Tree Attention still runs Flash Attention locally on each GPU's shard.

Exercises

ExerciseCore

Problem

8 devices hold local partial states with $(m, l)$ pairs $(2, 3), (1, 4), (3, 1), (0, 5), (2, 2), (1, 6), (4, 1), (0, 3)$ in that order. Compute the final merged $(m, l)$ using the ring schedule (fold left to right), then using a balanced binary tree, and confirm they agree.

Hint

Reveal Solution

ExerciseAdvanced

Problem

Using the peak-memory formulas $\text{Mem}_{\text{ring}} = 4btd + 2bd$ and $\text{Mem}_{\text{tree}} = 2btd + 2bd + 2bn_h$, find the condition on $n_h$ (number of attention heads) under which Tree Attention's peak memory is strictly lower than Ring Attention's, and check whether that condition holds for a typical configuration with $t = 8192$, $d = 4096$.

Hint

Reveal Solution

Mental Model

Think of the $P$ local partial states as $P$ numbers that need repeated pairwise merge operations. A ring is a minimum-parallelism schedule: it serializes $P - 1$ merges, one per round. A balanced tree performs up to $P/2$ merges per round and uses $\lceil \log_2 P \rceil$ rounds. This is the same reduction pattern used for a parallel sum or max; Tree Attention applies it to the online-softmax state from Flash Attention.

Scope

This module is about the reduction step's communication structure, not about sparse or linear attention approximations (those trade exactness for speed; ring and tree do not) and not about single-device tiling mechanics (covered in Flash Attention). Three current papers ground it: FlashAttention (Dao et al., NeurIPS 2022) for the online-softmax merge itself, Ring Attention (Liu, Zaharia, Abbeel, ICLR 2024) for the sequential-ring schedule, and Tree Attention (Shyam et al., 2024, revised 2025) for the topology-aware tree schedule and its measured speedup over Ring Attention.

References

Exact attention and online softmax:

• Dao, T., Fu, D. Y., Ermon, S., Rudra, A., Re, C. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." NeurIPS 2022, arXiv:2205.14135. §§3.1–3.3 for tiled attention and the online-softmax recurrence.
• Milakov, M., Gimelshein, N. "Online normalizer calculation for softmax." arXiv:1805.02867. Algorithm 3 for the stable online normalizer.
• Rabe, M. N., Staats, C. "Self-attention Does Not Need $O(n^2)$ Memory." arXiv:2112.05682. §3 for exact memory-efficient attention.
• Dao, T. "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning." ICLR 2024, arXiv:2307.08691. §§3–4 for parallel work partitioning within exact attention.

Distributed schedules:

• Liu, H., Zaharia, M., Abbeel, P. "Ring Attention with Blockwise Transformers for Near-Infinite Context." ICLR 2024, arXiv:2310.01889. Algorithm 1 for ring rotation and computation-communication overlap.
• Shyam, V., Pilault, J., Shepperd, E., Anthony, Q., Millidge, B. "Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU Clusters." arXiv:2408.04093. Theorem 1; §§6.1–6.3; Eqs. 8–9.