Context Parallelism (CP) splits a batch along the sequence-length dimension, allowing each CP rank to process only a chunk of the original sequence and reducing the memory required for training. I use CP to mean any technique that splits the context length and processes it independently, without relying on Tensor Parallelism, to clearly distinguish Sequence Parallelism from Context Parallelism.

Two of my favorite resources motivate CP in terms of long sequences.

stas00/ml-engineering:

Self-Attention, which is the key component of Transformers, suffers from quadratic memory requirements with respect to the sequence length, therefore when sequence length gets to a certain length, even a batch size of 1 might not be able to fit onto a single GPU and require additional partitioning along the sequence dimension. And once this is done, the sequence can be of any length.

The Ultra-Scale Playbook:

Context Parallelism (CP) specifically targets the challenge of training with very long sequences by sharding activations along the sequence dimension across GPUs. […] CP is particularly valuable when scaling to extreme sequence lengths (128k+ tokens) where, even when using full activation recomputation, the memory requirements for attention would be prohibitive on a single GPU.

That is all true, but CP has another substantial benefit: it can balance attention compute across Data Parallel ranks and improve end-to-end training time by reducing the variance of Flash Attention execution time.

Theorem

For a fixed global token batch, increasing the Context Parallel degree together with the micro-batch size reduces the variance of Flash Attention execution time. Ignoring communication overhead, the variance decreases proportionally to 1 / CP

The benefit comes from keeping the global batch size fixed. As CP increases, the number of DP groups decreases, so each remaining group processes proportionally more independent samples. Their attention costs are averaged across CP ranks, reducing variation.

Let’s try to prove this in this article! We will have a computational simulation and a benchmark with a real Flash Attention 3 kernel on an H100.

Load balancing is important.

In distributed training, iteration time is determined by the slowest rank. Even with Tensor, Pipeline, or Context Parallelism, gradients eventually have to be synchronized across Data Parallel groups. If one group reaches an all-gather or reduce-scatter later than the others, everybody else waits.

Stragglers can arise from anywhere: a throttled GPU, a slow network link, or a compute imbalance. Most blocks in a standard transformer scale linearly with sequence length, but attention remains quadratic. Consequently, even small fluctuations in sample length can create severe timing differences across ranks.

Ulysses

DeepSpeed Ulysses is one of the most widely adopted CP algorithms. It starts with the sequence split into CPequal-sized chunks. Before attention, an all-to-all changes the layout from sequence-sharded to head-sharded: each rank receives the full sequence for a subset of attention heads. A reverse all-to-all restores the sequence-sharded layout afterward.

This reduces per-rank activation memory by a factor ofCP, although the exact reduction depends on the attention architecture. With GQA, KV heads may need to be replicated once CP exceeds the number of KV heads.

Theory

Problem Setting

Let:

• CP be the Context Parallel group size;
• N be the total device count, giving N/CP independent DP/CP groups;
• H be the number of attention heads;
• S be the token budget processed by one GPU before CP is applied.

We assume independent samples drawn from a length distribution. As CP increases, we increase the micro-batch size (and hence the token budget of each remaining group) so the global batch remains fixed: (N / CP) · (CP · S) = N · S.

The approximate number of independent samples processed by one group is: Nₛ ≈ (CP · S) / E[L].

Since S / E[L] is constant, let: Nₛ = CP · c (1)

This assumes independent samples and ignores cross-sample packing correlation.

We model Flash Attention time as quadratic in sequence length and linear in the local query-head count. For causal attention, the exact pair count is L(L + 1) / 2, but its constant factor can be absorbed. Since Ulysses gives each rank H / CP query heads:

• T_flash ≈ γ · H_local · Σᵢ Lᵢ²
• T_flash = γ · (H / CP) · Σᵢ Lᵢ²

Here, γ depends on the GPU, head dimension, and Flash Attention version.

Mean, variance, and coefficient of variation.

Let Rᵢ = Lᵢ² be independent and identically distributed, with mean μᵣ and variance σᵣ². We require L to have a finite fourth moment. First, the mean is constant:

E[T_flash] ∝ (1 / CP) · Nₛ · E[R] 
∝ (1 / CP) · (CP · c) · μᵣ
= c · μᵣ

For the variance:

Var(T_flash) ∝ Var((1 / CP) · Σᵢ Rᵢ)
= (1 / CP²) · Var(Σᵢ Rᵢ)
= Nₛ · σᵣ² / CP²

Using equation (1):

Nₛ · σᵣ² / CP² = CP · c · σᵣ² / CP²
= c · σᵣ² / CP

Therefore:

• Var(T_flash) ∝ 1 / CP
• CV(T_flash) ∝ 1 / √CP

Which ends the proof.

Exponential vs. Normal distributions

The absolute variance depends heavily on the sample-length distribution. Compare:

• an exponential distribution, L ~ Exp(λ) where p(L) = λe^(-λL)
• and a normal distribution, L ~ Normal(μ, σ²)

Assume both have mean μ = 1 / λ, and σ is small enough that negative lengths are negligible. Since Var(T_flash) ∝ Var(L²) / CP and Var(L²) = E[L⁴] − E[L²]², we only need to compare Var(L²).

Exponential

For E[Lⁿ] = n! / λⁿ: Var(L²)_exp = 24μ⁴ − 4μ⁴ = 20μ⁴.

Normal

For a normal random variable: Var(L²)_normal = (μ⁴ + 6μ²σ² + 3σ⁴) − (μ⁴ + 2μ²σ² + σ⁴) = 4μ²σ² + 2σ⁴.

Let CV = σ / μ. The ratio then becomes Var(T_flash)_exp / Var(T_flash)_normal = 20 / (4CV² + 2CV⁴)

This gives:

• CV=1: 3.33x
• CV=0.5: 17.78x
• CV=0.25: 77.58x
• CV=0.1: 497.51x

For a tight normal distribution, variation in attention times is almost nonexistent compared with a right-skewed exponential distribution. Better data packing can therefore improve load balance. A dataset mixing many short samples with a few very long ones has a heavy right tail, dramatically increasing variance and potentially slowing each iteration.

Computational Simulation

We validate the result fir