Ring Attention

Ring attention computes exact self-attention over a sequence too long for one accelerator by splitting it into blocks across devices arranged in a ring. Each device keeps one block of queries and passes blocks of keys and values to its neighbor, overlapping that communication with the computation on the block it already holds, so the longest sequence a model can handle grows in proportion to the number of devices. Hao Liu, Matei Zaharia, and Pieter Abbeel of UC Berkeley introduced it in October 2023 in Ring Attention with Blockwise Transformers for Near-Infinite Context, building on their earlier blockwise parallel transformer, and reported training and inference with context windows of millions of tokens without approximating attention.

How Ring Attention Works#

A sequence is split into N blocks, one per GPU or TPU. Each device holds its query block for the whole computation. In each of N steps it computes attention between its queries and the key-value block it currently holds, sends that block to the next device, and receives another from the previous one:

After N steps, every query block has met every key and value block. No device ever stores the full attention matrix or the full set of keys and values, so memory per device depends on the block size instead of the sequence length. The transfer is hidden behind computation as long as computing attention on one block takes longer than sending one, which sets a minimum block size for given hardware. Partial results combine with the same running softmax that flash attention uses inside one GPU: each device keeps a running sum of exponentiated scores and of the weighted values, and divides one by the other at the end. Simulating four devices on one machine shows that the result is exact:

import torch

torch.manual_seed(0)
n_devices, block, dim = 4, 8, 16
q, k, v = torch.randn(3, n_devices * block, dim)
q_blocks, k_blocks, v_blocks = q.chunk(n_devices), k.chunk(n_devices), v.chunk(n_devices)

outputs = []
for i in range(n_devices):  # device i keeps query block i for the whole computation
    numerator, denominator = torch.zeros(block, dim), torch.zeros(block, 1)
    for step in range(n_devices):
        j = (i - step) % n_devices  # key-value block that has arrived from the previous device
        p = torch.exp(q_blocks[i] @ k_blocks[j].T / dim**0.5)
        numerator += p @ v_blocks[j]
        denominator += p.sum(dim=-1, keepdim=True)
    outputs.append(numerator / denominator)

full = torch.softmax(q @ k.T / dim**0.5, dim=-1) @ v
print(f"Largest score block per device: {block}x{block}, full matrix: {len(q)}x{len(k)}")
print("Matches full attention:", torch.allclose(torch.cat(outputs), full, atol=1e-5))

It prints Matches full attention: True: each device only ever held an 8 × 8 block of scores, and together they produced attention over the full 32 × 32 matrix. Production kernels also subtract a running maximum before exponentiating to keep the numbers stable. With a causal mask, devices holding early query blocks skip most key blocks and sit idle, which Striped Attention fixes by giving each device tokens spread evenly through the sequence, raising end-to-end training throughput up to 1.45 times at 256K tokens.

Applications#

Ring attention matters wherever one input is too long for a single device, and video is the clearest case. The same Berkeley group used it to train the Large World Model (LWM), a family of 7-billion-parameter vision-language models (VLMs) whose context grew step by step from 4K to 1 million tokens of interleaved text and video. LWM encodes each frame as 256 tokens, so an hour of video sampled at one frame per second comes to 921,600 tokens, and the model answered questions about hour-long YouTube compilations, a long-range video understanding task. Training large language models (LLMs) that read whole codebases or document collections in one prompt hits the same limit. PyTorch's context_parallel() API swaps scaled dot-product attention for ring attention inside its context, and the TorchTitan team used it to train Llama 3 8B on 1-million-token sequences across 32 H100 GPUs.

Ring Attention vs. Other Parallelism#

• Context parallelism: The general strategy of splitting one sequence across devices. Ring attention is the algorithm behind PyTorch's implementation, and DeepSpeed Ulysses is an alternative that uses all-to-all exchanges to give each device all tokens for a subset of attention heads.
• Tensor parallelism and pipeline parallelism: Split the model's weights within layers or across layers, which helps when the parameters outgrow one device's memory.
• Flash attention: A single-device kernel. Ring attention runs a blockwise kernel of this kind on each device for its local block.

Data parallelism splits the batch instead of the sequence, and it is the mode Ultralytics YOLO26 uses for multi-GPU training: passing device=[0, 1] launches PyTorch DistributedDataParallel, which keeps a full model copy on each GPU, gives each GPU an equal share of every batch of images, and averages gradients after each step.