Context Parallelism

Activation memory scales with sequence length, so a long-context run hits a wall that no amount of model sharding addresses — the activations for one 128k-token sequence do not fit on one device however small the weights are.

Splitting the sequence is the obvious response and attention is the obstacle. Position 100 000 must attend to position 5, which is on a different machine.

Ring attention

The insight comes from FlashAttention: attention is already computed blockwise, accumulating an output and a running softmax normaliser over key blocks. Nothing requires those blocks to be in the same memory.

Arrange the devices in a ring. Each holds its own query block permanently and its own KV block initially. Then, $N-1$ times: compute attention of the local queries against the currently held KV block, and pass that KV block to the next device while receiving one from the previous.

After $N-1$ rounds every query has seen every key. The result is exact — this is not an approximation of full attention, it is full attention with the accumulation distributed.

The causal imbalance

With a causal mask, a naïve contiguous split is badly unbalanced. Rank 0 holds queries 0–1023 and attends to almost nothing; rank 7 holds the last chunk and attends to everything. Rank 7 sets the pace and the rest wait.

Zigzag assignment fixes it: split into $2N$ chunks and give rank $i$ chunks $i$ and $2N-1-i$. Every rank then holds one early chunk and one late one, and the total masked-in work is equal across ranks.

Ulysses, the alternative

DeepSpeed Ulysses splits the sequence too, but instead of ringing the KV around, it performs an all-to-all before attention to redistribute the tensor by head rather than by position — each rank then computes complete attention for a subset of heads — and an all-to-all after to redistribute back.

Ulysses moves fewer bytes and cannot scale past the number of attention heads, which GQA makes a real constraint — eight KV heads means eight devices. Llama 3 uses a ring-based scheme for that reason.

import torch
import torch.distributed as dist
from torch import Tensor


def ring_attention(q: Tensor, k: Tensor, v: Tensor, group) -> Tensor:
    """q, k, v: this rank's chunk, [B, H, C, D]. Exact full attention."""
    world, rank = dist.get_world_size(group), dist.get_rank(group)
    out = torch.zeros_like(q)
    m = torch.full(q.shape[:-1] + (1,), -torch.inf, device=q.device)
    lse = torch.zeros_like(m)

    for step in range(world):
        # Start the next hop before computing on the current block, so the
        # transfer overlaps the matmul rather than following it.
        if step + 1 < world:
            next_k, next_v = torch.empty_like(k), torch.empty_like(v)
            reqs = [dist.P2POp(dist.isend, k, (rank + 1) % world, group),
                    dist.P2POp(dist.irecv, next_k, (rank - 1) % world, group),
                    dist.P2POp(dist.isend, v, (rank + 1) % world, group),
                    dist.P2POp(dist.irecv, next_v, (rank - 1) % world, group)]
            handles = dist.batch_isend_irecv(reqs)

        out, m, lse = flash_block(q, k, v, out, m, lse, source=(rank - step) % world)

        if step + 1 < world:
            for h in handles:
                h.wait()
            k, v = next_k, next_v

    return out

The source argument is not decoration — the block currently held came from a known rank, and the mask depends on where those keys sit in the global sequence. Getting it wrong produces a model that trains, with an attention pattern that is causal on each device and not globally.