Spot the masking bug in this packed sequence training setup

Spot the error·Medium·4.0 · 0·~2 min·Asked atBain Niki Ai Qdrant·Relevant atMicrosoft

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.

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

Concept explanation~2 min read

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

Packed-sequence training is one of the most common throughput optimizations in modern LLM pretraining and fine-tuning. The trick is simple: concatenate several short examples into one longer sequence to avoid wasted compute on padding tokens. The catch is that attention masking, which everyone learns as 'apply a causal mask', is no longer sufficient.

This bug is dangerous because it is silent. The model trains, the loss decreases, and quality regressions only appear in careful evaluation. The fix is conceptually simple (layer a document mask on top of the causal mask) but the production implementation matters because explicit (T, T) mask materialization defeats the memory advantage of FlashAttention.

This deep dive walks the bug, the fix, the production-grade API (FlashAttention-2's cu_seqlens), and the broader checklist of related concerns that come with packed-sequence training.

Mental model: packed sequences need both causal (block future) AND document (block cross-example) masking. The two constraints are orthogonal and both must be enforced.

The causal-only blind spot at example boundaries

A causal mask is lower-triangular: position t in a sequence can attend to positions 0 through t. The intuition is 'attend to the past'.

In a single example, 'the past' is well-defined: the previous tokens of the same document. The model learns to predict the next token from its own prefix, which is exactly the right inductive bias for autoregressive language modeling.

What happens at a packed boundary

Suppose we pack three examples into one sequence:

Example A: tokens at positions 0-16 (length 17)
Example B: tokens at positions 17-41 (length 25)
Example C: tokens at positions 42-88 (length 47)

Under a plain causal mask, the first token of example B (position 17) can attend to positions 0-17, which includes all of example A. The model treats example A as legitimate context for predicting the first token of example B.

This is wrong on two axes:

Semantic mismatch. Example A is unrelated to example B. The model learns spurious correlations between unrelated documents.
Distribution mismatch. At inference, the model sees only one example at a time. Training with cross-example context teaches the model to rely on context that does not exist at inference.

Why this is silent

The model still optimizes the loss, which still decreases because most predictions are within an example and benefit from correct local context. The leakage only affects predictions near example boundaries, which are a small fraction of tokens. Aggregate metrics like perplexity drop only slightly, but downstream eval (specifically tasks that exercise prompt formatting, instruction following, or document boundaries) can regress meaningfully.

The document mask as the fix

FlashAttention-2's cu_seqlens API

The full production checklist for packed-sequence training

Situations where this technique stops working.

2–4 min · Everything important, quickly.

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 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.

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?

Instead of computing attention across the entire packed sequence and then masking, FlashAttention-2 tiles attention per example by iterating over (cu_seqlens[i], cu_seqlens[i+1]) ranges. Each example becomes a separate but batched attention computation on the GPU. The (Q, K, V) tiles outside the diagonal blocks are never touched, so the compute and memory cost is O(sum n_i^2) instead of O(T^2).

2 more follow-ups an interviewer would ask next. Sign in to reveal them.

Red flags & common mistakes

The phrases 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 bullets to scan on the way to the call.

Why a plain causal mask is insufficient for packed sequences
What a block-diagonal (document) mask looks like as a 2D pattern

Primary sources. Browse if you want the original framing.