Pack sequences without joining their contexts

The problem: shared storage can change context

Right padding preserves unequal examples, but short rows can leave unused storage. Placing several short examples in one row can fill those cells. An ordinary causal mask, however, lets each query use every earlier key in that row. A later example would then receive the earlier example as additional context. Shifting concatenated source documents would also create next-token targets across their boundaries.

We need to preserve the examples while moving their storage. Sequence packing places complete, already formed input-target pairs into shared rows. Each occurrence keeps its own segment identity and local positions, so storage adjacency creates neither an attention edge nor a new target.

Form pairs inside each source range

Use three supplied encoded documents:

The existing layout uses 0 for BOS and 1 for EOS. These arrays are controlled valid-ID examples in the $266$-entry fixture vocabulary. Their token counts are $4$, $3$ and $5$; their prediction lengths are $3$, $2$ and $4$. A selected range supplies one fewer next-token pair than source tokens.

Form each pair before packing. A’s final input 100 still predicts its own EOS 1; it does not predict B’s BOS. Globally shifting the concatenated source documents would add an EOS-to-BOS pair at each join. Those targets do not belong to the examples we selected. An explicitly selected interior range follows the same rule: retain its existing neighboring pairs without adding BOS or EOS.

Fill only the current row

Capacity $6$ means six prediction pairs per storage row. Our next-fit policy reads examples in their supplied order. A fills columns zero through two, B fills columns three and four, and the last cell is padding. C does not fit in that one remaining cell, so it starts a new row and leaves two padding cells. The packer never revisits an earlier row, sorts by length or splits an example.

The two input rows are [0,99,100,0,101,-1] and [0,102,103,104,-1,-1]. Their targets are [99,100,1,101,1,-1] and [102,103,104,1,-1,-1]. The rectangle stores $12$ cells and contains $9$ real targets. PAD_CELL is -1, a storage marker rather than a vocabulary ID.

A, B and C receive occurrence segment IDs 0, 1 and 2. The segment rows are [0,0,0,1,1,-1] and [2,2,2,2,-1,-1], with NO_SEGMENT also -1. The position rows are [0,1,2,0,1,0] and [0,1,2,3,0,0]. B therefore begins at local position zero even though its storage column is three. Padding’s stored position zero is unused.

An occurrence is one supplied use of an example. Two uses of the identical source document and range still receive distinct occurrence ordinals and segments. Their source IDs, example IDs and ranges may match; their causal contexts remain separate.

Input validity, target validity and loss eligibility keep the meanings from Chapter 4. Every real target is supervised here, so the three row-zero fields are [true,true,true,true,true,false]; the row-one fields are [true,true,true,true,false,false]. Input validity says an input exists, target validity says its immediate successor exists, and loss eligibility says that target is supervised. Their agreement in this fixture does not make the fields interchangeable.

Permission depends on the example

The relation applies within one row. A query may use a key only when both positions are valid, both belong to the same occurrence segment, and the key’s local position is no later than the query’s. A binary value of one means permitted; it is not an attention probability.

This is block-causal permission: the allowed entries form one causal block per occurrence.

For B’s first query, storage column three, the only permitted key is column three. The earlier keys in columns zero through two belong to A and are excluded. B’s next query permits columns three and four. Merely checking that the key precedes the query in storage would admit A incorrectly.

Both validity checks matter. Padding cells share segment -1 and stored position zero, so segment equality and local causality alone would allow padding-to-padding edges. An inactive query instead skips softmax and produces a zero attention mixture. A later bias or residual may produce a nonzero padded hidden state; validity must prevent that state from influencing real predictions or loss.

What the permission symbols mean

For a contiguous segment, the allowed keys form an inclusive range from its first storage column through the query column. B’s segment starts at column three, so its two ranges are [3,4) and [3,5) in Rust’s half-open notation. The checked start and valid length represent the permission relation without allocating a separate dense mask for every row and attention head.

Position reset and segment isolation are separate requirements. Assigning B row-global positions [3,4] violates its local-position metadata. It does not necessarily change rotary attention scores: with fixed unrotated queries and keys, a common offset applied to all their positions in an isolated segment preserves ideal relative rotary dot products. We check the local positions directly. Resetting positions alone also cannot exclude A’s keys.

Keep the same nine-target objective

Packing retains the raw loss sum and valid-target count from Chapter 4. The token-weighted loss gives every eligible target equal weight and conditions on its own permitted segment prefix. Specify gold probability $1/2$ for all three A targets, $1/4$ for both B targets and $1/2$ for all four C targets. These are constructed loss-probe probabilities, rather than model predictions.

The negative log likelihood sums are $3\ln2$, $4\ln2$ and $4\ln2$ nats. Adding them and dividing once gives $11\ln2/9=(11/9)\ln2$ nats per valid target, printed as 0.847180 by the probe. The rows contain five and four targets: averaging their means equally would give $(6/5)\ln2$, printed as 0.831777. Averaging the three document means equally would give $(4/3)\ln2$, printed as 0.924196. Both change the weight assigned to individual targets.

The finite-logit probe uses vocabulary size $V=266$. For desired gold probability $0<q<1$, give each nongold class logit zero and the gold class logit $z_y=\ln((V-1)q/(1-q))$. Softmax assigns gold probability $q$ in ideal arithmetic. The Rust calculation uses f64, with rounding and the declared comparison tolerance.

In ideal arithmetic, for eligible target $i$ and vocabulary class $c$, the mean-loss logit derivative is $(p_{ic}-\mathbf{1}[c=y_i])/N$, where $p_{ic}$ is that class’s probability, $y_i$ is the gold class and $N$ is the total eligible-target count. Here $i$ indexes eligible targets across all rows. Excluded logit cells have zero contribution. The loss adapter skips them before target conversion or arithmetic; multiplying an already computed NaN by zero would not exclude it safely. Zero eligible targets is an error because the mean has no denominator.

Extend the existing storage and decoder

PackConfig declares prediction capacity, maximum rows per model call, cell capacity and metadata budget. pack_examples validates the existing SequenceExample inputs and their shared binding, then assigns whole examples to the current row. An oversized example fails instead of being silently split or truncated. An exact fill creates no extra empty row, and the final model call retains its actual smaller row count.

pub fn pack_examples(
    partition: Partition,
    examples: &[SequenceExample<'_>],
    expected: &BatchBinding,
    config: &PackConfig,
) -> Result<PackedBatchEpoch, VariableBatchError> {
    validate_pack_examples(partition, examples, expected, config.storage)?;
    if examples.len() > i32::MAX as usize + 1 {
        return Err(VariableBatchError::Invalid(
            "occurrence segment IDs must fit i32",
        ));
    }
    let planner_bytes = checked_metadata_add(
        checked_metadata_mul(examples.len(), size_of::<BatchPlacement>())?,
        size_of::<Vec<BatchPlacement>>() + ALLOCATION_ALLOWANCE,
    )?;
    enforce_metadata(planner_bytes, config.max_metadata_bytes)?;
    let mut placements = reserve(examples.len())?;
    let mut row = 0usize;
    let mut used = 0usize;
    for (occurrence, example) in examples.iter().enumerate() {
        if example.len() > config.capacity() - used {
            row = row.checked_add(1).ok_or(VariableBatchError::Invalid(
                "packed row count overflows usize",
            ))?;
            used = 0;
        }
        placements.push(BatchPlacement {
            occurrence,
            row,
            start: used,
        });
        used += example.len();
    }
    let row_count = row.checked_add(1).ok_or(VariableBatchError::Invalid(
        "packed row count overflows usize",
    ))?;
    let batch_count = row_count.div_ceil(config.max_rows_per_batch());
    let descriptor_bytes = checked_metadata_add(
        checked_metadata_mul(examples.len(), size_of::<PackedOccurrence>())?,
        checked_metadata_mul(batch_count, ALLOCATION_ALLOWANCE)?,
    )?;
    let mut metadata_peak = checked_metadata_add(size_of::<PackedBatchEpoch>(), planner_bytes)?;
    metadata_peak = checked_metadata_add(metadata_peak, descriptor_bytes)?;
    metadata_peak = checked_metadata_add(
        metadata_peak,
        checked_metadata_mul(batch_count, size_of::<PackedBatch>())?,
    )?;
    metadata_peak = checked_metadata_add(metadata_peak, ALLOCATION_ALLOWANCE)?;
    let mut first = 0usize;
    for batch in 0..batch_count {
        let first_row = checked_metadata_mul(batch, config.max_rows_per_batch())?;
        let rows = (row_count - first_row).min(config.max_rows_per_batch());
        let next_row = first_row + rows;
        let mut end = first;
        while end < placements.len() && placements[end].row < next_row {
            end += 1;
        }
        let cells = checked_metadata_mul(rows, config.capacity())?;
        if cells > config.max_cells() {
            return Err(VariableBatchError::Capacity {
                cells,
                maximum: config.max_cells(),
            });
        }
        // The batch Vec already owns the embedded VariableBatch struct; count its buffers once.
        let storage_bytes =
            VariableBatch::budget_placed([rows, config.capacity()], examples, first..end)?;
        metadata_peak =
            checked_metadata_add(metadata_peak, storage_bytes - size_of::<VariableBatch>())?;
        first = end;
    }
    enforce_metadata(metadata_peak, config.max_metadata_bytes)?;
    let mut batches = reserve(batch_count)?;
    let mut first = 0usize;
    for batch in 0..batch_count {
        let first_row = batch * config.max_rows_per_batch();
        let rows = (row_count - first_row).min(config.max_rows_per_batch());
        let mut end = first;
        while end < placements.len() && placements[end].row < first_row + rows {
            end += 1;
        }
        // Normalize row coordinates in the already budgeted planner, preserving global occurrence IDs.
        for placement in &mut placements[first..end] {
            placement.row -= first_row;
        }
        batches.push(PackedBatch::from_placements(
            partition,
            examples,
            expected,
            [rows, config.capacity()],
            &placements[first..end],
            first..end,
        )?);
        first = end;
    }
    Ok(PackedBatchEpoch::from_parts(
        batches,
        row_count,
        examples.len(),
        metadata_peak,
    ))
}

PackedBatch adds occurrence placement records around the existing VariableBatch. The shared storage builder retains source and example IDs, source ranges, partition, tokenizer identity and split-receipt identity. Its private arrays derive positions, segment starts, validity and targets together. The packed wrapper reuses the same valid-target loss methods. A receipt digest identifies the caller’s accepted evidence; it does not approve that evidence.

    pub(crate) fn from_placements(
        partition: Partition,
        examples: &[SequenceExample<'_>],
        expected: &BatchBinding,
        shape: [usize; 2],
        placements: &[BatchPlacement],
        span: Range<usize>,
    ) -> Result<Self, VariableBatchError> {
        let storage =
            VariableBatch::build_placed(partition, examples, expected, shape, placements, span)?;
        let mut occurrences = reserve(placements.len())?;
        for placement in placements {
            occurrences.push(PackedOccurrence {
                ordinal: placement.occurrence,
                row: placement.row,
                columns: placement.start..placement.start + examples[placement.occurrence].len(),
                segment: i32::try_from(placement.occurrence).expect("checked occurrence namespace"),
            });
        }
        Ok(Self {
            storage,
            occurrences,
        })
    }

The row-local allowed_key method exposes the full binary rule. The model receives a valid prefix length per row and a segment start per query. The existing causal-softmax loops use those ranges in forward and backward calculation; the same boundary reaches every attention head and layer.

    /// Eligibility is stored as lengths and validity, without a dense head mask.
    pub fn allowed_key(
        &self,
        row: usize,
        query: usize,
        key: usize,
    ) -> Result<bool, VariableBatchError> {
        if row >= self.shape[0] || query >= self.shape[1] || key >= self.shape[1] {
            return Err(VariableBatchError::Invalid(
                "attention coordinate is outside this batch",
            ));
        }
        let q = row * self.shape[1] + query;
        let k = row * self.shape[1] + key;
        Ok(self.valid_tokens[q]
            && self.valid_tokens[k]
            && self.segments[q] == self.segments[k]
            && self.positions[k] <= self.positions[q])
    }

pub fn causal_attention_with_ranges(
    context: AutogradContext,
    query: &TensorValue,
    key: &TensorValue,
    value: &TensorValue,
    lengths: &[usize],
    starts: &[usize],
) -> Result<CausalSelfAttentionForward, CausalMaskingError> {
    causal_attention_with_layout(context, query, key, value, Some(lengths), Some(starts))
}

fn causal_attention_with_layout(
    context: AutogradContext,
    query: &TensorValue,
    key: &TensorValue,
    value: &TensorValue,
    lengths: Option<&[usize]>,
    starts: Option<&[usize]>,
) -> Result<CausalSelfAttentionForward, CausalMaskingError> {
    let prepared = scaled_self_attention_scores_with_context(context, query, key, value)?;
    let tokens = query.shape()[1];
    let additive_mask = causal_additive_mask(tokens)?;
    let weights = match (lengths, starts) {
        (Some(lengths), Some(starts)) => prepared
            .scaled_scores
            .causal_softmax_with_ranges(context, lengths, starts),
        (Some(lengths), None) => prepared
            .scaled_scores
            .causal_softmax_with_lengths(context, lengths),
        (None, None) => prepared.scaled_scores.causal_softmax_with_context(context),
        (None, Some(_)) => {
            return Err(autodiff_error(CausalMaskingStage::MaskedSoftmax)(
                crate::autograd::model_ops::ModelOpError::CausalLengthCount {
                    expected: query.shape()[0],
                    actual: 0,
                }
                .into(),
            ));
        }
    }
    .map_err(autodiff_error(CausalMaskingStage::MaskedSoftmax))?;
    let output = weights
        .matmul_with_context(context, value)
        .map_err(autodiff_error(CausalMaskingStage::ValueMixture))?;

    Ok(CausalSelfAttentionForward {
        raw_scores: prepared.raw_scores,
        scaled_scores: prepared.scaled_scores,
        additive_mask,
        weights,
        output,
        scale: prepared.scale,
        key_width: prepared.key_width,
        value_width: prepared.value_width,
        range_starts: starts.map(<[usize]>::to_vec),
    })
}

The existing causal-only additive report still has shape $[T,T]$, where $T$ is row capacity. It describes future columns but does not describe the segment boundaries. The range-aware softmax enforces those boundaries. Dense score and attention-weight tensors remain in this scalar decoder; an implicit permission rule does not eliminate their rectangular work.

forward_packed delegates to the existing padded decoder entry with the shared storage. The embedding gather excludes PAD_CELL, rotary lookup uses local positions, and the decoder runs its existing blocks and tied vocabulary projection. loss_packed reuses the existing eligible-row gather and indexed NLL. Packing adds no tokenizer token, probability kernel or second decoder.

    pub fn forward_packed(
        &self,
        batch: &crate::training::packed_batch::PackedBatch,
        expected: &BatchBinding,
    ) -> Result<DecoderModelForward, PaddedDecoderError> {
        self.forward_packed_with_context(AutogradContext::recording(), batch, expected)
    }

    pub fn forward_packed_with_context(
        &self,
        context: AutogradContext,
        batch: &crate::training::packed_batch::PackedBatch,
        expected: &BatchBinding,
    ) -> Result<DecoderModelForward, PaddedDecoderError> {
        self.forward_padded_with_context(context, batch.storage(), expected)
    }

    pub fn loss_packed(
        &self,
        batch: &crate::training::packed_batch::PackedBatch,
        expected: &BatchBinding,
    ) -> Result<MaskedLoss, PaddedDecoderError> {
        self.loss_packed_with_context(AutogradContext::recording(), batch, expected)
    }

    pub fn loss_packed_with_context(
        &self,
        context: AutogradContext,
        batch: &crate::training::packed_batch::PackedBatch,
        expected: &BatchBinding,
    ) -> Result<MaskedLoss, PaddedDecoderError> {
        self.loss_padded_with_context(context, batch.storage(), expected)
    }

fn constructed_logits(batch: &PackedBatch) -> Result<Tensor, Box<dyn Error>> {
    let vocabulary = batch.binding().vocabulary_size();
    let mut values = vec![0.0; batch.inputs().len() * vocabulary];
    for (cell, &valid) in batch.loss_mask().iter().enumerate() {
        if valid {
            let q = if batch.segments()[cell] == 1 {
                0.25
            } else {
                0.5
            };
            values[cell * vocabulary + batch.targets()[cell] as usize] =
                ((vocabulary - 1) as f64 * q / (1.0 - q)).ln();
        }
    }
    Ok(Tensor::from_vec(
        vec![batch.shape()[0], batch.shape()[1], vocabulary],
        values,
    )?)
}

Run the bounded example

With Docker running and the repository’s Rust workspace setup available, run from the repository root. The command uses the bundled Chapter 4 prepared fixture and split receipt; it needs no corpus download:

COURSE_CORPUS=false ./course run cargo run --offline --locked \
  -p practical-ch05-packed-sequence-masks

The JSON trace contains both packed rows and occurrence records, the allowed-key lists, nine valid targets, and the constructed loss sum 7.624619 nats and mean 0.847180 nats per target. It separately compares a decoder with identical untrained parameters across separate, padded and packed execution. Those model loss sums are each printed as 50.551553 nats; the constructed loss probabilities do not describe this decoder’s predictions.

The model has width $8$, $2$ heads, feed-forward width $16$, $2$ layers, maximum prediction length $16$ and $3448$ parameter coordinates, with seed 42. Across the nine targets, $2394$ valid vocabulary logits and all $3448$ named parameter-gradient coordinates are checked for each comparison against separate execution: $4788$ logit comparisons and $6896$ gradient comparisons in total. Both printed maximum absolute differences round to 0.000000000000; the individual comparisons still use the stated f64 tolerance. No optimizer update or trained-model quality result is implied.

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let trace = practical_ch05_packed_sequence_masks::demo_trace()?;
    println!("{}", serde_json::to_string_pretty(&trace)?);
    Ok(())
}

Count storage and attention separately

For $B$ rows of width $T$, token storage has $BT$ cells, while dense attention has $BT^2$ query-key score positions per head. A causal segment of prediction length $l$ permits $l(l+1)/2$ semantic edges. A permission edge says which relationship is legal; it does not say whether the implementation computes other score positions before masking.

The main example has $75\%$ valid density in either layout. The second rises from $75\%$ to $100\%$, but its packed rectangle has more dense score positions. These counts establish storage and shape differences. They do not establish a speedup.

History: packing and the cost of moving attention data

Rectangular Transformer training batches allocate padding to accommodate unequal sequence lengths, leaving some stored positions without real tokens. The original 2021 version of Packing: Towards 2x NLP BERT Acceleration examines that problem. The authors study BERT pretraining and compare two deterministic methods based on sequence-length histograms. Using multiple sequences in one sample requires additional attention masking and a change to the masked-language-model loss. Those methods and BERT results do not specify our input-order next-fit policy, causal targets or decoder performance.

Later attention work addressed the cost of moving data during the attention calculation. The 2022 version-two FlashAttention paper describes an exact, IO-aware attention algorithm that uses tiling to reduce transfers between GPU high-bandwidth memory and on-chip SRAM. FlashAttention, version 2

Changing the computation order is a different decision from deciding which examples may attend to one another. Neither an execution tile nor shared storage permits a cross-segment edge.

The runnable Rust contrast counts padded and packed storage, then visits the same permitted edge inventory in query order and in tiles of width two. Sorting copies of the inventories checks that the traversal order retains the same relationships. The tile traversal is an illustration; it implements neither FlashAttention’s numerical kernel nor its GPU memory strategy and provides no timing benchmark.

    let tile_width = 2;
    let mut tile_edges = Vec::new();
    for row in 0..2 {
        for query_start in (0..6).step_by(tile_width) {
            for key_start in (0..6).step_by(tile_width) {
                for query in query_start..(query_start + tile_width).min(6) {
                    for key in key_start..(key_start + tile_width).min(6) {
                        if batch.allowed_key(row, query, key)? {
                            tile_edges.push([row, query, key]);
                        }
                    }
                }
            }
        }
    }
    let mut sorted_query = query_edges.clone();
    sorted_query.sort_unstable();
    let mut sorted_tile = tile_edges.clone();
    sorted_tile.sort_unstable();
    if sorted_query != sorted_tile {
        return Err("edge traversals disagree".into());
    }

For a causal decoder, packing preserves each example’s prediction and attention boundaries. A later attention execution strategy must preserve each example’s permitted context. The tiny Rust traversal illustrates edge ordering without reproducing FlashAttention or establishing a speedup.

Read the aligned cells and permission matrices

The figure places each row’s input, target, occurrence segment, local position and validity fields beside its permission matrix. A matrix row selects a query column and a matrix column selects a key column. The matrices expose B’s separate causal block even though A and B share storage.

Optional practice: inspect the preserved relationships

From the repository root, with the same Docker setup and bundled inputs as the example, run the bounded CPU comparisons:

COURSE_CORPUS=false ./course run cargo test --offline --locked \
  -p llm-from-scratch-practical --test practical_ch05_packed_sequence_masks \
  -- --nocapture

The integration suite runs $13$ tests and prints counts for $100$ fixed seeds in each CPU model profile. The smaller profile has width $4$, $2$ heads, feed-forward width $4$, one layer and maximum prediction length $4$. The larger profile has width $8$, $2$ heads, feed-forward width $16$, two layers and maximum prediction length $16$. Both use the $266$-entry vocabulary. The complete capacity-six example uses the larger profile.

Expect 13 passed and 0 failed. Within each case, the comparisons check every valid vocabulary logit, the raw token-summed loss and every named parameter-gradient coordinate across the three representations. Each finite comparison must satisfy absolute tolerance $10^{-10}$ plus relative tolerance $10^{-9}$ times the larger absolute value. The displayed maximum differences are rounded to twelve decimal places; a displayed zero does not prove exact equality. The tests use untrained parameters and take no optimizer step.

With the same working directory and setup, run the activation tests:

COURSE_CORPUS=false ./course run cargo test --offline --locked \
  -p llm-from-scratch-practical --lib packing_tests -- --nocapture

Expect 2 passed and 0 failed. The first test perturbs independent A activations while holding parameters, shapes, B inputs and the B-only objective fixed. B’s logits, B-only loss and all $3448$ named parameter-gradient coordinates remain within the same tolerance; derivatives into A and padding are exactly zero. Removing the segment range guard admits A’s influence. The second test checks exact zero derivatives into future and padding activations for B’s first query. Changing A can still change the whole-batch parameter gradient because A’s own loss contributes to it. Shared-parameter learning and an illegal cross-example context dependency are different relationships.

Keep the example contract as decoder depth grows

Packed rows retain the same selected source pairs, occurrence-local causal prefixes and token weighting as separate examples. The decoder consumes their checked local positions and key ranges through its existing block loop. The next chapter addresses deeper configurable decoder blocks. More layers must retain these boundaries at every attention operation; depth does not repair an incorrectly joined context or target.