Skip to content

Packed sequences: a mask that reads the batch

A proposal for M12, now built. It was written first, as M10's and M11's were, because it decides what the training figures mean. Before it, they counted the attention of one document filling the whole sequence. Pretraining is rarely run that way.

Pretraining does not pad. It concatenates documents until the sequence is full and trains on the result, so one row of a batch holds the end of one document, several whole ones, and the start of another. Whether attention may cross from one document into the next is a choice. Llama 3 chose not: its model section says it masks attention between documents in the same sequence (§3.2). It found the mask mattered little in ordinary pretraining and a great deal in continued pretraining on very long sequences. Zhao et al. measure the same choice and call it intra-document causal masking. Krell et al. packed BERT's sequences this way for the speed.

The choice also changes what a step costs, and nothing in the tool can say so yet.

What the engine assumes today

  • A training sequence is one document. The attention a step pays for is fixed by T, causal, window and a mask of positions. Every per-token figure, and so every training time and cost, is computed that way.
  • A mask reads positions, heads and the batch index, never a tensor. A score can read one since M11, T5's bias table. A mask cannot, because a mask is counted by evaluating it, and a tensor's values are not the engine's to know. attention-expressions.md lists "inputs at run time" among the things the language cannot say yet, with document masking as the example.
  • Training and serving are measured at the same attention. The operating point already separates them where it matters: a training dtype and a serving dtype. There is no way to say that the training batch is packed and a request is not.
  • The generated attention keeps its block masks. expression_attention caches one per mask function and shape. That is right for a mask of positions and wrong for one that reads the batch.

None of this is wrong for what it describes. It describes one document at a time.

What packing changes

The keys a query keeps

Under a document mask, a query keeps only the tokens before it in its own document. For documents of length L that tile the sequence exactly, a query keeps L/2 keys on average, where causal attention alone keeps T/2.

Llama-3-8B at 8,192 tokens:

Keys a query keeps Attention per token Share of the forward pass
One document 4,096 2.15 GFLOP 12.5%
1,024-token documents 512 268 MFLOP 1.8%

Training FLOPs per token fall by 11%. That is the training-time estimate getting an eighth of its attention term right, and it moves at every sequence length a long-context run uses.

The mean is not enough

A token is more likely to land in a long document than a short one, simply because a long document holds more tokens. For lengths with mean μ and coefficient of variation c, a query keeps on average

μ(1 + c²) / 2

keys, before the sequence's edges cut documents short. Fixed lengths (c = 0) give μ/2. Exponentially distributed lengths (c = 1) give μ, twice as much for the same mean. Real corpora are heavier-tailed than that. So the operating point has to describe a distribution, not just a length.

The kernel's blocks

FlexAttention computes the score matrix in 128 × 128 blocks. It skips a block the mask removes entirely, computes a block it keeps entirely, and computes and then masks a block that straddles a boundary. With a causal mask alone at 8k, only the sixty-four blocks on the diagonal straddle, and the waste is under 2%. With documents, every boundary makes straddlers:

Documents (aligned to blocks) Scores kept Scores computed Ratio
1,024 tokens 524,800 589,824 1.12×
256 tokens 32,896 49,152 1.49×

These are per document. Documents that do not start on a block boundary do worse. attention-expressions.md already names block granularity as a gap, and document masking is where it stops being small.

Phase 2 counted it. A stream starts its documents anywhere, so nearly every boundary cuts through blocks. Averaged over every phase, fixed documents at 8,192 cost the kernel:

Documents Scores computed ÷ scores kept
2,048 tokens 1.19×
1,024 tokens 1.37×
512 tokens 1.74×
256 tokens 2.48×
128 tokens 3.94×

The aligned table above understates these.

The design

A documents input

input gains a role. It is tokens by default, which is every design today. The new role is documents: a B T integer tensor holding each position's document index. A mask reads it the way M11's score reads a table, by the port's name:

doc(b, q) == doc(b, kv)

wired from the documents input into the attention, through every block above it as the tables were. A mask may read a tensor only when the engine knows how to make one up. A documents input qualifies; a learned table still does not.

Packing is a condition of training

The operating point gains packing. It is off by default: one document per sequence, today's numbers exactly. When on, it gives documents of mean length μ and spread c, drawn from a gamma distribution, so c = 0 is fixed and c = 1 exponential.

It applies to training only: trainPerToken, the training time, the cost and the token budget. The serving figures keep one document per sequence, because a request is one. The analysis reports the packed attention beside the unpacked, as flops.packed. Like perStream, it is omitted when there is nothing to say, so a design without a documents mask reports what it does today.

Packing a design that has no document mask changes nothing, since attention then crosses documents. A rule says so as a note, since Llama 3 found the mask optional at ordinary lengths.

Counting by evaluating, as before

The mask is still measured by sampling it, with the same seeded, stratified draws. For each sampled batch row, the engine first draws a packing from the operating point's distribution, then evaluates the mask against it. The closed forms keep the sampler honest: - exactly L/2 for fixed documents that tile the sequence; - the μ(1 + c²)/2 above, corrected for the edges, for the rest.

Beside the share, the engine reports what the kernel computes: the same draws, rounded out to whole 128 × 128 blocks.

Generated code

  • The model's inputs. The model takes the documents as an input: forward(self, ids, documents).
  • The mask. A mask that reads them is a factory, like a score that reads a table. Called with the documents, it returns the mask_mod.
  • The block mask. It is built once per forward pass and reused by every layer, as FlexAttention's FAQ advises, and not kept past the batch it was built for. The generated helper's cache is keyed by the mask function, and a factory makes a new one every call. Left alone, it would miss every time and grow by one mask per step.
  • The eager fallback indexes the same tensor.

Verification

  • Isolation. The runtime makes up a packing for the documents input. With it, a generated model must show that changing the tokens of one document moves no other document's outputs, in the fused form and the eager one.
  • Block counts. FlexAttention's own BlockMask, built for the same packing, must agree with the engine's block figure.

Positions that restart

Hugging Face's DataCollatorWithFlattening restarts position_ids at each packed example. FlashAttention's variable-length kernel then treats every example as its own sequence. The same belongs here: a positions role on input, read by rope in place of the index. It changes what the model computes and not what it costs, so it comes after the cost is right. Phase 3 found that a rotation under a document mask does not need it, and a learned position table does; see below.

The editor

The operating point gets the packing control. The mask preview in the inspector draws a sampled packing, so the block-diagonal shape is visible. The walkthrough says what the mask is for. Built in phase 4, below.

The phases

  1. Documents, measured.
  2. The documents role, a mask that reads it, packing in the operating point (engine, CLI and MCP), counting by sampling, and the split between training and serving.
  3. The test is the same as M11's first phase: every golden unchanged, because no design has a documents input and packing is off by default.
  4. The fixed-length closed form must reproduce exactly, and the gamma forms within the sampler's error.

Done. - The input and the mask. input has a role, tokens by default. A mask may read a tensor at two indices, a row and a position, and nothing else. sdpa, gqa_attention and transformer_block grow a B T integer input for it, as they grow one for a score's table. - The operating point. packing: { mean, spread } in the engine, --pack and --pack-spread on the command line, packing on the MCP tools. With it, flops.packed holds the training figures and the cost is counted from them. Everything else stays one document a row. It is present only when a mask reads the documents. - The rule. documents follows what reaches such a mask back through every stack it was handed into, and objects unless it starts at an input whose role is documents. Token ids have the same shape and element type, so nothing else would catch them wired there. - Golden files. Every analysis, rules and codegen golden is unchanged. Only the catalog's prose moved, for the new parameter.

What building it found: - A row begins inside a document. A stream cut into rows starts each row wherever the last one stopped, so the pieces at both ends are shorter documents as far as the mask can tell. Fixed 1,024-token documents at 8,192 keep 491.1 keys a query, averaged over every phase a row can begin at, not the 512 of a row that begins on a boundary. That is what the analysis reports and what the test holds it to, exactly; the proposal's 512 is the special case. - The long documents are the noise. The first sampler, 64 rows of 32 queries each, was off by 5% for exponential lengths, because which queries land in long documents varied from draw to draw. Now it is: - 1,024 rows cut from a circular stream of at least 1,024 documents, whose lengths are one from each stratum of the distribution, shuffled; - 16 queries a row; - for each query, keys in bands that widen by four going back from it.

 That is the same quarter of a million evaluations. Fixed lengths are
 within half a percent of the exact count. Spread-out lengths are within
 about two percent of a brute-force stream of two hundred thousand
 documents, which itself agrees with the closed form for exponential
 lengths to the token. It takes about 100 ms in the editor the first time
 and nothing after, being cached.
  • The reference needs documents, not rows. A brute-force reference cut from three thousand documents was off by two percent through its long documents alone, however many rows it cut.
  • Generated code came early. A mask that reads the documents is a factory called with them, as a score that reads a table is, and the eager form keeps documents apart. On CUDA each call still builds, and keeps, a block mask of its own. The generated model says so in a warning, and building one per batch is what phase 2 is for.
  • Generated and verified.
  • The factory, the block mask built once per forward pass, and the runtime's made-up packing.
  • The isolation test in both forms, and BlockMask's own count against the engine's block figure.

Done. - The block mask, once a batch. The generated helper's cache used to be keyed by the mask function, which a factory makes anew every call. Now: - A mask factory marks what it returns with the tensors it read. - _block_mask builds the block mask for the first layer that asks, and shares it with every later layer handed the same documents: the same tensor, at the same version. - It replaces that entry when the next batch's documents arrive, so it keeps one block mask per mask, not one per step. - The runtime's packing. tensorcad_runtime.packing draws documents the way the engine does: gamma-distributed lengths laid end to end, each row a window of the stream. verify feeds it to a documents input: - its lengths average a quarter of the row, exponentially spread; - the forward pass, the profiled FLOPs and the export all run on it. smoke-train and trace feed token ids alone. For a design with more than one input, they now say so rather than fail on the call. - The kernel's blocks. - flops.packed.fwdAttentionBlocks is the attention with every 128 × 128 block holding a kept score computed whole. - Whether a block holds one is decided for the whole block at once, over the ranges of positions it covers, by interval arithmetic on the mask. A row's document ids never fall, so reading them over a range is the first and last document in it, and for a mask of documents and positions the answer is exact. For any other mask it can only count a block the kernel would skip, never miss one it computes. - A test checks that against evaluating every point of thousands of random boxes, and against the exact count for fixed documents. - Held against PyTorch. On llama-3-8b scaled to two million parameters, with the mask in every layer, at 512 tokens in 96-token documents: - One document's tokens changed move no other document's outputs, and no other row's, by exactly zero, over a whole forward pass. - flex_attention, given the model's own mask and a block mask create_block_mask built, computes what the eager form computes, also to zero. - Over 256 rows of the runtime's packing, the scores kept, counted one by one, are within 0.7% of the engine's figure for fixed lengths and 2.6% for exponential ones. - The blocks create_block_mask says a kernel computes are within 0.1%. - The block figure is 4.9 times the scores kept there: 96-token documents against 128-token blocks. 3. Positions that restart. The positions role and rope reading it, held against Hugging Face's flattening collator.

Done. - What was added. A positions role on input, and positions on rope, gqa_attention and transformer_block, which carry a pos input down to both rotations. Generated code turns by it through a helper of its own, so a design that rotates by the index generates exactly what it did. ROPE-02 refuses positions on an attention with no rope to turn. - The rule. The phase 1 documents rule became input-roles, which also follows a rotation's positions back to an input that says it holds them. - The runtime. It draws the positions of the same packing it draws documents from, and positions_of is exactly the collator's position_ids for whole examples. - The cost. Nothing: restarting the positions changes no count.

What building it found is that the proposal had the reason slightly wrong. A rotary model under a document mask needs no positions that restart. Its attention sees only the distance between two tokens of one document, and the distance is the same whichever number either of them starts from. The probe packs a row as the collator does and runs it three ways. Held against each document run alone:

Positions Rotary Learned (pos_embedding)
Restarting 6e-7 8e-7
Counted straight through 7e-7 1.25, on outputs of 1.9

So pos_embedding takes positions too: a learned table is looked up by the position itself, and there restarting is the whole difference. For rotary positions it is precision. Counted from 131,072 rather than zero, a document's outputs differ by 7e-5 in float32, and by 4e-4 from a million. 4. The editor and the rules. - The packing control, the sampled preview and the walkthrough. - Two rules: packing without a document mask, and documents short enough against the kernel's blocks that the share understates the work by more than a quarter.

Done. - The packing control. The operating point takes a packing: docs, the mean document length, and spread, which starts at 1 because real corpora vary. It shows for a design with a documents input, and whenever a packing is set, so one carried over from another design can always be taken off. A stored packing the engine would refuse is dropped when the editor loads. - The readout. Under the training figure it shows the packed training figure, which the cost is counted from, and the packed attention with how much more a kernel computes in whole blocks. - The preview. The mask preview draws one row of the packing, marked as such. For Llama-3-8B at 8k in 1,024-token documents that is the block diagonal: every diagonal block, nothing above it, nothing more than four 256-token cells below it. - The walkthrough. Its attention step says the mask keeps packed documents apart, and that the positions restart when they do. Under a packing it adds what the attention costs and what the kernel computes. - packing-unused notes a packing on a design where every attention reads across documents. - document-blocks measures on the whole forward pass, not only on attention, because phase 2 showed the proposal's threshold would fire on nearly everything. It notes a ratio of whole blocks to kept scores past 1.25. It warns only once the difference is a twentieth of the forward pass. - Llama-3-8B in 1,024-token documents at 8k is noted: 1.38x, 0.6% of a token's work. - nano-sort in 128-token documents, whose attention is a large share of it, is warned about.

Done when: - Llama-3-8B with a document mask, at 8,192 tokens and fixed 1,024-token documents, is counted at exactly 512 keys a query. Phase 1 found that a row cut from a stream begins inside a document, which makes it 491.1, and the analysis is held to that. - Gamma-distributed packings agree with the size-biased mean. - A generated model keeps its documents apart in both forms. - FlexAttention's block count matches the engine's. - With packing off, every golden is unchanged.

What this leaves open

  • Real length distributions. A gamma distribution is two numbers. A corpus is a histogram, and FineWeb-Edu's is not a gamma. Reading one from a file is the obvious next step, and it is left out until the two numbers have been used.
  • Which figure is the headline. The scores a mask keeps are what the design asks for; the blocks are what the kernel does. This proposal reports both and keeps the first as the headline, as every mask is counted today. If the gap is large for realistic packings, the headline should move, for windows too.
  • Llama 3's own presets. The paper trained with the mask. Carrying it in llama-3-8b and its siblings would make them say what they did, and would change their generated forward to take documents. That is a change to three presets' code, for a mask the paper found made little difference at their length.
  • Training with real boundaries. smoke-train reads its corpus as one stream of tokens and does not know where documents end, so it cannot train a design that keeps them apart, and says so. Recording the boundaries as the corpus is tokenized would let it.
  • Context parallelism. Splitting a packed sequence across devices balances badly when documents are uneven. The planner does not model that, and this does not change it.

Sources