Ulysses Sequence Parallelism: A Temporary Change of Layout

The layout that saves memory breaks at attention

Token sharding works for almost everything—except the one operation that needs the full sequence.

Let a sequence have length L, with H attention heads, running across P ranks. For most Transformer operations, the natural layout splits the sequence:

rank 0: tokens 0 … L/P − 1, all H heads
rank 1: tokens L/P … 2L/P − 1, all H heads
...

Each rank holds only L/P token positions. QKV projections, layer norms, residual connections, output projections, and the MLP are all position-wise once their inputs are local. They run without reconstructing the rest of the sequence.

Attention is the exception. For one head:

O_h = softmax(Q_h K_hᵀ / √d_h) V_h

A query at one position needs keys and values from every permitted position in the sequence. A rank that owns only a token slice cannot compute complete attention for its queries on its own.

The obvious repair is to all-gather the sequence before attention. That produces the correct result, but it also puts all L tokens and all H heads on every rank. The memory saving from token sharding disappears at exactly the expensive point.

Use the other independent axis

Multi-head attention has a second dimension that can be partitioned safely: heads.

Heads interact only after their attention results are concatenated and passed through the output projection. During attention itself, head 0 does not need the Q, K, or V values of head 1. That permits a second layout—a head cut:

The state worth tracking is ownership: for each (token, head) cell, which rank holds it? Under the token cut, each rank owns a horizontal band—a slice of tokens, all heads. Under the head cut, each rank owns a vertical band—all tokens, a subset of heads.

For a standard multi-head attention layer, each rank begins with Q, K, and V shaped like [B, L/P, H, d_h]. The first all-to-all logically transposes the token and head partition:

[B, L/P, H, d_h]  →  all-to-all  →  [B, L, H/P, d_h]

Each sender splits its local heads into P groups. Each receiver collects one head group from every sender, then concatenates those token blocks along the sequence dimension. The result is not an all-gather. No rank holds the complete sequence and every head at the same time. Each rank has the full context only for its assigned H/P heads.

That is enough to run ordinary causal or bidirectional attention locally.

One attention block, end to end

Four steps. Two all-to-alls. The attention kernel itself is untouched.

The full sequence of events inside one attention block is short:

1. Keep tokens sharded. A rank owns L/P positions and computes its local Q, K, and V projections. Shape: [B, L/P, H, d_h].

2. Switch to head ownership. The first all-to-all turns local QKV tensors from "some tokens, all heads" into "all tokens, some heads." Shape becomes [B, L, H/P, d_h].

3. Run normal attention. Each rank has full-sequence Q, K, and V for its head subset. The attention kernel can be FlashAttention or another compatible implementation; the mask and attention math remain unchanged. Output: [B, L, H/P, d_h].

4. Switch back to token ownership. A second all-to-all returns the layout to [B, L/P, H, d_h]. Each rank again owns complete hidden states for its local tokens, so the output projection, residual path, and MLP continue in the token-sharded layout.

The shape trace, for each of Q, K, and V:

[B, L/P, H, d_h]  →  all-to-all  →  [B, L, H/P, d_h]
→  local attention  →  [B, L, H/P, d_h]
→  all-to-all  →  [B, L/P, H, d_h]

Attention changes what each cell means (from Q/K/V to O). All-to-all changes which GPU holds it. That separation is the whole method.

The two all-to-alls are Ulysses's own layout transitions. A larger training system may still have other collectives for FSDP, ZeRO, data parallelism, or tensor parallelism.

The interactive walkthrough below lets you step through each stage and watch ownership flip across the (token, head) grid. The static description above is self-contained; the animation verifies what you already understand.

What Ulysses costs, and where it stops helping

The exchange is not free. Know the constraints before you reach for it.

Communication

Ulysses moves packed QKV once before attention and the attention output once afterward. For standard multi-head attention, the original paper counts this as QKV-sized traffic plus output-sized traffic per layer. It trades activation memory for bandwidth, so network topology matters. Good interconnect bandwidth is assumed.

Its advantage is that the exchange preserves the useful sharding factor: when sequence length and Ulysses degree grow together, each rank does not need to communicate the entire global tensor. This is why Ulysses is attractive on fast, well-connected GPU groups.

Head-count constraints

The number of query heads must normally be compatible with the Ulysses degree: H_q % P == 0 is the typical requirement. GQA and MQA need extra care because there may be fewer KV heads than query heads. Some implementations require compatible divisibility; others replicate KV heads to support a larger sequence-parallel group. That relaxes the apparent limit, but adds memory and communication cost. The K/V-head layout may impose the tighter bound.

What Ulysses does not do

• It shards sequence activations. It does not shard model weights, optimizer state, or gradients by itself.
• It preserves exact dense-attention math. It does not make the total L² attention work disappear.
• If the desired parallelism degree exceeds the usable head layout, or if the all-to-all becomes the bottleneck, a ring- or context-parallel approach may be a better fit.

Ulysses is one point in a larger parallelism design, not a replacement for every other strategy.