Sequence parallelism — splitting the token dimension across GPUs

Over the last few lessons, we have been slicing large LLM computations in several directions.

Data parallelism splits the batch. Tensor parallelism splits large matrix operations. Pipeline parallelism splits the layers. Yesterday we saw that all of these choices have to respect the physical topology of the GPU cluster.

Today we add another axis:

\[ \boxed{\text{split the sequence itself across GPUs}} \]

This family of techniques is usually called sequence parallelism or, in long-context systems, context parallelism.

The motivation is increasingly important: even if the model parameters fit comfortably across GPUs, a very long sequence can create enormous activation and attention-state memory.

The key intuition is simple:

If every GPU does not need to hold every token's intermediate state, distribute the tokens too.

Start with the hidden-state tensor¶

After token embedding, a Transformer commonly works with a tensor shaped roughly like:

\[ X\in\mathbb{R}^{B\times T\times d}, \]

where:

\(B\) is batch size,
\(T\) is sequence length,
\(d\) is hidden dimension.

Suppose:

\[ B=8,\qquad T=32{,}768,\qquad d=8192. \]

With BF16 values at 2 bytes each, just one hidden-state tensor contains:

\[ 8\times32{,}768\times8192\times2 \]

bytes.

That is:

\[ 4{,}294{,}967{,}296\text{ bytes} \approx4\text{ GiB}. \]

And a training step does not have only one such tensor. Many intermediate activations may need to survive long enough for backward propagation.

So long context creates a memory problem even before we discuss model weights.

Split tokens across four GPUs¶

Instead of giving every GPU all \(T\) positions, divide the sequence:

32,768 tokens

GPU 0: tokens     1–8192
GPU 1: tokens  8193–16384
GPU 2: tokens 16385–24576
GPU 3: tokens 24577–32768

Each GPU now stores only about:

\[ \frac{T}{4} \]

token positions for sequence-partitioned activations.

For our 4 GiB hidden-state example, the local portion becomes roughly:

\[ 1\text{ GiB per GPU}. \]

That looks wonderful.

But Transformers contain an operation that immediately makes the problem more interesting:

\[ \boxed{\text{attention mixes information across token positions}} \]

The easy operations are token-local¶

Many Transformer operations can be applied independently to each token.

For example, an MLP maps each token representation through something like:

\[ \operatorname{MLP}(x) = W_2\,\sigma(W_1x). \]

If GPU 0 owns tokens 1–8192, it can apply the MLP to those tokens without needing token 20,000 from GPU 2.

Likewise, many normalization and elementwise operations are mostly local to each token's hidden vector.

Conceptually:

GPU 0 tokens ── local MLP ──→ GPU 0 tokens
GPU 1 tokens ── local MLP ──→ GPU 1 tokens
GPU 2 tokens ── local MLP ──→ GPU 2 tokens
GPU 3 tokens ── local MLP ──→ GPU 3 tokens

For these operations, sequence partitioning is natural.

Attention is where the communication appears.

Attention wants keys and values from other tokens¶

For one attention head:

\[ Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V. \]

Attention computes:

\[ \operatorname{Attention}(Q,K,V) = \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}} \right)V. \]

Suppose GPU 3 owns a query corresponding to token 30,000.

In a causal Transformer, that token may attend to:

\[ 1,\ldots,30{,}000. \]

But GPU 3 locally owns only:

\[ 24{,}577,\ldots,32{,}768. \]

The keys and values for earlier tokens live on other GPUs.

So we cannot simply partition the sequence and pretend the devices are independent.

The distributed system must somehow make the relevant \(K\) and \(V\) information available.

The naive solution: gather everything¶

One straightforward strategy is:

each GPU computes local K,V
          │
          ▼
all-gather K,V across GPUs
          │
          ▼
each GPU performs attention for its local queries

After the gather, every GPU can see the full sequence's keys and values.

That preserves correct attention.

But it partly defeats the memory benefit because each device temporarily materializes global data.

It also creates substantial communication.

For long contexts, we would prefer to avoid repeatedly constructing a full copy everywhere.

A better mental model: move blocks past the queries¶

Imagine each GPU keeps its local queries stationary.

The key/value blocks move between GPUs:

round 1:
GPU0 sees KV0
GPU1 sees KV1
GPU2 sees KV2
GPU3 sees KV3

round 2:
GPU0 sees KV3
GPU1 sees KV0
GPU2 sees KV1
GPU3 sees KV2

round 3:
GPU0 sees KV2
GPU1 sees KV3
GPU2 sees KV0
GPU3 sees KV1

...

Each GPU computes the contribution from the currently available key/value block, then the block moves again.

Eventually every query has encountered every key/value block it is allowed to attend to.

This is the intuition behind ring-style context parallel attention.

Instead of:

\[ \text{copy all KV everywhere}, \]

we do something closer to:

\[ \text{stream KV blocks through the devices}. \]

That trades additional communication and scheduling complexity for lower peak memory.

But softmax seems to require all attention scores at once¶

There is a subtle mathematical problem.

For one query, attention weights are:

\[ p_i = \frac{e^{s_i}} {\sum_j e^{s_j}}. \]

If the scores \(s_j\) are distributed across GPUs, how can we normalize them without first collecting every score?

The answer is that softmax can be computed incrementally.

This is closely related to the online-softmax idea behind FlashAttention.

Suppose we have processed one block of scores and maintain:

the maximum score seen so far, \(m\),
the normalization sum, \(\ell\),
the weighted value accumulator, \(o\).

When a new block arrives with maximum \(m'\), we update the running statistics while rescaling the previous contribution.

Conceptually:

\[ m_{\text{new}} = \max(m,m'). \]

The old exponentials are rescaled by:

\[ e^{m-m_{\text{new}}}, \]

and the new ones by:

\[ e^{m'-m_{\text{new}}}. \]

This lets us combine attention blocks without storing the entire:

\[ T\times T \]

attention matrix.

So distributed long-context attention relies on the same deeper principle we encountered with FlashAttention:

\[ \boxed{\text{attention can be tiled without changing the mathematical result}} \]

provided the running softmax statistics are maintained correctly.

Causal attention gives us another optimization¶

In causal language modeling, token \(t\) cannot attend to future tokens.

So the attention matrix is lower triangular:

query →
key ↓

X . . . .
X X . . .
X X X . .
X X X X .
X X X X X

Suppose sequence blocks are ordered:

GPU 0 → earliest tokens
GPU 1 → next tokens
GPU 2 → next tokens
GPU 3 → latest tokens

GPU 0 never needs keys from GPUs 1–3 for its queries because those keys correspond to future positions.

GPU 1 needs GPU 0's keys, but not GPU 2's or GPU 3's.

GPU 3 may need all earlier blocks.

A topology-aware implementation can exploit this causal structure to avoid useless computation or communication.

The logical attention pattern therefore influences the physical communication pattern.

Sequence parallelism is not exactly context parallelism¶

The terminology can be confusing because different systems use these names somewhat differently.

A useful conceptual distinction is:

Sequence parallelism often refers to partitioning token-wise activations for operations that do not require full-sequence attention state on every device.

Context parallelism usually emphasizes distributing the actual attention computation for long contexts across devices.

They share the same core idea:

\[ \boxed{\text{the sequence dimension becomes a distributed dimension}} \]

but context parallelism has to solve the harder cross-token attention problem.

When reading system documentation, it is therefore worth checking what a particular framework means by the term rather than relying only on the label.

Why long context makes this increasingly valuable¶

Consider increasing context length from:

\[ 8{,}192 \]

tokens to:

\[ 131{,}072. \]

That is a factor of:

\[ 16. \]

Token-wise activation memory grows roughly linearly with \(T\):

\[ O(T). \]

Naive attention-score storage would grow quadratically:

\[ O(T^2). \]

Modern memory-efficient attention avoids materializing that full quadratic matrix, which is crucial.

But even after eliminating the explicit \(T^2\) matrix, the remaining activations, KV tensors, and computation associated with 131k tokens are enormous.

So long-context scaling increasingly becomes a distributed-systems problem.

A model may fit.

A single sequence may not.

Training and inference stress this dimension differently¶

During training, long sequences create large activation memory and expensive backward passes.

Sequence/context parallelism can distribute those activations and the attention computation across devices.

During inference, the dominant long-context structure is often the KV cache.

For a long conversation, the server may need to hold keys and values for hundreds of thousands or even millions of previous tokens.

That creates a related possibility:

GPU 0 → one portion of KV history
GPU 1 → another portion
GPU 2 → another portion
GPU 3 → another portion

The new query can be evaluated against distributed KV blocks, and partial attention results can be combined.

The same basic idea therefore appears in both training and inference:

distribute the token history rather than requiring one accelerator to own all of it.

Sequence parallelism composes with the other dimensions¶

Now our distributed Transformer has several possible axes:

data parallelism      → split examples
tensor parallelism    → split hidden/model dimensions
pipeline parallelism  → split layers
sequence parallelism  → split tokens

A large training run might therefore use something like:

\[ \text{DP}\times\text{PP}\times\text{TP}\times\text{CP}. \]

For example:

\[ 8\times4\times8\times2 = 512\text{ GPUs}. \]

That does not mean every combination is sensible.

Each dimension creates its own communication pattern.

The topology lesson from yesterday still applies: tightly communicating groups should be mapped onto appropriately fast links.

The parallelism configuration is therefore becoming a mapping between two structures:

Transformer dimensions
        ↓
batch / layers / hidden width / tokens

Hardware dimensions
        ↓
GPU / local fabric / server / network / rack

Efficient large-scale training is largely about matching those structures well.

A concrete mental model: four librarians and one enormous book¶

Imagine a book with 100,000 pages.

One librarian cannot conveniently keep working notes for every page on one desk.

So four librarians divide the book:

Librarian 0: pages     1–25,000
Librarian 1: pages 25,001–50,000
Librarian 2: pages 50,001–75,000
Librarian 3: pages 75,001–100,000

For tasks that concern only one page, each librarian works independently.

That is like an MLP applied token-by-token.

But now librarian 3 is asked:

On page 90,000, what earlier passages does this sentence refer to?

The answer may depend on notes held by every other librarian.

One solution is to photocopy all 100,000 pages onto every desk.

That works, but defeats the reason for splitting the book.

A better system circulates boxes of notes:

box 1 → librarian 0 → librarian 1 → librarian 2 → librarian 3
box 2 → librarian 1 → librarian 2 → librarian 3 → librarian 0
...

Each librarian updates a running answer as relevant notes pass by.

No desk needs to contain the entire book at once.

That is the core intuition behind distributed long-context attention.

The deeper lesson¶

Early in this series, sequence length looked like an input parameter:

\[ T. \]

As models scale, \(T\) becomes something more important.

It becomes a parallelization dimension.

That is a significant shift in perspective.

We are no longer asking only:

How much context can the Transformer mathematically attend to?

We must also ask:

Where do the representations for all those tokens live, and how does attention move information between them?

The maximum useful context length of a production model is therefore not determined only by positional encodings or attention algorithms.

It is also constrained by:

memory capacity
memory bandwidth
GPU interconnects
communication algorithms
parallelism layout
latency requirements

So today's main principle is:

\[ \boxed{ \text{long context is not just a modeling problem; at scale, it is a distributed-memory problem} } \]

Once token history itself becomes distributed, the next interesting question is what happens when the model architecture is sparse too — for example, when only a few expert networks activate for each token. That leads to another kind of routing problem: Mixture-of-Experts, where the network must move tokens to the GPUs holding the experts they selected.