Core ML

Ring Attention

To handle a sequence too long for one GPU, ring attention shards it across devices arranged in a ring, streaming key-value blocks around the ring while each device computes its local attention against whichever block passes through — every device sees all keys and values without storing them all.

Why Does This Exist?

FlashAttention solves the memory-bandwidth problem for attention on a single GPU. But for very long sequences — 100K tokens or more — even the $O(N)$ memory cost of storing Q, K, and V becomes prohibitive on a single device. You can't tile your way out of a problem where the raw tensors don't fit.

Ring attention (Liu et al., 2023) is a distributed algorithm that takes FlashAttention's tiling logic and extends it across devices. Instead of tiling within one GPU's SRAM, you tile across GPUs: shard the sequence evenly, pass the KV blocks around a device ring, and let each device compute its local slice of the attention output as each KV block arrives. At the end of one full ring pass, every query has attended to every key, and the output is assembled without any device ever storing the full sequence.

Think of It Like This

Passing a reference binder around a table

Eight researchers sit at a circular table, each with their own notes (the query slice they own). A single reference binder (one KV block) starts with researcher 1, who reads the relevant pages and makes notes. Then the binder passes to researcher 2, who does the same — consulting the same binder against their own question set. After the binder completes one circuit of the table, every researcher has seen all of it, but only one copy ever existed and nobody carried more than one person's worth of notes at a time.

Ring attention is that binder passing around a table of GPUs.

How It Actually Works

Sequence sharding and the ring topology

With $D$ devices and a sequence of length $N$ , each device holds a slice of the Q, K, and V tensors of length $N/D$ . Device $d$ owns $Q_d$ , $K_d$ , $V_d$ .

The devices form a logical ring: device 0 → device 1 → ... → device $D-1$ → device 0. At each step:

Each device sends its current KV block to the next device in the ring.
While sending (communication), each device also computes attention between its own $Q_d$ and the KV block it currently holds (compute).

Communication and compute overlap: the send of the current KV block happens in parallel with computing attention against it, hiding most of the communication latency.

Online softmax across ring steps

After $D$ ring steps, each device has accumulated partial attention outputs from all $D$ KV blocks. These partials are combined using the same online softmax recurrence that FlashAttention uses within a single GPU. Each device maintains a running numerator and denominator across ring steps, correcting for the row-wise maximum as each new block arrives. The final output on each device is exact.

Memory per device

Each device holds $N/D$ rows of Q and accumulates $N/D$ rows of output. The KV block in transit is also $N/D$ rows. Total memory per device: $O(N/D)$ — linear in the local slice, not the global sequence length. Scaling to 8 devices cuts per-device memory 8×.

Show Me the Code

Simulating two-device ring attention and verifying against a single-device standard pass.

1import numpy as np2
3rng = np.random.default_rng(5)4n, d = 8, 16   # n must be even for 2-device split5
6Q = rng.standard_normal((n, d)) * 0.17K = rng.standard_normal((n, d)) * 0.18V = rng.standard_normal((n, d)) * 0.19
10# Device 0 owns first half, device 1 owns second half11Q0, Q1 = Q[:n//2], Q[n//2:]12
13def softmax_attn(q, k, v):14    scores = q @ k.T / d**0.515    scores -= scores.max(-1, keepdims=True)16    w = np.exp(scores)17    return w @ v, w.sum(-1)18
19# --- Ring step 1: each device attends to its own KV ---20# Device 0: Q0 × K0,V021out0_a, w0_a = softmax_attn(Q0, K[:n//2], V[:n//2])22# Device 1: Q1 × K1,V123out1_a, w1_a = softmax_attn(Q1, K[n//2:], V[n//2:])24
25# --- Ring step 2: KV blocks rotate (device 0 now sees K1,V1 and vice versa) ---26out0_b, w0_b = softmax_attn(Q0, K[n//2:], V[n//2:])27out1_b, w1_b = softmax_attn(Q1, K[:n//2], V[:n//2])28
29# Combine: weight partial outputs by their softmax sums30out0 = (out0_a * w0_a[:, None] + out0_b * w0_b[:, None]) / (w0_a + w0_b)[:, None]31out1 = (out1_a * w1_a[:, None] + out1_b * w1_b[:, None]) / (w1_a + w1_b)[:, None]32ring_out = np.vstack([out0, out1])33
34# Reference: single-device standard attention35scores = Q @ K.T / d**0.536scores -= scores.max(-1, keepdims=True)37w = np.exp(scores)38w /= w.sum(-1, keepdims=True)39ref_out = w @ V40
41print("Max difference:", np.abs(ring_out - ref_out).max())42# ~ 1e-15 (floating-point rounding only)

The results agree to floating-point precision — ring attention is exact, not approximate.

Watch Out For

Assuming ring attention works with causal masking naively

Causal masking in ring attention is more involved than in single-device attention. A device at position $d$ attending to a KV block from device $d' > d$ has to mask future tokens, but "future" is defined globally, not per-block. The masking logic needs to know the global position of each KV block to apply the right triangular mask. Getting this wrong produces non-causal outputs that are hard to detect by loss alone, since the training signal is still there — just contaminated.

Treating ring attention as sequence parallelism generally

Ring attention is one specific form of sequence parallelism — the form where the sequence is sharded and KV blocks rotate. Other forms (e.g., Ulysses sequence parallelism) shard differently, using all-to-all collectives rather than a ring pass. The right choice depends on your network topology: ring is efficient on a ring-structured interconnect; Ulysses is more efficient on all-to-all fabrics like NVLink. They're not interchangeable.

The Quick Version

Ring attention shards the sequence across $D$ devices and passes KV blocks around a ring, letting each device compute partial attention outputs against each block as it arrives.
Communication (sending the KV block) overlaps with compute (attending against it), hiding latency.
Memory per device is $O(N/D)$ — scaling with the local slice, not the full sequence.
The final output is exact, combining partial results with an online softmax recurrence.
Causal masking requires careful global-position tracking per KV block during ring steps.