Context parallel attention: exchange keys and merge softmax globally

For full attention, each query compares against all valid keys in its sequence; for causal attention, it compares against the permitted global prefix. A rank owning token positions must still access remote KV needed by those queries. Implementations can all-gather, exchange chunks in a ring, or use other supported strategies. The exact traffic and overlap depend on placement, head sharing, chunk size, and algorithm. There is no universal collective count implied by saying context parallel.

Global token positions determine causal masks and positional encodings. Treating each local chunk as a new sequence can expose future keys or hide required earlier keys. Packed variable-length sequences also need boundaries that prevent attention across different examples. GQA can reduce the number of KV features exchanged but does not remove the remote-key dependency. Backward must return and combine contributions to the owners of K and V; a correct forward result alone does not establish correct parameter training.

A chunk can summarize its maximum score m, scaled exponential denominator ℓ, and scaled value numerator z. To combine two chunks, rescale both summaries to the larger maximum before adding denominator and numerator. Divide only after all required contributions are included. This produces the same mathematical attention as a global softmax in exact arithmetic while avoiding overflow from unshifted exponentials. Floating-point rounding still depends on block order and accumulation precision.

An all-masked chunk needs a defined neutral-summary handling so expressions such as minus-infinity minus minus-infinity do not create NaNs. A query with no permitted keys requires the model’s explicit convention rather than this ordinary positive-denominator formula. Dropout and backward intermediates also need matching global randomness and reconstruction rules. Do not reuse a simplified scalar summary as if it specified every implementation detail of a production fused attention kernel.

A simple contiguous causal partition can make early query chunks do less attention work than later ones. A supported zigzag or balanced layout can redistribute that work while preserving global positions. Such a layout changes ownership, not the allowed attention relation. Record both the logical token order and the physical chunk assignment. Memory savings depend on retained KV, intermediate summaries, communication buffers, and backward recomputation; exchanging chunks does not imply that every rank stores every full-sequence activation permanently.

The network lab shows only one local computation. Reset, set x=2 and w=1, predict sigmoid(1.5)≈0.818, then Step. Its single scalar graph has no attention softmax, KV exchange, masking, or distributed summary merge. It helps distinguish an inspectable computation from the additional dependencies a distributed version must preserve. The worked summary arithmetic, rather than the diagram, proves the correct normalization across the two chunks.