Note: Sequence/Context Parallelism, Splitting One Long Sequence Across GPUs

The problem, from first principles

06-tensor-parallel-matmul could split a matmul cleanly because a matmul's output elements don't depend on each other. Attention is different: softmax is a genuinely GLOBAL operation over the whole key/value sequence — every output row needs to know the max score and the total normalizer across every key, not just the ones on its own GPU — so naively splitting the sequence across GPUs and computing softmax independently on each piece gives the WRONG answer. Sequence/context parallelism needs a smarter algorithm to fix this.

From theory to code

Implement attention_chunk_stats (the "online softmax" running statistics for one key/value chunk), merge_chunk_stats (the rescale-and-combine rule for merging two chunks' statistics), and ring_attention, which processes key/value chunks one at a time — as if each arrived from the next GPU around a ring — while never materializing the full attention matrix over the whole sequence at once.

Constraints

• attention_chunk_stats(query, key_chunk, value_chunk) returns (chunk_output, chunk_sum, chunk_max): the un-normalized weighted value sum, the un-normalized softmax denominator, and the row-wise max score, computed AS IF this chunk were the only one that existed.
• merge_chunk_stats must rescale BOTH the accumulator and the new chunk onto a shared new max (max(acc_max, chunk_max)) before adding them — never add stats computed relative to different maxes directly.
• ring_attention(query, key_chunks, value_chunks) must produce the exact same result as computing ordinary attention over the full, concatenated key/value sequence directly, for any way of splitting the sequence into chunks (even chunks of size 1).

Hints

Open one at a time. Each gives away a little more than the last.