Sequence packing with block-diagonal masks

Summary

Sequence packing concatenates multiple training examples back-to-back into a single fixed-length sequence and uses a block-diagonal attention mask so each example only attends within itself, eliminating the FLOPs and memory wasted on padding tokens.

Most NLP corpora have heavily skewed length distributions: many short examples, few long ones. With naive padding to the longest example in the batch, the wasted-token ratio is

1 - \frac{ℓ ˉ}{ℓ _{m a x}}

This ratio counts padded token positions, not total attention FLOPs or end-to-end throughput. Dense padded implementations can still compute on masked positions; efficient variable-length kernels can avoid that work. How much packing helps depends on lengths, bucketing, kernels, and the rest of the training step.

The mechanism

Pick a fixed packed length $L$ (e.g., 8192).
Concatenate examples $e_{1}, e_{2}, \dots$ until adding the next would exceed $L$ . Record the boundaries (cumulative sequence lengths, often called cu_seqlens).
Build a block-diagonal attention mask: token $i$ in example $e_{a}$ cannot attend to any token in $e_{b} \neq = e_{a}$ .
Compute attention with a kernel that respects cu_seqlens (FlashAttention-2 supports this natively via the varlen API).
Apply the loss only on response tokens within each example (mask the boundaries and any prompt tokens for SFT).

The worked construction resets position IDs at each boundary. The block-diagonal mask provides attention isolation; position handling must separately match the model and kernel contract.

Learning objective: map packed-example boundaries to reset position IDs and a causal block-diagonal attention mask that permits no cross-example attention.

Learning objective

Earlier physical position does not grant attention across example boundaries

Derivation

1 · PACK TOKENS, KEEP BOUNDARIES

slot	0	1	2	3	4	5
example	e1	e1	e1	e2	e2	e3
position	0	1	2	0	1	0

Boundaries: `cu_seqlens = [0, 3, 5, 6]`
The intervals [0, 3), [3, 5), and [5, 6) recover the three examples without padding.

2 · MASK BY EXAMPLE AND TIME

q \ k	e1:0	e1:1	e1:2	e2:0	e2:1	e3:0
e1:0	A	F	F	X	X	X
e1:1	A	A	F	X	X	X
e1:2	A	A	A	X	X	X
e2:0	X	X	X	A	F	X
e2:1	X	X	X	A	A	X
e3:0	X	X	X	X	X	A

A = attend · F = future · X = other example
Only the three lower-triangular diagonal blocks participate in attention.

first use the cumulative boundaries to split the physical row into logical examples and restart positions at each boundary. Then read each mask row as one query: it may attend backward inside its own example (A), never forward in causal training (F), and never across an example boundary (X). Packing removes padding; boundaries preserve the same attention problem each example had before concatenation. Loss masking is a separate per-token decision. Original worked example informed by Krell et al. (2021) and the FlashAttention variable-length interface.

Check: Why may e2:1 attend to e2:0 but not to the earlier physical slot e1:2?

Read it this way: Tokens 0 and 1 belong to A; tokens 2 and 3 belong to B. B cannot attend to A even though A appears earlier in the packed buffer.
Can B1 read A1 merely because A1 appears earlier in the packed buffer?

Numbers

For the illustrated lengths [3,2,1], padding three examples to length three allocates nine token positions for six real tokens. The six-token packed row has no slack, reducing physical positions from nine to six. That 1.5x position ratio is not a measured GPU speedup.

Compare packed and unpacked steps with matched examples, labels, position semantics, and loss weighting. Benchmark tokens per second, peak memory, and task quality on the actual model, precision, hardware, and kernel. Packing can otherwise change per-example weighting even when the attention mask is correct.

Common implementation pitfalls

Confusing positions with isolation. A correct block-diagonal mask forbids cross-example attention regardless of numbering. Standard fixed-frequency RoPE preserves within-example attention under a common position shift; absolute embeddings or other scaling rules may not. Follow the specific model/API contract.
Wrong target boundaries. Mask any next-token target that would train one example to predict the first token of another. Attention masking and label masking solve different problems.
Mixing prompt and response in SFT without masking. For SFT, mask out prompt tokens from the loss within each packed example.
Assuming the mask saves computation. A dense implementation can still materialize the full $L \times L$ matrix. Use a boundary-aware kernel and measure end-to-end savings rather than inferring them from zero attention weights.

Predict before revealing: keep the causal block-diagonal mask, but change positions from [0,1,2,0,1,0] to [0,1,2,3,4,5]. With standard fixed-frequency RoPE, does this create cross-example attention?

Identify which contract changed

No. The mask still excludes other examples. Each example received only a common position shift, so relative rotary positions and its attention weights are unchanged. Removing the mask would break isolation even with reset positions. Do not extend this invariance claim to arbitrary positional embeddings or scaling variants.

When not to pack

Very long single examples that fill or exceed $L$ on their own (no concatenation possible; padding is already minimal).
When examples have inter-document context that should attend across boundaries (rare).
During inference, use the serving engine’s multi-request batching and KV-cache contracts; training-style concatenation is not a drop-in replacement.

FlashAttention. The underlying kernel that makes varlen attention efficient.
Gradient accumulation. Orthogonal way to grow effective batch size.