Spot the masking bug in this packed-sequence training setup

Spot the error·Medium·4.0 · 0·~2 min

Attempt it

Click any words you think contain an error. Click again to unmark.

Mark at least one word to submit.

**Packed sequences** are a memory-saving trick used during training. Instead of padding every example to the maximum length in a batch (which wastes compute on padding tokens), you concatenate several short examples end to end into one longer sequence and process them together. The trap is the attention mask. A standard **causal mask** is lower-triangular: position t can attend to positions 0 through t. That works fine for a single example, where 'positions 0 through t' are all part of the same context. But in a packed sequence, positions 0 through t may include tokens from a completely different example that was concatenated earlier. The causal mask happily lets position t attend backward to those tokens, treating them as legitimate context. The model learns from contaminated context and the next-token loss is computed against the wrong distribution. The fix is a **document mask** (also called a block-diagonal mask or example mask). It zeroes out attention between any two positions that belong to different examples. Combined with the causal mask, you get a block lower triangular pattern where each block on the diagonal corresponds to one packed example. FlashAttention-2 exposes this via the `cu_seqlens` (cumulative sequence lengths) API: you pass the example boundaries as a 1D tensor, and the kernel handles the block-diagonal masking efficiently without materializing the full mask matrix.

The bug is the assumption that the causal mask alone suffices for packed sequences. It does not. A **causal mask** blocks attention to **future** positions. In a single example, that is exactly what you want: position t learns from its own past, predicts the next token, and the loss is well-formed. In a **packed sequence**, example B sits to the left of example A. Under a plain causal mask, the first token of example A can attend backward to the entire content of example B, because example B is in A's past. The model treats B as legitimate context for predicting A's next token, which is wrong on two axes: B is unrelated to A, and B is generally a complete document that already ended. The fix is to layer a **block-diagonal document mask** on top of the causal mask. The combined mask zeroes attention between any two positions in different examples, producing a block lower triangular pattern where each example forms a self-contained causal triangle on the diagonal. **FlashAttention-2's `cu_seqlens` API** is the production solution. You pass cumulative sequence lengths as a 1D tensor (e.g., `[0, 17, 42, 89]` for three packed examples of lengths 17, 25, and 47), and the kernel computes attention per example without materializing the full mask. The memory and compute cost is the same as separate per-example attention but with the throughput of batched processing. **Why this matters.** A model trained with leakage learns to use cross-example context that is not available at inference, when the model sees only one example at a time. The mismatch shows up as quality regressions that scale with packing density.

Quick reference

TL;DR

A plain causal mask lets each example attend backward into earlier packed examples, leaking signal. Layer a block-diagonal document mask on top, or use FlashAttention-2's cu_seqlens API.

Memory aid

Easy to grasp

Imagine printing five different children's stories on the same long roll of paper, end to end, to save paper. You then hand the roll to a student and ask them to predict the next word at any point, only looking to the left of where they are. Without any marks separating the stories, the student happily looks left and gets to read the end of story 4 as 'context' for the beginning of story 5, which is completely the wrong context. The fix is to draw a thick vertical line between every story so the student knows they can only look left as far as the most recent line. Packed-sequence attention works the same way: the lines between examples must be enforced, not just inferred.

Key concepts

Causal mask Document mask (block-diagonal)Packed sequences FlashAttention

Detailed answer & concept explanation

~7 min read

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.

Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example.

Quick scan · 2–4 min

Everything important, quickly.

Explain the causal-only mask's blind spot at example boundaries, motivate the document mask as a block-diagonal overlay, show how the masks compose via elementwise AND, walk through FlashAttention-2's cu_seqlens implementation, and close with the production checklist for packed-sequence training.

Code example

python

import torch
from flash_attn import flash_attn_varlen_func

# Three examples packed into one sequence of length 17 + 25 + 47 = 89.
example_lengths = [17, 25, 47]
cu_seqlens = torch.tensor(
    [0] + list(torch.tensor(example_lengths).cumsum(0).tolist()),
    dtype=torch.int32,
)  # [0, 17, 42, 89]

# Q, K, V are packed: shape (total_tokens, n_heads, d_head)
total_tokens = sum(example_lengths)  # 89
n_heads, d_head = 8, 64
Q = torch.randn(total_tokens, n_heads, d_head, dtype=torch.float16, device='cuda')
K = torch.randn(total_tokens, n_heads, d_head, dtype=torch.float16, device='cuda')
V = torch.randn(total_tokens, n_heads, d_head, dtype=torch.float16, device='cuda')

# FlashAttention-2 applies causal + block-diagonal (per cu_seqlens) masking.
output = flash_attn_varlen_func(
    Q, K, V,
    cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
    max_seqlen_q=max(example_lengths), max_seqlen_k=max(example_lengths),
    causal=True,
)
# output has shape (total_tokens, n_heads, d_head)
# Each example's tokens attend only to earlier tokens in the same example.

Real-world examples

Real products, models, and research that use this idea.

FlashAttention-2 (Dao 2023) introduced the cu_seqlens API specifically to support efficient packed-sequence training.
Megatron-LM, DeepSpeed, and most production training stacks support packed sequences with document masking via cu_seqlens or equivalent.
Llama 4 Maverick and similar 2026 frontier models train with packed sequences at high density (8x to 16x) using document masking.
Mistral and Qwen 3.5 use packed-sequence training during pretraining and SFT phases.
PyTorch's FlexAttention API (2024) exposes arbitrary mask functions including block-diagonal patterns as first-class inputs.

Interviewer follow-ups

What an interviewer would ask next. Try answering before peeking at the approach.

QHow does FlashAttention-2's cu_seqlens API avoid materializing the (T, T) document mask?

QWhat goes wrong if you reset positional encoding at example boundaries but forget the document mask?

QWhy does packed training give such a large throughput win when it requires this extra masking machinery?

Don't say this

Red flags and common mistakes that signal junior thinking. Click to expand.

Most common mistake

Assuming the causal mask alone is sufficient because pad tokens are gone. The boundary between concatenated examples is exactly the leak point a causal mask cannot see.

60-second cheat sheet

The night-before-the-interview bullets. Scan these on the way to the call.

References

Primary sources. Skim if you want the original framing.