Context Parallelism (CP)

Definition

Context parallelism partitions the sequence dimension itself across GPUs, so each holds only $s/N$ query/key/value tokens. It is the axis for long-context training: at million-token context the KV cache and the $O(s^2)$ attention working set blow up regardless of how you shard weights, and TP/SP/PP all still leave attention over the full sequence on a single GPU. Ring Attention (Liu, Zaharia, Abbeel, arXiv:2310.01889) is the enabling algorithm.

Key math / How it works

The problem: attention is not token-local — query block $i$ must attend to all KV blocks, which live on other GPUs. Ring Attention arranges the CP GPUs in a ring: each GPU computes attention of its local queries against its local KV block, then passes its KV block to the next GPU while simultaneously computing on the KV block it just received. Over $N$ steps every query has seen every KV block.

Because attention is computed blockwise with online softmax (FlashAttention), partial results accumulate correctly without ever materializing the full $s\times s$ score matrix. Both activation-memory terms then scale as $s/\text{CP}$.

The key win — communication fully hidden. The KV send/recv for step $k+1$ is launched while the blockwise attention for step $k$ computes. As long as (compute per block) ≥ (time to ship one KV block), the ring communication is completely overlapped — Ring Attention adds no wall-clock overhead vs. single-device, up to $N\times$ longer sequences. (Same overlap philosophy as FSDP prefetch, applied to the sequence dimension.)

Causal-mask load imbalance. With a causal mask, later query blocks attend to more KV blocks than earlier ones, so a naive ring leaves some GPUs idle. Striped Attention (Brandon et al., 2023) [recent] permutes tokens so each GPU gets a balanced early/late position mix, restoring even load. Production CP (Megatron-Core, Transformer Engine) uses this kind of balancing.

In practice: Llama 3 kept CP=1 at 8K context and raised it to CP=16 only for the 128K long-context phase — you pay for CP only when $s$ is large. CP composes with SP (both split the sequence).

Why it matters

CP is the standard axis for long-context training today, turning the quadratic attention memory/compute of million-token sequences into a tractable, overlap-hidden ring. It slots into N-D parallelism between TP and PP in the rank ordering, activated only for the long-context phase of a run.

Taught in

• 10-model-parallelism — §6 Ring Attention.
• 11-sota-large-scale-training — §4 CP in the Llama 3 config.

See also

• flash-attention — blockwise online softmax that makes the ring correct
• online-softmax — the accumulation primitive
• sequence-parallelism — the sibling sequence-splitting axis
• tensor-parallelism / pipeline-parallelism — the other model-parallel axes
• 3d-parallelism — CP in the rank ordering
• gpu-systems-for-llms