Sequence parallelism — splitting the token dimension across GPUs

Over the last few lessons, we have learned to cut a Transformer along several axes.

Data parallelism cuts the batch.

Tensor parallelism cuts the hidden dimensions and matrix operations.

Pipeline parallelism cuts the layers.

FSDP cuts the stored model state.

Today we will cut another dimension that becomes increasingly important as context windows grow:

\[ \boxed{\text{the sequence itself}} \]

This family of techniques is usually called sequence parallelism or, for attention specifically, context parallelism.

The intuition is simple:

If a request contains 128,000 tokens, why should every GPU hold the activations for all 128,000 positions?

Instead, different GPUs can own different stretches of the sequence.

The catch is that attention allows tokens in one stretch to interact with tokens in another. So splitting the sequence saves memory, but creates a new communication problem.

Start with the activation tensor¶

A Transformer layer typically receives something shaped like:

\[ X\in\mathbb{R}^{B\times T\times d_{\text{model}}}. \]

Here:

\(B\) is batch size,
\(T\) is sequence length,
\(d_{\text{model}}\) is hidden width.

Suppose:

\[ B=1,\qquad T=128{,}000,\qquad d_{\text{model}}=8192. \]

In BF16, the raw hidden-state tensor alone contains:

\[ 1\times128{,}000\times8192 = 1{,}048{,}576{,}000 \]

values.

At 2 bytes each:

\[ \approx2.1\text{ GB}. \]

And that is only one hidden-state tensor.

Training may need many intermediate activations across many layers.

Long context therefore puts enormous pressure on activation memory.

Split the sequence across four GPUs¶

Instead of every GPU holding all \(128{,}000\) token positions, divide them:

GPU 0 → tokens      0–31,999
GPU 1 → tokens 32,000–63,999
GPU 2 → tokens 64,000–95,999
GPU 3 → tokens 96,000–127,999

Each GPU now owns:

\[ \frac{T}{4} = 32{,}000 \]

positions.

For tensors whose computation is independent across token positions, this is easy.

Suppose we apply RMSNorm:

\[ y_t = \frac{x_t} {\sqrt{\frac{1}{d}\sum_{j=1}^{d}x_{t,j}^2+\epsilon}} \odot\gamma. \]

The normalization for token \(t\) depends on that token's hidden dimensions, not on other token positions.

So GPU 0 can normalize its 32,000 tokens while GPU 1 normalizes its own 32,000 tokens.

No token-to-token communication is required.

The same is true for many elementwise operations.

Why this pairs naturally with tensor parallelism¶

Yesterday's tensor parallelism often split the hidden dimension.

Imagine four GPUs holding:

GPU 0 → hidden channels 0–2047
GPU 1 → hidden channels 2048–4095
GPU 2 → hidden channels 4096–6143
GPU 3 → hidden channels 6144–8191

Each GPU may therefore hold all tokens but only part of each token.

For some operations, that is awkward.

Layer normalization or RMSNorm wants all hidden channels belonging to a token.

Sequence parallelism can effectively transpose the ownership pattern:

tensor-parallel layout:

GPU 0 → all tokens, hidden shard 0
GPU 1 → all tokens, hidden shard 1
GPU 2 → all tokens, hidden shard 2
GPU 3 → all tokens, hidden shard 3

             ↓ redistribute

sequence-parallel layout:

GPU 0 → token shard 0, full hidden width
GPU 1 → token shard 1, full hidden width
GPU 2 → token shard 2, full hidden width
GPU 3 → token shard 3, full hidden width

Now each GPU can perform token-local operations independently.

So sequence parallelism is often less about making the model mathematically different and more about choosing a better layout for intermediate activations.

A tiny matrix mental model¶

Imagine the activation matrix for one batch item:

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

Draw it as:

                  hidden dimension d
          ┌─────────────────────────────┐
token 1   │                             │
token 2   │                             │
token 3   │                             │
...       │                             │
token T   │                             │
          └─────────────────────────────┘

Tensor parallelism often cuts vertically:

          ┌───────┬───────┬───────┬───────┐
          │ GPU 0 │ GPU 1 │ GPU 2 │ GPU 3 │
          │       │       │       │       │
          │       │       │       │       │
          └───────┴───────┴───────┴───────┘

Sequence parallelism cuts horizontally:

          ┌─────────────────────────────┐
GPU 0     │                             │
          ├─────────────────────────────┤
GPU 1     │                             │
          ├─────────────────────────────┤
GPU 2     │                             │
          ├─────────────────────────────┤
GPU 3     │                             │
          └─────────────────────────────┘

Same tensor.

Different cut.

This is a recurring theme in distributed ML:

\[ \boxed{\text{tensor layout determines communication}} \]

But attention breaks the easy story¶

For an MLP or normalization, token positions can often be processed independently.

Self-attention explicitly couples them.

Recall:

\[ 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 the final 32,000 tokens.

Its queries may need to attend to keys and values from earlier tokens owned by:

GPU 0
GPU 1
GPU 2

For causal attention, a query near token 120,000 may depend on nearly every previous token.

So simply saying:

GPU 3 owns token 120,000

does not mean GPU 3 has everything needed to compute that token's attention.

This is the central challenge of context parallelism.

Why long context makes this worth solving¶

Ignoring optimized kernels for a moment, the attention score matrix has shape:

\[ T\times T. \]

If:

\[ T=128{,}000, \]

then:

\[ T^2 = 16{,}384{,}000{,}000 \]

possible query-key score positions.

Modern kernels such as FlashAttention avoid materializing that entire matrix in HBM, which is crucial.

But even with memory-efficient attention, the model must conceptually process interactions between distant queries and keys.

Long context therefore creates two separate scaling pressures:

activation storage
+
attention computation / communication

Sequence and context parallelism help distribute both across GPUs.

One strategy: move K and V around the ring¶

Suppose four GPUs each own a sequence chunk:

GPU 0 → Q0, K0, V0
GPU 1 → Q1, K1, V1
GPU 2 → Q2, K2, V2
GPU 3 → Q3, K3, V3

GPU 3's queries \(Q_3\) need earlier keys and values.

Instead of gathering every \(K\) and \(V\) onto every GPU at once, the GPUs can circulate K/V blocks.

Conceptually:

round 1:

GPU 0 uses K0,V0
GPU 1 uses K1,V1
GPU 2 uses K2,V2
GPU 3 uses K3,V3

        ↓ rotate K,V blocks

round 2:

GPU 0 receives another K,V block
GPU 1 receives another K,V block
GPU 2 receives another K,V block
GPU 3 receives another K,V block

        ↓ compute partial attention

        ↓ rotate again

After enough rounds, each query block has interacted with every key/value block it is allowed to attend to.

This family of ideas is often called ring attention.

The key insight is:

\[ \boxed{\text{move pieces of context through the GPUs instead of replicating the whole context everywhere}} \]

But softmax seems to require all scores at once¶

There is a subtle problem.

For one query, attention probabilities are:

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

If the scores \(s_j\) are spread across GPUs or arrive block by block, how can we compute the denominator without storing every score?

The answer is an online softmax.

Suppose we have processed some score block and maintain:

current maximum \(m\),
current exponential sum \(\ell\),
current weighted value accumulator \(o\).

When a new block arrives with maximum \(m_{\text{new}}\), we can update the running statistics using a new global maximum:

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

The old normalization contribution can be rescaled:

\[ \ell' = e^{m-m'}\ell + \sum_{j\in\text{new block}}e^{s_j-m'}. \]

The weighted value accumulator can be updated similarly.

This means attention can process K/V blocks incrementally without materializing all scores simultaneously.

This same mathematical trick is one of the foundations behind FlashAttention.

A concrete four-chunk example¶

Suppose a sequence has only eight tokens:

A B C D E F G H

Split it across four GPUs:

GPU 0 → A B
GPU 1 → C D
GPU 2 → E F
GPU 3 → G H

Consider the query for token H.

Because attention is causal, H may attend to:

A B C D E F G H

GPU 3 begins with local keys:

G H

It computes the partial attention contribution from those tokens.

Then it receives:

E F

and updates its online softmax statistics.

Then:

C D

and finally:

A B.

At no point does GPU 3 necessarily need one giant score vector stored in memory.

It streams through blocks and maintains the sufficient statistics needed to produce the exact attention output.

That is a useful mental model for distributed long-context attention:

The query stays home; the relevant memory pages travel past it.

Causal masking creates useful asymmetry¶

For bidirectional attention, every query may need every K/V block.

For causal language modeling, token \(t\) only attends to positions:

\[ j\le t. \]

That means early sequence chunks need less information than late chunks.

For example:

GPU 0 queries → only chunk 0
GPU 1 queries → chunks 0,1
GPU 2 queries → chunks 0,1,2
GPU 3 queries → chunks 0,1,2,3

A good context-parallel implementation can exploit this triangular structure rather than performing useless work on masked future tokens.

The exact scheduling becomes a systems optimization problem.

Sequence parallelism versus context parallelism¶

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

A useful conceptual distinction is:

Sequence parallelism often refers to sharding token-position activations for operations that are naturally token-local, frequently in conjunction with tensor parallelism.

Context parallelism usually emphasizes distributing the sequence dimension through attention itself, where tokens on one GPU must interact with K/V information owned by other GPUs.

So:

sequence parallelism
→ "split activation positions"

context parallelism
→ "make attention work across those split positions"

The exact implementation vocabulary varies, but this mental model is robust.

Training memory improves¶

Suppose an activation tensor costs:

\[ M \]

bytes when all \(T\) positions are local.

With sequence parallel degree \(N_{\text{SP}}\), the sharded part ideally falls toward:

\[ \frac{M}{N_{\text{SP}}}. \]

If a long-context training run has 40 GB of sequence-dependent activations per GPU, four-way sequence sharding might reduce the sharded portion toward:

\[ \frac{40}{4} = 10\text{ GB}. \]

Not every activation or temporary buffer scales perfectly this way, so real savings differ.

But the principle is powerful:

\[ \boxed{\text{longer context can be paid for with more GPUs rather than only more HBM per GPU}} \]

This becomes increasingly important as context windows grow.

During autoregressive inference, we store keys and values from previous tokens.

For one layer, a simplified KV-cache size scales like:

\[ M_{\text{KV}} \propto B \times T \times N_{\text{KV heads}} \times d_{\text{head}} \times 2 \times \text{bytes}. \]

The factor of 2 is for keys and values.

The crucial term is:

\[ T. \]

Double context length and, all else equal, KV-cache memory roughly doubles.

So distributing sequence-related state can also matter enormously for inference serving.

Long context is not merely a compute problem.

It is a memory-placement problem.

Why sequence parallelism is not free¶

As always, splitting something creates a boundary.

Without sequence sharding:

token activations
→ local

With sequence sharding:

token activations
→ distributed

Whenever an operation needs information across sequence shards, communication appears.

So the trade becomes:

\[ \boxed{\text{less activation memory per GPU} \leftrightarrow \text{more communication}} \]

This should feel familiar.

Activation checkpointing traded memory for compute.

FSDP traded memory for parameter communication.

Tensor parallelism traded large local compute for frequent partial-result communication.

Sequence parallelism trades sequence-local memory for redistribution and cross-context communication.

Put all the cuts on one tensor¶

Consider again:

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

We now recognize three obvious axes:

\[ B,\qquad T,\qquad d. \]

Different parallelism strategies can attack each:

\[ B \rightarrow \text{data parallelism}, \]

\[ T \rightarrow \text{sequence/context parallelism}, \]

\[ d \rightarrow \text{tensor parallelism}. \]

And the Transformer stack itself has another axis:

\[ L \rightarrow \text{pipeline parallelism}. \]

So a useful abstraction is:

\[ \boxed{ \text{LLM parallelism is largely the art of partitioning }B\times T\times d\times L } \]

That is not the whole story—parameters, experts, optimizer state, and other structures add more axes—but it captures a surprising amount.

A 1,024-GPU thought experiment¶

Suppose a training job uses:

\[ N_{\text{DP}}=16, \]

\[ N_{\text{PP}}=8, \]

\[ N_{\text{TP}}=4, \]

and:

\[ N_{\text{CP}}=2. \]

Then:

\[ N_{\text{GPU}} = 16\times8\times4\times2 = 1024. \]

What are those GPUs doing?

16 data-parallel groups
→ different training examples

within each replica:

8 pipeline stages
→ different layer ranges

within each stage:

4 tensor-parallel shards
→ different hidden/matrix pieces

within those groups:

2 context-parallel shards
→ different sequence regions

A single logical training example may therefore involve many GPUs simultaneously.

Yet mathematically we still describe the model simply as:

\[ p_\theta(x_{t+1}\mid x_{\le t}). \]

This gap between the simplicity of the mathematical definition and the complexity of the physical execution is one of the defining features of modern LLM systems.

The deeper mental model: tensors have geography¶

Early in learning neural networks, we think about tensors mostly in terms of:

shape
dtype
values

At large scale, a tensor has another property:

\[ \boxed{\text{where is each piece physically located?}} \]

For:

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

you should increasingly imagine a map:

batch dimension    → which replica?
sequence dimension → which context shard?
hidden dimension   → which tensor-parallel rank?
layer dimension    → which pipeline stage?

Two mathematically adjacent values may live on different GPUs.

An operation is cheap when the values it needs are already colocated.

An operation becomes expensive when the required values cross machine boundaries.

That gives us today's main principle:

\[ \boxed{\text{at scale, tensor layout is part of the algorithm}} \]

The algebra tells us which values must interact.

The distributed system decides where those values live and how they meet.

And long-context LLMs make this especially visible:

When the sequence becomes too large to live comfortably in one place, the context itself becomes a distributed object.