Navigation
🌐 Distributed GPU Systems· 10 min read

Context Parallelism

Ring Attention and Ulysses — splitting long-sequence attention across GPUs when the sequence is too long for one GPU's memory.

distributedcontext-parallelismring-attentionulysseslong-context

Context parallelism (CP) splits the sequence dimension of attention across GPUs — enabling sequences too long for a single GPU’s KV cache memory. While TP splits the model and DP splits the batch, CP splits the context.

1. The Problem: Long Sequences Break Single-GPU Memory

For a 7B model with MHA serving 128K context: KV cache = 64 GB. That’s nearly all of an H100’s memory — no room for the model weights.

For training with full attention: the attention matrix is S2=128K2=16S^2 = 128K^2 = 16 billion elements × 2 bytes = 32 GB per head. With 32 heads: 1 TB. Impossible on any single GPU.

Context parallelism distributes this across GPUs: each GPU handles S/PS/P tokens of the sequence.

Ring Attention uses the same online softmax technique as FlashAttention to merge partial attention results across tiles.

2. Ring Attention

How It Works

Split the sequence into PP chunks, one per GPU. Each GPU has Q for its chunk. K and V are rotated through the ring:

Step 0: GPU i computes attention(Q_i, K_i, V_i)    — local KV
Step 1: Rotate K,V: GPU i receives K_{i-1}, V_{i-1} from left neighbor
        GPU i computes attention(Q_i, K_{i-1}, V_{i-1})
Step 2: Rotate again: GPU i receives K_{i-2}, V_{i-2}
        ...
Step P-1: All KV chunks have rotated through all GPUs

After P-1 rotations, each GPU has computed attention against all KV positions — using online softmax to merge partial results (same technique as FlashAttention).

Communication

Each step: send S/PS/P tokens of K and V to the next GPU.

  • Per step: 2×(S/P)×dk×22 \times (S/P) \times d_k \times 2 bytes (K + V, BF16)
  • Total: (P−1)×2×(S/P)×dk×2≈2Sdk×2(P-1) \times 2 \times (S/P) \times d_k \times 2 \approx 2 S d_k \times 2 bytes

Key: Communication of K/V rotation can be overlapped with the attention computation of the current step. While GPU ii computes attention with KjK_{j}, it’s simultaneously receiving Kj−1K_{j-1} for the next step.

Memory per GPU

  • Q: S/PS/P tokens (local only)
  • K, V: S/PS/P tokens (current rotation — only one chunk at a time)
  • Attention: (S/P)2(S/P)^2 per head (within each tile)

Total attention memory: O(S2/P2)O(S^2/P^2) per head — quadratically reduced by CP degree.

Worked Example

7B model, MHA (32 heads, dk=128), S=128K, CP=4:

  • Without CP: KV cache per request = 128K × 512 KB/token = 64 GB (doesn’t fit on one H100)
  • With CP=4: each GPU handles S/4 = 32K tokens → KV cache per GPU = 32K × 512 KB = 16 GB ✅
  • Attention memory per GPU: (S/P)2=32K2=1B(S/P)^2 = 32K^2 = 1B elements × 32 heads × 2 bytes = 2 GB (vs 128 GB without CP)
  • Communication per layer: rotating 32K × 128 × 2 bytes × 2 (K+V) = 16 MB per step, 3 steps = 48 MB
  • On NVLink (450 GB/s): 48 MB / 450 GB/s = 0.1 ms per layer — easily overlapped with attention compute

3. Ulysses (DeepSpeed)

An alternative to Ring Attention that uses All-to-All instead of rotation:

Step 1: Each GPU has Q_i, K_i, V_i for its S/P tokens
Step 2: All-to-All on K and V: redistribute by head instead of by sequence
        → Each GPU now has full S tokens but only N_h/P heads
Step 3: Each GPU computes full attention for its heads (all S tokens)
Step 4: All-to-All on output: redistribute back by sequence

Ring Attention vs Ulysses

Ring AttentionUlysses
CommunicationPoint-to-point rotation (P-1 steps)Two All-to-All operations
Overlap potentialExcellent (rotate while computing)Limited (All-to-All is blocking)
Memory per GPUO(S2/P2)O(S^2/P^2) attentionO(S2/P)O(S^2/P) attention (full sequence, fewer heads)
Best forVery long sequences, limited NVLinkModerate sequences, fast All-to-All
Used byLlama 3.1, variousDeepSpeed

4. When to Use Context Parallelism

ConditionUse CP?
Sequence ≤ 8KNo — single GPU handles it
Sequence 8K-32KMaybe — if KV cache is tight
Sequence 32K-128KYes — single GPU can’t hold KV cache + model
Sequence > 128KYes — CP=4-8 typical

CP is typically the innermost parallelism dimension (within a TP group), because it has the highest communication frequency (every attention layer).

Composition: DP × TP × CP. Example: 64 GPUs = DP=8 × TP=4 × CP=2.

Self-Check

  • Explain why long sequences break single-GPU memory (calculate KV cache for 128K)
  • Describe Ring Attention’s rotation mechanism and how it uses online softmax
  • Explain why Ring Attention communication can overlap with attention compute
  • Compare Ring Attention vs Ulysses (communication pattern, memory, best use case)
  • Explain how CP composes with TP and DP

References