06 — Sequence / Context Parallelism (Ring Attention)

Papers: Ring Attention with Blockwise Transformers (Liu, Zaharia, Abbeel, 2023 — arxiv:<phone_number_or_numberic_id_or_random_id_85>), Striped Attention (Brandon et al. 2023 — arxiv:<phone_number_or_numberic_id_or_random_id_86>).

The problem

Attention memory is O(N²) in sequence length. FlashAttention makes it O(N) in materialized memory but the KV cache and activations still scale linearly, and eventually a single sequence’s activations don’t fit on one GPU. Beyond a few hundred thousand tokens, sequence-level sharding becomes mandatory.

Ring Attention idea

Split sequence N into P blocks of size N/P, one per GPU. Each GPU holds:

  • Its own Q, K, V blocks.

  • Rolling KV blocks from neighbors.

Algorithm:

for step in range(P):
    compute local Q_i · (currently-held K, V)
    accumulate into running online-softmax state
    async-send K, V to next rank, receive K, V from previous rank

After P iterations, each rank has seen all K, V from every rank exactly once. The online softmax trick from FlashAttention lets the accumulation happen incrementally without materializing full attention.

Key observation: the ring pass overlaps with compute. If block-attention compute per step ≥ K,V transfer time per step, you hide the comm perfectly. FA’s blockwise structure is what makes this work.

Communication cost

Per ring iteration: 2 × (N/P) × d_head × num_kv_heads × bytes sent per rank.

Total per attention layer: P × 2 × (N/P) × ... = 2N × d_head × num_kv_heads × bytes per rank.

So communication scales with total N, not with P. That’s what makes it scalable in P.

Striped Attention — the causal fix

Ring Attention as originally described has a load imbalance for causal masking: early Q blocks see many K blocks (bulk of compute); late Q blocks see fewer. Striped Attention (Brandon et al., MIT, 2023) permutes the sequence to distribute causal work evenly across ranks, giving ~1.5–2× throughput improvement for causal training. This is the version modern frameworks (Megatron-CP, torchtitan) actually use.

When you’ll hit this

Training:

  • Above 32k context: you probably want CP=2 to keep activations manageable.

  • At 128k: CP=4 or 8.

  • At 1M+ (Gemini, Kimi Linear class training): CP=32+ is standard.

Inference:

  • Long-context serving (Kimi, Claude, Gemini’s 1M+ context products) uses forms of context parallelism at inference too. Kimi Linear (2025, arxiv:<phone_number_or_numberic_id_or_random_id_87>) reports 6× decode throughput at 1M context using MLA + linear-attention hybrid + CP-style sequence sharding.

  • For most inference workloads (<32k context), CP is overkill — TP handles it fine.

Composition

CP composes with TP: CP shards the sequence, TP shards the heads. They’re orthogonal. Megatron’s --context-parallel-size composes with --tensor-model-parallel-size cleanly.

CP does not compose freely with data parallelism — CP is a form of sharding within a single sequence, whereas DP is across sequences. In practice you build a mesh (dp, cp, tp, pp) and CP is a dimension inside the sequence group.

Kernel reality

Ring Attention needs a paged / block-aware attention kernel that can accept K, V arriving in ring iterations. FlashAttention-2’s block structure is compatible; Tri Dao and colleagues extended FA2 for the ring case in 2024. In practice, most implementations wrap FA2 in a Python-level ring loop with async NCCL send/recv, and only the truly extreme long-context work goes deeper.

Relevant kernels:

  • flash-attn’s flash_attn_varlen_func with CP support (2024+).

  • Megatron’s context_parallel module.

  • torchtitan’s Ring Attention example.

  • vLLM/SGLang: CP support for inference is nascent as of late 2025; watch the PRs.

Awareness-level takeaway (what you actually need to know)

  1. Ring Attention is the canonical way to shard attention across the sequence dim.

  2. Striped Attention fixes causal imbalance and is the practical default.

  3. Communication scales with N, not P — that’s why it’s viable.

  4. You compose it with TP inside a device mesh.

  5. Long-context inference (>128k) increasingly requires it — this is where the frontier will keep pushing.

  6. For most of your work at 8k–32k contexts, TP alone is fine. Don’t over-engineer.

Exercises

  1. Given N=128k, P=8 ring stages, hidden=8192, GQA with 8 KV heads, bf16: compute bytes sent per ring iteration per rank per attention layer.

  2. Explain why online softmax is a prerequisite for Ring Attention to work at all.

  3. Sketch the causal-mask imbalance for vanilla Ring Attention with P=4, then show how a stripe permutation [0, 4, 1, 5, 2, 6, 3, 7] (for P=8) rebalances it.

Not covered in depth (but exists, awareness only)

  • DeepSpeed Ulysses — an alternative to Ring Attention using all-to-all rather than ring pass. Wins at small P; loses at large P.

  • LongContext-Serving-in-One-Attention (LCSA) and related 2024–2025 inference-time context-parallel papers.

  • Linear/hybrid attention (Kimi Linear, Mamba variants) which change the memory equation before needing CP.