Ring attention and sequence parallelism

Setup#

pip install numpy

Written against NumPy 1.26.4, Python 3.10. This simulates the ring on one CPU, so you can read the algorithm without a cluster.

The whole algorithm, simulated#

import numpy as np

rng = np.random.default_rng(0)
N, d, P = 16, 8, 4                       # 16 tokens, head dim 8, a ring of 4 "devices"
S = N // P                               # 4 tokens per device
Q, K, V = (rng.normal(size=(N, d)).astype(np.float32) for _ in range(3))
scale = 1.0 / np.sqrt(d)

def reference(Q, K, V):
    s = (Q @ K.T) * scale
    s = s - s.max(-1, keepdims=True)
    p = np.exp(s)
    return (p / p.sum(-1, keepdims=True)) @ V

# Each device owns one shard of Q, K, V and never sees the whole sequence.
q = [Q[r*S:(r+1)*S] for r in range(P)]
k = [K[r*S:(r+1)*S] for r in range(P)]
v = [V[r*S:(r+1)*S] for r in range(P)]

O = [np.zeros((S, d), np.float32) for _ in range(P)]
m = [np.full(S, -np.inf, np.float32) for _ in range(P)]
l = [np.zeros(S, np.float32) for _ in range(P)]

peak_kv = 0
for hop in range(P):                     # P steps: each device sees every KV shard once
    for r in range(P):
        Kb, Vb = k[r], v[r]              # after `hop` rotations this is shard (r - hop) % P
        peak_kv = max(peak_kv, Kb.nbytes + Vb.nbytes)
        s = (q[r] @ Kb.T) * scale
        m_new = np.maximum(m[r], s.max(-1))
        corr = np.exp(m[r] - m_new)
        p = np.exp(s - m_new[:, None])
        l[r] = corr * l[r] + p.sum(-1)
        O[r] = corr[:, None] * O[r] + p @ Vb
        m[r] = m_new
    # the ring rotation: every device hands its KV shard to its neighbour
    k = k[-1:] + k[:-1]
    v = v[-1:] + v[:-1]

ring = np.concatenate([O[r] / l[r][:, None] for r in range(P)])
ref = reference(Q, K, V)
print("ring output matches full attention:", np.allclose(ring, ref, atol=1e-5))
print("max difference:", float(np.abs(ring - ref).max()))
print()
print(f"devices in the ring          : {P}")
print(f"tokens held per device       : {S}  (of {N})")
print(f"largest KV block in flight   : {peak_kv:,} bytes")
print(f"whole-sequence KV would be   : {K.nbytes + V.nbytes:,} bytes")
print(f"communication steps          : {P} rotations, each of one KV shard")

print()
print("what this buys at scale, one head, bf16")
print(f"{'devices':>8}{'seq len':>12}{'per-device KV':>16}{'full KV':>14}")
for P2 in (8, 64, 512):
    n = P2 * 32768
    per = 2 * (n // P2) * 128 * 2
    print(f"{P2:>8}{n:>12,}{per/2**20:>13.1f} MB{2*n*128*2/2**30:>11.1f} GB")

ring output matches full attention: True
max difference: 1.7881393432617188e-07

devices in the ring          : 4
tokens held per device       : 4  (of 16)
largest KV block in flight   : 256 bytes
whole-sequence KV would be   : 1,024 bytes
communication steps          : 4 rotations, each of one KV shard

what this buys at scale, one head, bf16
 devices     seq len   per-device KV       full KV
       8     262,144         16.0 MB        0.1 GB
      64   2,097,152         16.0 MB        1.0 GB
     512  16,777,216         16.0 MB        8.0 GB

Reading the output#

The ring result equals full attention to 1.8e-07. This is the property that matters. Ring attention is not an approximation, a windowing scheme or a sparsification. It is the same function, evaluated in a different order across machines.

The running statistics m, l and O are per device. Compare them to the FlashAttention loop: identical structure. Ring attention is FlashAttention's tiling where the outer loop crosses a network instead of a memory hierarchy. That is the whole idea in one sentence.

largest KV block in flight is one shard, not the sequence. Per-device memory is set by shard size. That is why the last table shows a flat 16 MB per device while total sequence length grows to 16 million.

The last column grows and the third does not. Add devices, get context. This is the only technique in this section with that property.

Where the rotation happens for real#

In a distributed implementation the list rotation becomes a point-to-point send and receive:

# conceptual, using torch.distributed
next_k = torch.empty_like(k)
req_s = dist.isend(k, dst=(rank + 1) % world_size)     # start sending now
req_r = dist.irecv(next_k, src=(rank - 1) % world_size)
out, m, l = attend_block(q, k, v, out, m, l)           # compute while it travels
req_s.wait(); req_r.wait()
k = next_k

The isend/irecv before the compute call is the entire performance story. If one block of attention takes longer than one shard's network hop, communication is fully hidden and the ring costs nothing beyond the arithmetic.

What to use instead of writing this yourself#

• PyTorch context parallel (torch.distributed.tensor.experimental.context_parallel) — documented as a prototype, and used by torchtitan. The PyTorch team reports scaling Llama3-8B to 1M sequence length on 32 H100s, composed with FSDP and torch.compile.
• ring-flash-attention — community kernels wrapping flash-attn with the ring loop.
• DeepSpeed-Ulysses — a different sequence-parallel scheme that all-to-alls over the head dimension rather than passing KV around a ring.

The PyTorch implementation offers two variants: an all-gather "pass-KV" algorithm and an all-to-all one. Their benchmarks found all-to-all rarely beats all-gather, and all-gather is the default. Llama 3 training used the all-gather pass-KV form.

The load-balance problem, concretely#

With causal masking and contiguous shards, device r computes attention for its queries against shards 0..r. Device 0 does one block of work; device P-1 does P blocks. Average utilisation is about 50%.

The standard fix is zigzag or striped assignment: give each device two half-shards, one from the front of the sequence and one from the back. Every device then holds a mix of early and late queries and the per-hop work evens out.

Any ring implementation that does not mention load balancing is leaving roughly half the throughput on the floor.

Common mistakes#

Confusing sequence parallelism with tensor or pipeline parallelism. Tensor parallelism splits the weights of a layer. Pipeline parallelism splits layers across devices. Sequence parallelism splits the tokens. They compose, and they solve different shortages.

Ignoring the causal imbalance. Covered above. Symptom: adding devices scales context but throughput per device drops by half.

Choosing a shard size that leaves the GPU idle. If one block of attention finishes faster than a network hop, communication stops being hidden. Larger shards, fewer devices, or a faster interconnect.

Running a ring across slow links. Inside a node over NVLink this works well. Across nodes over standard Ethernet, the hop dominates and you have built a slow machine. Check your interconnect bandwidth against your per-block compute time before designing around this.

Assuming it helps decoding. Ring attention shines when there are many queries to process at once, which means prefill and training. Single-token decoding has one query row and nothing to overlap.

Try it yourself#

Add causal masking to the simulation: at hop h, device r should skip shards from the future entirely. Count the blocks each device computes and confirm the 1, 2, 3, 4 imbalance. Then implement zigzag sharding and watch the counts even out.

What to learn next#

• Advertised context vs usable context — whether the context you paid to train is the context you get.
• FlashAttention — the tiling argument ring attention lifts onto a network.
• Docker for ML — packaging a distributed training job so it runs the same everywhere.