On Context Parallelism

31 Aug, 2026

Introduction

The goal of this blog post is to explain four implementations of context parallelism. In this post at least, context parallelism refers primarily to the communication required within the attention computation. Activations remain sequence-sharded at transformer-block boundaries, and attention communicates across ranks so that each local query can attend to the required remote keys and values. Two things to note about the following attention-block forward function are the shape of x_local and the fact that context_parallel_attention accepts already sequence-sharded activations.

Throughout the post, B is the batch size, S is the global sequence length, H is the hidden size, N is the number of attention heads, and D = H/N is the head dimension. The examples assume causal attention and zigzag sequence sharding.

    def forward(self, x_local):
        """
        LOCAL input:
            x_local: [B,s,H]

        LOCAL projection:
            q, k, v: [B,s,N,D]

        LOCAL CP-attention result:
            attended: [B,s,N,D] -> flattened [B,s,H]

        LOCAL output:
            [B,s,H]
        """
        batch, local_sequence, _ = x_local.shape
        qkv = _linear_projection(x_local, self.qkv)
        qkv = qkv.view(
            batch, local_sequence, 3, self.config.num_heads, self.head_dim
        )
        q, k, v = qkv.unbind(dim=2)
        q, k = apply_rope(q), apply_rope(k)
        attended = context_parallel_attention(q, k, v)
        attended = attended.reshape(batch, local_sequence, self.config.hidden_size)
        return _linear_projection(attended, self.output)

All-gather

K/V all-gather is probably the simplest implementation of context parallelism. Each rank keeps its local query tokens, all-gathers the key and value tokens from the other context-parallel ranks, and then computes attention for only its local queries against the (after the all-gather) replicated K/V sequence.

To ground ourselves a bit:

C = context-parallel size
S = global sequence length
s = S/C                    local sequence length
t = S/(2C) = s/2           length of one "zigzag" chunk

Before attention, rank r owns two noncontiguous "zigzag" chunks of each of Q, K, and V, defined like:

rank r owns:
    chunk r
    chunk 2C-r-1

local Q, K, V:
    [B,s,N,D] = [B,2t,N,D]

# example with C=4
rank 0: [chunk 0, chunk 7]
rank 1: [chunk 1, chunk 6]
rank 2: [chunk 2, chunk 5]
rank 3: [chunk 3, chunk 4]

Transformer Engine (TE) calls this DualChunkSwap, or "zigzag" sharding. Assigning one early and one late chunk to each rank balances causal-attention work because early queries attend to much shorter prefixes than later queries. Megatron constructs this input layout for use in TE here.

However, the all-gather concatenates each rank's complete local tensor in process-group rank order, which requires reordering the result into the true global sequence order. For instance, in the C=4 example above, the all-gather would produce chunks in the order [chunk 0, chunk 7, chunk 1, chunk 6, chunk 2, chunk 5, chunk 3, chunk 4]. TE performs this reordering here.

Forward

Once K/V are in global sequence order, the conceptual forward pass is small:

def forward(ctx, q, k, v, local_positions, cp_group, scale):
    """
    LOCAL inputs:
        q, k, v:         [B,s,N,D]
        local_positions: [s]

    REPLICATED temporaries:
        k_full, v_full:  [B,S,N,D]
        global_positions:[S]

    LOCAL output:
        output:           [B,s,N,D]
    """
    cp_size = dist.get_world_size(cp_group)

    k_full = _all_gather_sequence(k, cp_group)
    v_full = _all_gather_sequence(v, cp_group)
    k_full = _rank_order_to_sequence_order(k_full, cp_size)
    v_full = _rank_order_to_sequence_order(v_full, cp_size)
    global_positions = torch.arange(k_full.size(1), device=k_full.device)

    output, lse = causal_attention(
        q,
        k_full,
        v_full,
        q_positions=local_positions,
        kv_positions=global_positions,
        scale=scale,
    )

    ctx.cp_group = cp_group
    ctx.cp_size = cp_size
    ctx.scale = scale
    ctx.save_for_backward(q, k, v, output, lse, local_positions)
    return output

Local Q contains two (generally) noncontiguous portions of the global sequence, while K/V are in global sequence order. TE processes these two local Q chunks separately. The code around here effectively boils down to:

out_early = attention(q_early, causally_visible_kv_for_q_early)
out_late = attention(q_late, causally_visible_kv_for_q_late)
output_local = concat(out_early, out_late)

# On rank 1 in the C=4 example
First attention call:
    Q = chunk 1
    K/V = chunks 0..1

Second attention call:
    Q = chunk 6
    K/V = chunks 0..6

The output is thus [B,s,N,D] and no output communication is necessary because the rank that owns a Q token also owns that token's output.

Backward

Because each rank's local queries may attend to the K/V tokens owned by every other rank, the attention backward on one rank produces:

dQ:                  [B,s,N,D]
partial dK, partial dV: [B,S,N,D]

dQ is already complete because each Q token participated in attention only on its owning rank. The rank-local partial dK and dV buffers, however, contain only the contributions from some of this rank's queries. Specifically, a K/V token at position p receives gradient contributions from every local query position q where q >= p. Another way to say this is that the largest local query position determines the outer boundary and K/V positions later than that boundary receive no gradient contribution from this rank.

Using our same running C=4 example:

# on rank 1
K/V chunk 0,1 receive contributions from Q chunks 1 and 6
K/V chunks 2..6 receive contributions from Q chunk 6
K/V chunk 7 receives no contribution

TE's backward can be found here. Given the shape of the forward pass, it re-gathers and reorders K/V at the beginning of backward here, accumulates the partial dK/dV contributions from both local Q chunks into global-sequence buffers here, and applies the inverse zigzag permutation before reduce-scattering dK/dV here.

The communication pattern is therefore:

Forward:
    all-gather K
    all-gather V

Backward:
    all-gather K
    all-gather V
    reduce-scatter dK
    reduce-scatter dV

The main advantage to this implementation is simplicity, i.e. after K/V have been gathered and reordered, each rank performs otherwise ordinary attention for its local queries. Memory use is the main disadvantage. Every rank materializes full [B,S,N,D] K and V tensors, and backward additionally requires full-sequence partial dK/dV buffers.

All-to-all

The all-to-all (A2A) implementation is effectively DeepSpeed-Ulysses, which I discussed in this previous blog post, so I won't fully go in depth into it. TE exposes the same attention layout as Ulysses whereby activations remain sequence-sharded at the transformer-block boundary, but attention temporarily exchanges sequence sharding for head sharding.

A2A context parallelism is best understood as temporarily transposing which dimension is distributed. Outside of attention the sequence is distributed, whereas inside attention the heads are distributed.

Each rank partitions its local Q, K, and V tensors into head groups and sends each group to its destination rank. After the all-to-all, rank r owns one head group from every sequence shard, giving it the full sequence for N/C heads. It can therefore compute ordinary attention independently for those heads. A second all-to-all reverses the transformation, restoring all N heads for the rank's original S/C tokens. As in the all-gather implementation above, TE reorders the chunks into global sequence order before attention. Its forward pass can be found here.

The pseudocode is:

def a2a_context_parallel_attention(q, k, v, positions, cp_group=None):
    """
    FORWARD local shapes:
        input Q,K,V:       [B,S/C,N,D]
        after QKV A2A:     [B,S,N/C,D]
        local attention O: [B,S,N/C,D]
        after output A2A:  [B,S/C,N,D]

    BACKWARD local shapes:
        input dO:          [B,S/C,N,D]
        after dO A2A:      [B,S,N/C,D]
        local dQ,dK,dV:    [B,S,N/C,D]
        returned gradients:[B,S/C,N,D]

    """
    q_heads = SequenceToHeadsAllToAll.apply(q, cp_group)
    k_heads = SequenceToHeadsAllToAll.apply(k, cp_group)
    v_heads = SequenceToHeadsAllToAll.apply(v, cp_group)

    q_heads, k_heads, v_heads = reorder_heads(q_heads, k_heads, v_heads)
    output_heads = _causal_attention(
        q_heads, k_heads, v_heads
    )
    return HeadsToSequenceAllToAll.apply(output_heads, cp_group)

Unlike K/V all-gather, A2A does not increase the number of Q/K/V elements held by a rank. The sequence dimension becomes C times larger while the head dimension becomes C times smaller. The tradeoff is that the Q, K, and V head counts must divide evenly across the A2A group. This limits the maximum A2A group size, particularly for GQA or MQA models with relatively few K/V heads.

Backward reverses the same layout transformations. The backward of the final output A2A first converts sequence-sharded [B,S/C,N,D] dO into full-sequence, head-sharded [B,S,N/C,D] dO. Attention backward then computes dQ, dK, and dV for the rank's full-sequence head group. Finally, the backward of the three input A2As returns those gradients to [B,S/C,N,D].

The complete attention communication pattern is:

Forward:
    A2A Q, K, and V from sequence sharding to head sharding
    attention over the full sequence for N/C heads
    inverse A2A output back to sequence sharding

Backward:
    A2A dO from sequence sharding to head sharding
    local attention backward
    inverse A2A dQ, dK, and dV back to sequence sharding

Peer-to-peer

The peer-to-peer (P2P) implementation keeps Q fixed on its owning rank and circulates K/V shards through a ring. Each rank computes attention between its local queries and one visiting K/V shard at a time, merging the partial results with an online softmax. After every rank has seen every K/V shard, it holds the complete attention output for its local queries without ever materializing full-sequence K and V tensors.

Unlike the all-gather, these sends and receives do not constitute a global collective because each rank communicates only with its two ring neighbors. The tradeoff for lower memory use is a larger number of point-to-point communication rounds where attention is split into C ring steps, making performance depend on overlapping the next K/V transfer with the current attention computation.

Forward

Transformer Engine's complete implementation is AttnFuncWithCPAndKVP2P. It constructs the ring neighbors here, and the forward ring begins here. TE starts the next asynchronous K/V send and receive before computing attention on the current shard, allowing communication to overlap with attention.

The conceptual forward pass is:

def forward(ctx, q, k, v, local_positions, cp_group, scale):
    """
    LOCAL inputs:
        q, k, v:         [B,S/C,N,D]
        local_positions: [S/C]

    LOCAL resident ring state:
        current_k, current_v: [B,S/C,N,D]

    LOCAL online-softmax state:
        output: [B,S/C,N,D]
        lse:    [B,N,S/C]

    LOCAL output:
        output: [B,S/C,N,D]
    """
    cp_size = dist.get_world_size(cp_group)
    rank = dist.get_rank(cp_group)
    current_k, current_v = k, v

    output = None
    lse = None

    for step in range(cp_size):
        kv_owner = (rank - step) % cp_size
        current_positions = _zigzag_positions_for_rank(kv_owner)

        tile_output, tile_lse = attention(
            q,
            current_k,
            current_v,
            q_positions=local_positions,
            kv_positions=current_positions,
            scale=scale,
        )
        output, lse = _merge_attention_tiles(
            output, lse, tile_output, tile_lse
        )

        if step + 1 < cp_size:
            current_k, current_v = _ring_exchange(
                current_k, current_v, cp_group
            )

    ctx.cp_group = cp_group
    ctx.scale = scale
    ctx.save_for_backward(q, k, v, output, lse, local_positions)
    return output

At a given step, each rank has a different kv_owner (i.e. it receives from one neighbor and sends to the other). Conceptually, the exchange looks like:

def _ring_exchange(tensors, group, send_offset=1, recv_offset=-1):
    ...
    global_ranks = dist.get_process_group_ranks(group)
    send_peer = global_ranks[(group_rank + send_offset) % group_world_size]
    recv_peer = global_ranks[(group_rank + recv_offset) % group_world_size]
    outputs = [torch.empty_like(x) for x in tensors]
    ops = []
    for source, target in zip(tensors, outputs):
        ops.append(dist.P2POp(dist.isend, source, send_peer, group))
        ops.append(dist.P2POp(dist.irecv, target, recv_peer, group))
    for request in dist.batch_isend_irecv(ops):
        request.wait()
    return outputs

I feel like most of us are probably familiar with the FlashAttention-style softmax, but perhaps not this LSE-based merge of completed attention calls (at least I wasn't), so I think it is worth spending a few extra moments on it. Conceptually, the merge looks like:

def _merge_attention_tiles(output, lse, tile_output, tile_lse):
    ...
    if output is None:
        return tile_output, tile_lse
    merged_lse = torch.logaddexp(lse, tile_lse)
    old_scale = torch.exp(lse - merged_lse)
    tile_scale = torch.exp(tile_lse - merged_lse)
    output = old_scale * output + tile_scale * tile_output
    return output, merged_lse

Now, we know that the FlashAttention epilogue takes the accumulated output and normalizes it by the running sum. After each round, TE receives this normalized K/V-shard output as well as an lse value for every query row and head, where lse is the log-sum-exp calculated as tile_max + log(tile_sum). The lse can be interpreted as the logarithm of the tile's total unnormalized softmax weight. Looking into it further:

tile_max = max_j(s_j)
tile_sum = sum_j(exp(s_j - tile_max))
tile_lse = tile_max + log(tile_sum)
         = m + log(sum_j{e^(s_j - m)})
         = log( e^m * sum_j{e^(s_j - m)})
         = log(sum_j{e^s_j})

So, logaddexp (which calculates log(e^x + e^y)) merges our previous and current unnormalized attention weights and keeps us in log-space. Then we simply use torch.exp to find the relative strengths of previous and new tiles and scale by those factors.

So, if we...

# suppose that...
previous_total_weight = exp(lse)
tile_total_weight     = exp(tile_lse)
# then...
merged_lse = torch.logaddexp(lse, tile_lse)
# which is simply the numerically stable
# log space equivalent to...
merged_total_weight = previous_total_weight + tile_total_weight
merged_lse = log(merged_total_weight)
# and so...
old_scale
    = exp(lse - merged_lse)
    = previous_total_weight / merged_total_weight

tile_scale
    = exp(tile_lse - merged_lse)
    = tile_total_weight / merged_total_weight
# which means that...
old_scale + tile_scale == 1
# and finally...
output = old_scale * output + tile_scale * tile_output

Thus, a tile containing stronger query-key matches has more total softmax weight and consequently makes a larger contribution to the merged output.

Backward

As in the all-gather implementation, d_q is accumulated locally because Q never leaves its owning rank. K/V shards, however, were used by queries on many ranks, and so their dK/dV contributions must be accumulated as the corresponding K/V shard travels through the ring. Then after a complete rotation, each K/V shard and its fully accumulated dK/dV return to their owning rank.

TE's backward ring begins here, but the more pedagogical backward is:

def backward(ctx, grad_output):
    q, k, v, output, lse, local_positions = ctx.saved_tensors
    cp_size = dist.get_world_size(ctx.cp_group)
    rank = dist.get_rank(ctx.cp_group)

    current_k, current_v = k, v
    current_dk = torch.zeros_like(k)
    current_dv = torch.zeros_like(v)
    d_q = torch.zeros_like(q)

    for step in range(cp_size):
        kv_owner = (rank - step) % cp_size
        current_positions = _zigzag_positions_for_rank(kv_owner)

        tile_dq, tile_dk, tile_dv = attention_tile_backward(
            q, current_k, current_v, output, grad_output,
            lse, q_positions=local_positions, kv_positions=current_positions, scale=ctx.scale,
        )
        d_q.add_(tile_dq)
        current_dk.add_(tile_dk)
        current_dv.add_(tile_dv)

        current_k, current_v, current_dk, current_dv = _ring_exchange(
            current_k, current_v,
            current_dk, current_dv, ctx.cp_group,
        )
    return d_q, current_dk, current_dv

Peer-to-peer + All-to-all

The final implementation combines the previous two approaches over a logical two-dimensional process mesh. If we let A be the A2A group size and R be the P2P group size, then the total context-parallel size is C = A * R. After the initial A2A, each rank owns S/R query positions and N/A heads. An entire P2P group collectively covers the full sequence S, but only for its assigned N/A heads.

For C=8, A=2, and R=4, under our rank mapping, the A2A groups are [0,1], [2,3], [4,5], and [6,7], while the P2P groups are [0,2,4,6] and [1,3,5,7], which handle the first and last N/2 heads, respectively.

When entering context-parallel attention, each rank holds Q,K,V: [B, S/(A*R), N, D] == [B, S/C, N, D]. As in the all-to-all implementation above, the ranks within each A2A group collect only the heads for which they are responsible before computing attention. Each rank within the P2P group then has Q,K,V: [B, S/R, N/A, D] (i.e. it receives A times as many sequence positions and 1/A as many heads). Thus, in our C=8, A=2, R=4 example, two separate P2P groups compute attention, with each rank responsible for one quarter of the queries and one half of the heads.

The complete composition would look something like this:

def a2a_p2p_context_parallel_attention(
    q, k, v, positions,
    a2a_group, p2p_group, cp_group,
):
    q_heads = SequenceToHeadsAllToAll.apply(q, a2a_group)
    k_heads = SequenceToHeadsAllToAll.apply(k, a2a_group)
    v_heads = SequenceToHeadsAllToAll.apply(v, a2a_group)

    positions_long = _recreate_group_positions(
        positions, a2a_group, cp_group,
    )

    output_heads = p2p_context_parallel_attention(
        q_heads, k_heads, v_heads,
        positions_long, p2p_group,
    )

    return HeadsToSequenceAllToAll.apply(output_heads, a2a_group)

Transformer Engine accepts the groups as [a2a_group, p2p_group] and separates the two mesh dimensions here. Its pre-attention A2A is here, and its inverse output A2A is here.

Results and Conclusion

The benchmarks ran on 8x H100 SXM GPUs with batch size 1 and bf16 Q, K, and V tensors. The attention shape was 64 heads with a head dimension of 128 (H=8192), using ordinary multi-head attention rather than GQA. Only the core attention operation was timed, so the results exclude the QKV and output projections. In the Torch version used for these measurements, the all-gather implementation's K/V-reordering index_select failed at sequence lengths of 256K and above (it's a known bug in that Torch version). In the tables, 2×4 and 4×2 mean A2A group size × P2P group size. The code can be found here.

Note that this experiment does not measure cross-node communication. When multiple nodes are required, the strategies' different uses of all-to-all and peer-to-peer communication would likely change the timing results significantly.

Forward time

Forward + backward time

The most surprising result to me is how competitive pure A2A is. It produces the fastest forward-plus-backward time at 16K and 32K, and remains within roughly 5% of the fastest strategy at every longer sequence length. If a model's head geometry supports the desired context-parallel degree and training remains within a fast local communication domain (i.e. NVLink/NVSwitch), Ulysses would seem to provide most of the benefit of context parallelism without the additional complexity of a P2P ring. The drawback to A2A, though, is that it's constrained by K/V-head divisibility, while P2P can scale the context-parallel group independently of the number of heads.