Context parallelism — splitting the sequence across GPUs

Yesterday we explored activation checkpointing: saving GPU memory by recomputing intermediate activations during backward. Before that, FSDP and ZeRO showed how to shard training state instead of replicating it.

Both are powerful, but neither fully solves a different problem: what happens when the sequence itself becomes enormous?

Today's idea is context parallelism (CP).

Instead of placing the entire 32,000- or 128,000-token sequence on one GPU, give each GPU a portion of the tokens. Then coordinate attention so every token can still see the history it needs.

The catch is that Transformer attention connects positions across the sequence. Splitting tokens is easy; preserving those connections efficiently is the hard part.

Why long contexts stress training memory¶

Consider the hidden states entering one Transformer layer. For microbatch size \(B\), sequence length \(T\), hidden width \(d\), and \(b\) bytes per element, their storage is

\[ M_{\mathrm{hidden}}=BTdb. \]

Double the sequence length and this tensor doubles. During training, backward may require many such tensors or other saved intermediates across layers.

Attention introduces a second challenge. For each query position \(i\), the model computes scores against eligible key positions \(j\):

\[ s_{ij}=\frac{q_i^\top k_j}{\sqrt{d_h}}, \qquad o_i=\sum_{j\le i} \frac{e^{s_{ij}}}{\sum_{u\le i}e^{s_{iu}}}v_j. \]

For full causal attention, the total score computation grows approximately as \(O(T^2)\).

FlashAttention avoids storing the entire \(T\times T\) score matrix in GPU memory, but it does not eliminate the quadratic attention computation. Activation checkpointing reduces saved intermediates, but the current computation still needs somewhere to run.

Context parallelism distributes that work and its sequence-shaped tensors across devices.

What exactly gets sharded?¶

Imagine a 32,768-token training sequence and four GPUs. A simple contiguous partition gives each GPU 8,192 token positions:

Full sequence: [        32,768 tokens         ]

GPU 0:         [  0–8191   ]
GPU 1:                    [ 8192–16383 ]
GPU 2:                                   [ 16384–24575 ]
GPU 3:                                                    [ 24576–32767 ]

Each GPU owns the queries and output activations for its local positions.

But a query on GPU 3 may attend to keys and values on all four GPUs. A query on GPU 0 needs only earlier positions in its own segment.

So this is not four independent 8,192-token attention problems. The model must still produce the same attention result it would have produced for the unsplit sequence.

The core trick: circulate keys and values¶

One implementation family uses ring attention. Each GPU keeps its local queries while blocks of keys and values circulate between GPUs.

Step 1: each GPU computes attention against its local K,V
Step 2: each GPU receives another GPU's K,V block
Step 3: compute another contribution to local queries
Step 4: repeat until all relevant K,V blocks are covered

A GPU does not need to collect all attention scores into one giant matrix. It can accumulate a softmax result incrementally as key/value blocks arrive.

This is the same IO-aware principle behind FlashAttention, now extended across GPUs: move blocks and combine partial results rather than materializing the entire attention matrix.

Other context-parallel implementations use different communication patterns, including all-to-all exchanges. The ring is a useful mental model, not a requirement of CP.

How do partial softmax results combine?¶

This is the most interesting mathematical detail.

For one query, suppose GPU \(r\) processes a block of eligible keys. Define

\[ m_r=\max_{j\in r}s_j, \qquad \ell_r=\sum_{j\in r}e^{s_j-m_r}, \qquad z_r=\sum_{j\in r}e^{s_j-m_r}v_j. \]

Here \(m_r\) is the block's largest score, \(\ell_r\) is its normalized exponential sum, and \(z_r\) is its weighted value sum.

After processing all relevant blocks, set

\[ m=\max_r m_r. \]

The exact full-attention output, apart from floating-point rounding, is

\[ \boxed{ o= \frac{\sum_r e^{m_r-m}z_r} {\sum_r e^{m_r-m}\ell_r} }. \]

The rescaling factors put every block on the same softmax scale. Without them, independently normalized block outputs could not simply be averaged.

Causal masking still applies: future positions must contribute nothing. The algebra changes how we compute attention, not which tokens the model is allowed to attend to.

Worked example: four GPUs and a 32K context¶

Consider a BF16 Transformer with

\[ B=1,\quad T=32768,\quad d=4096,\quad H_{kv}=8,\quad d_h=128. \]

One full hidden-state tensor occupies

\[ M_{\mathrm{hidden}} =1(32768)(4096)(2) =256\text{ MiB}. \]

With four-way context parallelism, each GPU owns about

\[ 256/4=64\text{ MiB} \]

of that tensor.

For a 32-layer model, retaining just one such full-sequence hidden-state tensor per layer would amount to 8 GiB. Sharding those tensors four ways reduces this illustrative category to 2 GiB per GPU, before checkpointing and other memory costs.

Now consider the keys and values at one layer:

\[ M_{KV}=2TH_{kv}d_hb =2(32768)(8)(128)(2) =128\text{ MiB}. \]

Each GPU initially holds a 32 MiB KV shard. In a simple four-GPU ring that visits every block, it receives the other three shards: approximately

\[ 3\times32=96\text{ MiB} \]

of KV traffic per GPU for that layer's forward attention.

At a hypothetical effective transfer bandwidth of 100 GB/s, transferring 96 MiB alone takes about 1 ms. That is not the total layer time: synchronization, computation, overlap, and backward communication still matter. A causal-aware implementation may skip future blocks and transfer less.

The tradeoff is clear: less persistent sequence memory per GPU, but more inter-GPU communication.

Why causal attention complicates load balancing¶

Suppose we use the contiguous partition above. Queries near the beginning of the sequence have little history to examine; queries near the end have much more.

GPU 0 queries: short histories
GPU 1 queries: longer histories
GPU 2 queries: still longer histories
GPU 3 queries: longest histories

If each GPU owns one contiguous quarter, GPU 3 may do much more attention work than GPU 0. The whole distributed operation can end up waiting for the busiest worker.

One remedy is to distribute positions in striped or zigzag chunks, so each GPU receives a mixture of early and late positions. This balances causal attention work better, although it complicates position bookkeeping and communication.

A perfectly even number of tokens per GPU is therefore not necessarily an even amount of computation.

How CP differs from the parallelism we've learned¶

Technique	What is distributed?	Main coordination cost
Tensor parallelism	Weight-matrix computation within a layer	Activation collectives
Pipeline parallelism	Consecutive model layers	Activations crossing stages
FSDP / ZeRO-3	Persistent training state	Parameter gathers and gradient scattering
Context parallelism	Token positions within a sequence	Attention information across sequence shards

These can be combined. A large training system might shard parameters with FSDP, split a layer's matrix operations with tensor parallelism, and distribute a very long sequence with context parallelism.

But each additional dimension introduces communication. The topology of the GPU interconnect and the order of collectives become important design constraints.

What about inference?¶

Context parallelism can help with very long prompt prefill, where many token positions are available for parallel computation.

Autoregressive decode is different. A request usually adds just one query position at a time. Distributing its historical KV cache across GPUs can increase memory capacity, but the new query still needs information from remote KV shards. Communication may outweigh the benefit for shorter contexts.

This connects to our earlier lessons: prefill and decode have different bottlenecks, so the best parallelism strategy for one phase may not be the best for the other.

Mental model: a research team reading one enormous book¶

Imagine four researchers assigned different quarters of a 1,000-page book. Each researcher is responsible for writing notes about their assigned pages.

But to interpret page 900, the fourth researcher may need facts introduced on pages 50, 300, and 600. The researchers therefore circulate relevant source material rather than pretending their quarters are independent books.

Context parallelism divides responsibility for the pages. Ring attention circulates the reference material. Online softmax combines evidence from different batches without storing every individual comparison.

The team saves desk space and shares the reading workload, but now coordination speed matters.

The deeper lesson¶

FSDP reduces replicated model state. Activation checkpointing reduces saved intermediate state. Context parallelism reduces how much of one long sequence must live and be processed on each GPU.

\[ \boxed{\text{Long-context training is a data-placement problem as well as a compute problem.}} \]

Context parallelism does not make full attention's quadratic work disappear. It distributes that work while introducing communication and load-balancing challenges.

The practical question is: At what sequence length does dividing attention across GPUs save enough memory and computation time to justify moving keys, values, and gradients between them?