Ring Attention and Context Parallelism: How Long-Context LLMs Are Trained Across GPUs

A single GPU cannot hold the activations of a one-million-token sequence, no matter how clever your kernel is. FlashAttention removed the quadratic attention-score matrix from memory, but it did not remove the linear terms: the keys, values, hidden states and checkpointed activations of every layer still scale with sequence length, and at 128K to 1M tokens they overflow an 80 GB device. Ring attention is the idea that breaks this ceiling. It shards one sequence across many GPUs and rotates key-value blocks around a ring, so attention stays exact while per-device memory shrinks with the number of devices.

This matters now because long context stopped being a research stunt. Frontier models ship with 128K to multi-million-token windows, and every one of them had to be trained, extended and served somehow. The answer in production is a family of techniques usually called context parallelism or sequence parallelism.

What this covers: the memory wall, the ring mechanism, DeepSpeed-Ulysses, the USP hybrid, causal load balancing, communication math with worked numbers, how Llama 3, Megatron-Core and PyTorch implement it, how it combines with tensor, pipeline and FSDP sharding, inference-time use, and the failure modes that bite in practice.

Context and Background

Self-attention has two costs that scale differently. Compute grows as O(S²) in sequence length S, and the naive implementation also materialises an S×S score matrix per head. FlashAttention fixed the second problem by tiling the computation so scores live only in on-chip SRAM, using an online softmax that keeps a running maximum and normaliser. Memory for attention itself becomes O(S) instead of O(S²). That is the foundation everything below builds on.

But O(S) per layer, summed over dozens of layers and multiplied by batch, is still a large number. Three things grow linearly with S and cannot be tiled away on one device: the hidden-state activations kept for the backward pass, the K and V tensors consumed by attention, and, at inference, the persistent KV cache. Activation recomputation trades FLOPs for memory, but you still need at least the layer inputs.

Before sequence sharding, the standard toolkit was tensor parallelism (TP), pipeline parallelism (PP) and data parallelism with sharded optimizer state (ZeRO or FSDP). None of them splits the sequence. TP splits weight matrices and attention heads, so activations per GPU shrink only partially. PP splits layers but each stage still holds full-length activations. FSDP shards parameters, gradients and optimizer states, not activations. When the sequence itself is the problem, you need a fourth axis.

The research line is short. Liu, Zaharia and Abbeel proposed Ring Attention with Blockwise Transformers in 2023, distributing a sequence across devices and overlapping the communication of key-value blocks with blockwise attention computation. Jacobs et al. published DeepSpeed-Ulysses the same autumn, sharding the sequence and using all-to-all collectives to switch between sequence-sharded and head-sharded layouts. Brandon et al. followed with Striped Attention, fixing a causal-mask imbalance in the ring. Fang and Zhao unified the two camps in USP. Framework teams then absorbed these ideas: Megatron-Core exposes a context-parallel size, PyTorch ships an experimental Context Parallel API, and Meta’s Llama 3 report describes a production variant.

Terminology is messy. “Sequence parallelism” originally meant something narrower in Megatron (sharding LayerNorm and dropout activations along the sequence within a TP group), while the long-context literature uses it for splitting attention itself. “Context parallelism” is the cleaner term for the latter, and this article uses it, noting the overlap where it matters. For how models evaluate such long windows afterwards, see long-context benchmarks like RULER.

The Memory Wall at 128K to 1M Tokens

Concrete numbers make the problem tangible. The figures below are illustrative arithmetic for a Llama-3-8B-shaped model (32 layers, hidden size 4096, 32 query heads, 8 KV heads, head dimension 128, bf16). They are derived from those shapes, not measured benchmarks.

KV tensors. One token’s K and V in one layer occupy 2 × 8 heads × 128 × 2 bytes = 4,096 bytes. Across 32 layers that is 128 KiB per token. At 128K tokens (131,072) the KV state is 16 GiB; at 1M tokens (1,048,576) it is 128 GiB. That already exceeds one 80 GB GPU, before weights.

Checkpointed activations. With full activation recomputation you still store each layer’s input: 4096 × 2 bytes = 8 KiB per token per layer, or 256 KiB across 32 layers. At 1M tokens that is 256 GiB. At 128K it is 32 GiB. Add the temporary working set inside one layer, such as the MLP intermediate (14,336 wide in Llama-3-8B, which is 28 KiB per token in bf16), and a single layer’s recompute needs about 28 GiB at 1M tokens for that tensor alone.

Attention compute. FLOPs for one attention layer forward are about 4 × S² × (heads × head_dim) for the two matrix multiplications QKᵀ and PV, halved by the causal mask. At S = 1M, hidden 4096: 4 × (1.05e6)² × 4096 ≈ 1.8e16 FLOPs per layer, 9e15 with causal skipping. That is roughly 9 PFLOP per layer forward, hundreds of PFLOP across layers, and multiplied by roughly 3.5 for forward plus backward with recomputation. On one GPU sustaining 400 TFLOP/s, a single layer’s forward attention takes about 22 seconds, and a full forward and backward pass over all 32 layers lands around 40 minutes for one sequence. Compute alone argues for parallelism.

So the wall has two faces: memory (linear terms overflow one device) and time (quadratic compute is too slow on one device). Sharding the sequence across N GPUs divides both by roughly N. The cost is that attention is the one operation where every query needs every key, so some data must cross the network. The entire design space is about how to move that data and hide its cost.

Figure 1: A long sequence is split along the sequence axis. Everything except attention is token-local and needs no communication; attention requires K and V from other ranks.

The figure shows the key structural fact: most of a transformer is per-token. LayerNorm, the MLP, the QKV projections, residual connections, and rotary embeddings (given correct position indices) act on each token independently. Sharding the sequence is therefore free for those layers. Only the attention score computation couples tokens, and that is where ring attention, Ulysses and all-gather schemes differ.

Core Architecture: Three Ways to Share K and V

Context parallelism splits the sequence across a group of N GPUs and uses one of three communication patterns to let every query see every key: a ring of point-to-point transfers (ring attention), an all-gather of K and V, or an all-to-all that re-shards from sequence to heads (Ulysses). All three compute exact attention. They differ in communication volume, overlap, mask flexibility and scaling limits.

Ring Attention: Pass KV Blocks Around the Ring

Each rank holds one contiguous block of Q, K and V: S/N tokens each. Rank r keeps its Q block fixed and computes attention of Q_r against whichever KV block it currently holds, accumulating the result with the online-softmax rescaling trick from FlashAttention. Meanwhile it sends its current KV block to the next rank and receives a new one from the previous rank. After N−1 rotations every query block has met every KV block.

The correctness relies on a property FlashAttention already exploits: softmax attention over a concatenation of key blocks can be computed block by block, provided you carry the running row maximum m, the running denominator l, and the unnormalised output o. When a new block arrives, you compute local scores, update m, rescale the old o and l by exp(m_old − m_new), and add the new contribution. The order of blocks does not matter mathematically, which is why a ring (arbitrary rotation order) works and why the result is exact up to floating-point reassociation. Liu et al. state the method makes no approximation and, with overlap, adds no extra cost relative to standard blockwise attention.

Overlap is the performance idea. While the GPU runs the attention kernel on block j, the NIC or NVLink engine is already moving block j+1. If compute time per step exceeds transfer time per step, communication is fully hidden. The next section quantifies when that holds.

Figure 2: Ring attention. Queries stay put; KV blocks hop one rank per step while the attention kernel runs on the block already received.

In the backward pass the ring runs again, but now dK and dV partials must travel with the KV blocks (or be accumulated in a second ring), roughly doubling communication relative to forward while compute grows by about 2.5x, so the overlap condition stays comfortable.

DeepSpeed-Ulysses: All-to-All and Head Sharding

Ulysses starts with the same sequence-sharded layout but takes a different route. Before attention, an all-to-all on Q, K and V converts each tensor from [S/N tokens, all heads] on each rank to [S tokens, H/N heads]. Each GPU now holds the full sequence for a subset of heads and runs completely standard, unmodified FlashAttention on them, because heads are independent. A second all-to-all on the output converts back to sequence-sharded.

Three properties follow. First, there is no ring bookkeeping and no special masking logic: any mask works because each GPU sees the whole sequence for its heads, and any optimised kernel works. Second, communication per GPU shrinks as N grows. The Ulysses authors argue the per-device communication volume stays constant when sequence length and device count increase proportionally, and they report about 2.5x faster training with 4x longer sequences than the previous state of the art in their experiments. Third, the parallel degree is capped by the number of heads: N must divide the head count. With grouped-query attention and 8 KV heads you cannot exceed 8-way Ulysses without replicating KV heads, which wastes memory and bandwidth.

Figure 3: Ulysses attention swaps the sharded dimension twice per layer, so the attention kernel itself needs no distributed logic.

All-Gather KV and the Hybrid USP Design

A third pattern simply all-gathers K and V (not Q) onto every rank, then computes local Q against the full KV. It sacrifices peak memory for simplicity: the gathered KV is a full-sequence tensor. Under GQA, KV is small, so this is cheaper than it sounds. PyTorch’s experimental Context Parallel API uses all-gather-based pass-KV as its default rotation and offers an all-to-all variant selectable with a setter; Megatron-Core describes its all-gather and reduce-scatter as implemented via point-to-point operations in a ring.

USP, from Fang and Zhao, observes that Ulysses and ring attention have complementary weaknesses. Ulysses is head-limited and wants a fast fabric because all-to-all is bandwidth-hungry and topology-sensitive. Ring attention scales with sequence length and tolerates slower links because transfers are point-to-point and overlappable. USP arranges the CP group as a two-dimensional grid: Ulysses degree u inside a fast domain (typically an NVLink node), ring degree r across slower links, with total degree u × r. The authors report training LLaMA3-8B at a 208K sequence length on two 8-GPU A800 nodes at 47% MFU. I did not independently reproduce that figure; it is the paper’s own claim.

The rule of thumb that emerges: use Ulysses where bandwidth is abundant and heads are plentiful, ring where you must go wide or cross slow links, and combine them when you need both.

Deeper Analysis: Communication Math, Causal Balance, and Real Implementations

When does the ring actually hide its communication?

The overlap condition is a simple inequality: the time to run attention on one KV block must exceed the time to move the next block. Take one ring step with b tokens per rank, h_q query heads, n_kv KV heads and head dimension d. The attention FLOPs are 4·b²·h_q·d (two matmuls, each 2·b²·d per head). The bytes moved are K plus V for one block: 2 tensors × b × n_kv × d × 2 bytes = 4·b·n_kv·d. Divide one by the other and the heads and dimension cancel:

FLOPs per byte = b × (h_q / n_kv).

The ring is compute-bound whenever that exceeds the machine’s FLOPs-to-bandwidth ratio. Use illustrative hardware figures: an H100-class GPU with a dense bf16 peak near 990 TFLOP/s of which we assume 400 TFLOP/s is achieved in attention kernels, an InfiniBand NDR link at 400 Gb/s (50 GB/s per direction), and NVLink at roughly 450 GB/s per direction. These are round numbers for reasoning, not measurements.

• Across InfiniBand: 400e12 / 50e9 = 8,000 FLOP per byte needed.
• Inside an NVLink node: 400e12 / 450e9 ≈ 890 FLOP per byte needed.

With Llama-3-style GQA (32 query heads over 8 KV heads, ratio 4), the ring is compute-bound over InfiniBand when 4·b ≥ 8,000, so b ≥ 2,000 tokens per rank. Over NVLink, b ≥ about 220. With plain multi-head attention (ratio 1) the InfiniBand threshold jumps to 8,000 tokens per rank. This is why GQA and MQA quietly make context parallelism cheaper: Megatron’s documentation makes the same qualitative point, noting that fewer KV heads reduce communication volume.

Now plug in the worked case. A 1M-token sequence over N = 8 ranks gives b = 131,072. One step’s attention is 4 × 131,072² × 32 × 128 ≈ 2.8e14 FLOPs, about 0.70 s at 400 TFLOP/s. The KV block is 131,072 × 4,096 bytes = 537 MB, which takes about 11 ms at 50 GB/s. Compute outweighs communication by roughly 65 to 1. The ring could widen until b falls near 2,000, which for 1M tokens is around 512 ranks, before communication starts to show. In practice other costs (stragglers, collectives for loss and gradients, launch latency) bite long before that, but the fundamental arithmetic is generous.

The same calculation exposes where ring attention hurts: short sequences. At S = 16K over 8 ranks, b = 2,048 and you sit right on the break-even line over InfiniBand with GQA. Context parallelism is a tool for long sequences; using it for 8K-token training mostly adds overhead.

Ulysses volume versus ring volume

Ulysses moves different quantities. For the same 1M-token, N = 8 example, each rank holds 131,072 tokens. Q is 8,192 bytes per token (32 heads × 128 × 2), so the Q all-to-all moves about 1.07 GB; K and V are 2,048 bytes per token each, about 268 MB apiece; the output all-to-all moves another 1.07 GB. Each all-to-all sends (N−1)/N of the local tensor off-rank, so total forward traffic per layer per rank is about 2.68 GB × 7/8 ≈ 2.35 GB. The ring’s forward traffic is 7 × 537 MB ≈ 3.76 GB per rank.

The scaling behaviour differs more than the totals. Ring volume per rank stays near S × (KV bytes per token) regardless of N, because every rank must eventually receive every other rank’s KV. Ulysses volume per rank is proportional to S/N, so it falls as you add GPUs, which is the constant-communication-under-proportional-scaling claim in the paper. The catch is that all-to-all is exposed (not overlapped with the attention kernel in the basic design) and stresses bisection bandwidth, so it prefers NVLink or NVSwitch domains and degrades over oversubscribed networks. The 2.35 GB above takes roughly 5 ms at 450 GB/s on NVLink but over 45 ms at 50 GB/s on InfiniBand, against several seconds of compute per layer at this length, so even exposed it is tolerable at 1M tokens; at 32K tokens the same fraction becomes painful.

Causal masks and load balancing

Language-model training uses a causal mask: token i attends only to tokens ≤ i. With contiguous sharding, rank 0 owns the earliest tokens and, in ring step k, finds most of the blocks it receives are entirely in the future and fully masked. The last rank has to compute every block. Work per rank grows linearly with rank index, 1 to N block-units, while the step time is set by the slowest rank.

Quantify: for N = 8 the total work is 36 block-units (1+2+…+8), so the average is 4.5, yet the critical path is 8. Utilisation is (N+1)/(2N) = 56%. The ideal gain from perfect balance is 2N/(N+1) = 1.78x for the attention part at N = 8. Skipping fully masked blocks helps compute but leaves ranks idle waiting.

Striped Attention fixes this by interleaving: rank r takes tokens r, r+N, r+2N, and so on, so every rank holds tokens spread across the whole sequence and every block has roughly the same triangular shape. The authors report up to 1.45x end-to-end training throughput over Ring Attention at 256K sequence length on A100 GPUs, and 1.65x at 786K on 16 TPUv4 chips. Those are whole-model figures, lower than the attention-only ideal because MLPs and projections are already balanced.

Production systems use a close cousin, often called zigzag or load-balanced sharding. Split the sequence into 2N chunks and give rank i chunk i and chunk 2N−1−i. Rank 0 holds the first and last chunks, rank N−1 holds two middle chunks. Counting causal work per chunk (chunk c attends c+1 chunk-blocks including its diagonal), every rank gets (i+1) + (2N−i) = 2N+1 units, perfectly equal. Llama 3’s report describes exactly this: the input is cut into 2 × CP chunks and the i-th and (2×CP − 1 − i)-th go to rank i. Zigzag has an advantage over fine striping: chunks stay contiguous, which keeps kernels efficient and mask logic simple.

Llama 3: why Meta chose all-gather over a ring

The Llama 3 paper (Meta, 2024) documents the most widely cited production recipe. For long-context pre-training with 131,072-token sequences, it uses context parallelism on top of tensor, pipeline and FSDP data parallelism. The notable choice is that it all-gathers K and V rather than using a ring that overlaps communication with compute. The authors give two reasons. An all-gather handles irregular masks such as the document mask (needed when several documents are packed into one sequence) more easily. And because GQA makes K and V much smaller than Q, the exposed all-gather latency is small, while attention compute grows quadratically against linearly growing communication.

The paper’s configuration table lists a 16,384-GPU run with TP = 8, CP = 16, PP = 16, DP = 8 (product 16,384), sequence length 131,072, 16M tokens per batch, 380 TFLOPs per GPU and 38% BF16 MFU. It also states the dimensions are ordered [TP, CP, PP, DP] by network demand: TP is most bandwidth-hungry and stays within a server, DP is outermost because it tolerates latency through asynchronous prefetching. That ordering is a good default for your own mesh layout.

The lesson is not that rings are bad. It is that with GQA and a heavily masked, packed-document workload, the simpler collective won. The ring’s strengths show up when KV is large (MHA) or when you run at extreme lengths over slow links.

A runnable simulation of the ring

The following NumPy code reproduces ring attention’s arithmetic on one machine: each simulated rank keeps its Q block and processes KV blocks in ring order using the online-softmax update, then the result is compared with ordinary full attention. It is non-causal for clarity.

import numpy as np

def full_attention(Q, K, V):
    s = Q @ K.T / np.sqrt(Q.shape[1])
    s -= s.max(axis=1, keepdims=True)
    p = np.exp(s)
    return (p / p.sum(axis=1, keepdims=True)) @ V

def ring_attention_sim(Q, K, V, n):
    S, d = Q.shape
    b = S // n
    out = np.zeros_like(Q)
    for r in range(n):                       # one iteration per rank
        q = Q[r*b:(r+1)*b]
        m = np.full(b, -np.inf)              # running row max
        l = np.zeros(b)                      # running denominator
        o = np.zeros((b, d))                 # running unnormalised output
        for step in range(n):                # ring rotation
            src = (r - step) % n             # whose KV block we hold now
            k = K[src*b:(src+1)*b]
            v = V[src*b:(src+1)*b]
            s = q @ k.T / np.sqrt(d)
            m_new = np.maximum(m, s.max(axis=1))
            p = np.exp(s - m_new[:, None])
            scale = np.exp(m - m_new)        # rescale old accumulators
            l = l * scale + p.sum(axis=1)
            o = o * scale[:, None] + p @ v
            m = m_new
        out[r*b:(r+1)*b] = o / l[:, None]
    return out

rng = np.random.default_rng(0)
S, d, n = 512, 64, 8
Q, K, V = (rng.standard_normal((S, d)) for _ in range(3))
err = np.abs(ring_attention_sim(Q, K, V, n) - full_attention(Q, K, V)).max()
print("max abs error:", err)   # ~1e-15 in float64

The error sits at floating-point noise level, which demonstrates the “exact attention” claim. Real implementations keep the accumulators (m, l, o) in float32 even when tensors are bf16, because repeated rescaling across many blocks otherwise accumulates error.

The load-balanced token assignment is equally small:

def zigzag_indices(S, N, rank):
    """Token indices owned by `rank` under 2N-chunk zigzag sharding."""
    c = S // (2 * N)
    first = range(rank * c, (rank + 1) * c)
    j = 2 * N - 1 - rank
    second = range(j * c, (j + 1) * c)
    return list(first) + list(second)

Using PyTorch’s Context Parallel API

PyTorch 2.7 and later include an experimental context-parallel API under torch.distributed.tensor.experimental. It shards tensors you pass as buffers along a chosen sequence dimension, in place, and replaces scaled-dot-product attention with a ring implementation inside the context manager. The tutorial’s pattern is:

import torch, torch.nn.functional as F
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from torch.nn.attention import sdpa_kernel, SDPBackend
from torch.distributed.tensor.experimental import context_parallel
from torch.distributed.tensor.experimental._attention import context_parallel_unshard

world = dist.get_world_size()
mesh = init_device_mesh("cuda", (world,), mesh_dim_names=("cp",))
q = torch.randn(2, 8, 65536, 128, device="cuda", dtype=torch.bfloat16)
k, v = torch.randn_like(q), torch.randn_like(q)

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    with context_parallel(mesh, buffers=(q, k, v), buffer_seq_dims=(2, 2, 2)):
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    (out,) = context_parallel_unshard(mesh, [out], [2])

Two caveats from the documentation: tensors that depend on position along the sequence, such as Llama 3’s rotary frequency table, must also be passed in buffers, or the rotary embedding will be wrong on every rank except the first; and the output stays sharded until you unshard it. The module is prototype, and underscore-prefixed imports signal the API may change, so pin your PyTorch version. The default rotation is all-gather based pass-KV, with an all-to-all variant available via set_rotate_method.

Megatron-Core and the 4D mesh

In Megatron-Core you enable context parallelism with --context-parallel-size (the config field context_parallel_size; default 1, which disables it). The documentation states the GPU count equals TP × CP × PP × DP, that activation memory per GPU shrinks by the CP degree, and that it works with MHA, MQA and GQA and with unidirectional or bidirectional masks. It lists Megatron Core 0.5.0 and Transformer Engine 1.1 as the minimum versions, though those requirements may have moved on, so check your release notes. I also recall newer releases adding a configurable CP communication type (point-to-point, all-gather, all-to-all and a hierarchical variant); I did not verify that against current docs, so treat it as something to confirm.

A worked mesh makes the interactions concrete. Suppose 64 GPUs on eight 8-GPU nodes, training a 70B-class model at 128K sequence length. A plausible layout is TP = 8 inside each node, CP = 2, PP = 2, DP = 2 (8×2×2×2 = 64). Each CP rank holds 64K tokens. TP all-reduces stay on NVLink, CP rings or all-gathers cross node boundaries every layer, PP sends activations only at stage boundaries, and DP/FSDP gathers happen once per layer across the outermost dimension. This is an illustrative layout, not a tuned recommendation; the right split depends on measured throughput.

Figure 4: A decision flow for choosing a context-parallel scheme from model shape and network topology.

Interaction with tensor parallelism and FSDP

TP and Ulysses both consume attention heads: TP splits heads across its group, and Ulysses then splits the remaining heads across the CP group. If TP = 8 on a model with 8 KV heads, there is one KV head per TP rank and no heads left for Ulysses; you would have to use a ring or all-gather scheme for CP. This head budget is the single most common reason Ulysses stops being available in practice. USP’s hybrid sidesteps it by capping the Ulysses degree and filling the remainder with ring degree.

FSDP composes more easily. CP ranks hold identical weights because they process different slices of the same sequence, so parameter sharding and gradient reduction must span the CP dimension as well as the data-parallel dimension: gradients are computed from different token slices and need to be summed across CP ranks. Frameworks usually fold the CP and DP axes together for the gradient reduction and for FSDP all-gathers, which means FSDP’s parameter all-gather now crosses more ranks. That is real extra traffic, and it is why the ordering [TP, CP, PP, DP] puts CP inside the DP axis.

Inference: context parallelism for prefill

Training is not the only consumer. The prefill of a 1M-token prompt has the same quadratic compute problem. Yang et al. (Meta) describe context parallelism for million-token inference with two lossless ring variants, pass-KV and pass-Q, covering full prefill, prefill with a persistent KV cache, and decode. They report scaling up to 128 H100 GPUs across 16 nodes, with a 1M-token prefill of Llama 3 405B completing in 77 seconds at 93% parallelisation efficiency and 63% FLOPS utilisation, and a 128K-token prefill in 3.8 seconds. They also report similar scaling over RDMA and TCP, which suggests the approach tolerates modest interconnects.

The reason for two variants follows the same bytes-versus-FLOPs logic: pass-KV sends keys and values around the ring, which is cheap when queries are numerous; pass-Q sends queries instead, which wins when only a few new query tokens attend to a large cached prefix. That is my reading of the design; the paper’s text should be consulted for the exact selection heuristic. Decode is different again: one query token per step means almost no compute to hide communication behind, so latency, not throughput, dominates. For the broader inference-side memory picture, see KV cache optimisation and multi-head latent attention versus GQA and MQA, both of which shrink the KV bytes that context parallelism has to move.

Trade-offs, Gotchas, and What Goes Wrong

Context parallelism is simple to describe and unforgiving to deploy. The failures below are the ones that cost real debugging time.

Position handling. Sharding tokens does not shard the meaning of their positions. Rotary embeddings, learned position tables and any mask built from absolute indices must use global positions on each rank. A rank that restarts positions at zero produces a model that trains without errors and quietly learns garbage. PyTorch’s documentation calls this out explicitly for the rotary frequency table.

Packed documents and masks. Training pipelines concatenate documents and mask attention across boundaries. A ring that skips blocks by position alone cannot know where documents end, so irregular masks can break block-skipping logic and re-introduce imbalance. This is one of the two reasons Llama 3 chose all-gather. Ulysses handles any mask because each rank sees the full sequence for its heads, but pays in head-count limits.

Head budget exhaustion. Ulysses degree times TP degree cannot exceed the number of (KV) heads without replication. With GQA at 8 KV heads, a TP of 8 leaves nothing. Replicating KV heads works but multiplies the K and V memory and traffic that GQA was designed to save.

Short sequences. Below a few thousand tokens per rank the ring falls under the compute-versus-bandwidth break-even computed earlier. Teams that enable CP globally for a mixed-length curriculum often see short-sequence stages slow down. Make CP size a per-stage setting, and fold CP back into DP when sequences are short.

Numerics. Ring attention reassociates floating-point sums, so results differ from single-device attention at the last bits and across different CP sizes. Keep log-sum-exp and accumulators in fp32. Do not expect bitwise reproducibility when you change the CP degree, which complicates loss-curve comparisons and checkpoint resumption.

Loss and gradient reduction. Each CP rank computes loss over its own tokens. Normalisation must divide by the global token count, and gradients must be reduced over the CP axis. Averaging per-rank means over unequal valid-token counts (padding, loss masks) biases the objective.

Stragglers and collectives. A ring is a chain of dependencies; one slow GPU or one congested link stalls every step. Mismatched send and receive ordering in point-to-point code can deadlock. All-to-all schemes suffer instead from oversubscribed spine links.

Memory spikes. All-gather designs transiently materialise full-sequence K and V: about 4.3 GB per layer for the 1M-token example (1,048,576 tokens × 4,096 bytes). That is fine once, but not if allocator fragmentation or overlapping prefetch keeps several layers’ buffers alive.

Sparse and sliding-window attention. Local attention patterns need only neighbours’ KV, so a full ring wastes communication. Windowed designs want halo exchange instead, and they break the symmetry that zigzag balancing assumes.

Evaluation gap. Training at 1M tokens does not mean the model uses 1M tokens well. Effective context is usually shorter than advertised, as the RULER-style evaluations show. Systems work solves feasibility, not quality.

Practical Recommendations

Start from the constraint, not the technique. If the sequence fits in memory with FSDP plus TP and activation recomputation, do not add CP. If it does not, calculate the per-rank block size S/N and compare 4·(S/N) (or S/N times your query-to-KV head ratio) with the FLOPs-to-bandwidth ratio of your slowest link. That one inequality predicts whether a ring will hide its communication.

Choose the scheme by topology and head budget. Inside one NVLink node with spare heads, Ulysses is the least intrusive option, because the attention kernel is untouched. Across nodes, or when heads are exhausted by TP, use a ring or an all-gather scheme. When you need both, use a USP-style hybrid. For causal training, always use zigzag or striped sharding; contiguous sharding leaves up to 44% of attention capacity idle at N = 8.

Order the mesh by bandwidth demand: TP innermost, then CP, then PP, then DP, as Llama 3 does. Keep CP groups within as few nodes as possible. Make CP size a per-stage parameter so short-sequence stages do not pay for it.

Prefer GQA or MLA when you control the architecture. Fewer KV bytes make every scheme here cheaper.

A pre-flight checklist:

• Confirm global position ids and rotary tables are sharded consistently with the tokens.
• Verify the loss is normalised by the global valid-token count and gradients reduce over the CP axis.
• Run a single-device versus CP equivalence test on a small model and tolerate only fp32-level differences.
• Profile one step and confirm attention kernels overlap with send/recv, not just that throughput looks acceptable.
• Check that Ulysses degree divides the KV head count after TP.
• Pin framework versions, since the PyTorch API is prototype and Megatron options evolve.
• Test document-mask behaviour with packed samples before long runs.

Frequently Asked Questions

What is ring attention?

Ring attention is an exact, distributed form of attention that shards a long sequence across GPUs arranged in a ring. Each GPU keeps its query block while key-value blocks rotate from neighbour to neighbour, and each GPU accumulates results with the online-softmax update from FlashAttention. Communication overlaps with computation, so per-device memory falls roughly in proportion to device count without approximating attention, as proposed by Liu, Zaharia and Abbeel in 2023.

What is the difference between context parallelism and sequence parallelism?

The terms overlap. In Megatron’s original usage, sequence parallelism shards LayerNorm and dropout activations along the sequence inside a tensor-parallel group, leaving attention untouched. Context parallelism shards the entire sequence, including attention, across a separate group of GPUs and adds communication for K and V. In the long-context literature, “sequence parallelism” often means the latter, which is why ring attention, Ulysses and USP are described under either name.

Ring attention or DeepSpeed-Ulysses: which is better?

Neither dominates. Ulysses uses all-to-all, runs an unmodified attention kernel, and has per-GPU communication that shrinks as devices are added, but its degree is capped by head count and it wants fast, high-bisection links. Ring attention scales to many more ranks and tolerates slower links through overlapped point-to-point transfers, but needs more careful masking and load balancing. USP combines them, placing Ulysses inside nodes and the ring across nodes.

Does ring attention approximate attention?

No. The output equals standard softmax attention up to floating-point reassociation. Online softmax lets blocks be processed in any order while carrying a running maximum, denominator and output, so the final normalised result matches full attention. The simulation earlier in this article shows an error near float64 round-off. Differences you may observe across CP sizes come from summation order and bf16 rounding, not from algorithmic approximation.

How does Llama 3 handle long context during training?

Meta’s Llama 3 report describes context parallelism on top of tensor parallelism, pipeline parallelism and FSDP data parallelism, ordered TP, CP, PP, DP. It splits sequences into 2×CP chunks for causal load balance and uses an all-gather of K and V instead of a ring, citing easier support for document masks and small exposed latency thanks to GQA. A listed configuration trains 131,072-token sequences on 16,384 GPUs at 38% BF16 MFU.

Can context parallelism speed up inference?

Yes, mainly for prefill. Prefill of a very long prompt is compute-bound and quadratic, so splitting it across nodes reduces time to first token. Meta researchers reported a 1M-token Llama 3 405B prefill in 77 seconds on 128 H100 GPUs. Decode benefits less, because one query token per step leaves little computation to hide communication behind, so decode latency stays dominated by communication and memory reads.

Further Reading

Internal:

• Long-context LLM benchmarks: RULER and effective context length for how to measure whether a long window is actually used.
• GPU memory physics: DRAM vs SRAM and FlashAttention kernel fusion for the tiling and online-softmax foundation.
• Multi-head latent attention vs GQA vs MQA for how KV compression changes communication volume.
• KV cache optimisation for LLM inference for the serving-side memory picture.

External primary sources:

• Liu, Zaharia, Abbeel, Ring Attention with Blockwise Transformers for Near-Infinite Context.
• Jacobs et al., DeepSpeed Ulysses, and Fang and Zhao, USP.
• Brandon et al., Striped Attention; Meta, The Llama 3 Herd of Models; Yang et al., Context Parallelism for Scalable Million-Token Inference.
• Megatron-Core context parallelism documentation and the PyTorch Context Parallel tutorial.

Verification note: paper claims above come from the arXiv abstracts and the Llama 3 paper text fetched on the day of writing. Hardware throughput and bandwidth numbers in worked examples are illustrative assumptions. The Megatron communication-type option and the pass-Q selection heuristic are flagged in the text as not verified.

By Riju — about