Another View on Context Parallel for LLMs: Balancing Compute

[Vlad]

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 first with a quadratic cost model, then with a real kernel. Samples are packed into a token budget of 8192; a sample crossing the boundary is split, with its remainder used in the next batch.

import numpy as np

MAX_SEQ_LEN = 8192

def read_sample():
    return min(int(np.random.exponential(MAX_SEQ_LEN // 8)), MAX_SEQ_LEN)

def read_batches(batch_size: int, n_batches: int):
    saved_sample = None
    batches = []
    for _ in range(n_batches):
        batch, current_size = [], 0
        while current_size < batch_size:
            sample = saved_sample if saved_sample is not None else read_sample()
            saved_sample = None
            if current_size + sample > batch_size:
                take = batch_size - current_size
                saved_sample, sample = sample - take, take
            current_size += sample
            batch.append(sample)
        batches.append(batch)
    return batches

def compute_flash_cost(batch: list[int]) -> float:
    return sum(s * s for s in batch)

def simulate(dp: int, cp: int, n_steps: int = 100) -> np.ndarray:
    groups = dp // cp
    local_seqlen = cp * MAX_SEQ_LEN  # keep global batch fixed
    costs = []
    for _ in range(n_steps):
        batches = read_batches(local_seqlen, groups)
        costs.append([compute_flash_cost(b) / cp for b in batches])
    return np.array(costs)

We define imbalance as (max - min) / meanacross DP groups.

With CP=1, the maximum time is seven times the minimum. With CP=8, the max-to-min ratio is less than two.

Benchmarking Real Kernels​

Following my CUDA benchmarking guide, I benchmarked Flash Attention 3 on a single H100. I used DP = 128, a base sequence length of 8192 and scaled the micro-batch size linearly with CP.

The measured results align almost ideally with the theoretical prediction.

Results​

Context Parallelism provides an excellent opportunity to balance load across ranks. Other approaches include packing batches beforehand or avoiding heavy-tailed data distributions, but CP provides a strong and simple baseline. Many CP implementations add communication that is too slow or too difficult to overlap. So while setting CP = DP is possible, it may make sense only for certain variants, such as Magi Attention.

Code​

The complete, runnable implementation is available in the simulation repository on GitHub.