📄 Striped Attention: Faster Ring Attention for Causal Transformers

Summary

Ring Attention (Liu et al., 2023) makes very long contexts trainable by splitting queries, keys and values along the sequence across accelerators arranged in a ring. Key/value blocks circulate between devices while compute overlaps communication. The authors show this scheme cannot benefit from causal masking. Each device holds a contiguous chunk, so after the first round some devices have fully masked tiles and others fully unmasked ones. Each round runs at the pace of the slowest device, so a causal run costs as much as a bidirectional one, even though it needs only about half the arithmetic.

Their fix, Striped Attention, reorders tokens so that every device holds positions spread evenly across the sequence (for example, every N-th token). Each device then skips roughly half its work in every round. Because attention is permutation-equivariant, the output stays exact, and the layout is set once before the first layer. Built on the JAX Ring Attention codebase and tested on 1B–7B causal models, it trains up to 1.65× faster on long sequences. This makes very long-context generative models cheaper to train.

Key points

• Imbalance: in causal Ring Attention, the work per round is either fully needed or fully discarded depending on the device. Masking therefore saves no wall-clock time.
• Striped layout: the token permutation is applied once before the model, with no extra per-layer communication, and results remain exact.
• A100 results: on a server with 8 A100 80GB GPUs, end-to-end training throughput improved by up to 1.45× over Ring Attention at a sequence length of 262,144 (3B model).
• TPU results: on 16 TPUv4 chips, speedups reached 1.65× at 786,432 tokens (3B model). On 8 TPUv3 chips, speedups reached up to 1.47×.
• Limitations and release: at short sequences (8,192 tokens on A100), the method is slightly slower than Ring Attention (0.92–0.98×). The code is open source at github.com/exists-forall/striped_attention.