Chapter 19 · Content revision 1

Accumulate attention without keeping the full grid

Follow a stable running attention summary, reconstruct tiled gradients and compare numerical results and logical storage with the dense reference.

Keep a summary instead of the whole attention grid

The dense reference computes a score and a probability for every query-key coordinate. With twice as many sequence cells, each of those stored grids has four times as many values. Sharing K/V heads reduces the learned and cached K/V lanes; a local window restricts eligible keys. Neither change automatically removes the dense grids from that reference implementation.

An attention output needs the normalized weighted sum of values. It need not retain every weight at once. We can consume keys and values in tiles, carrying a score maximum, a normalization sum and a weighted value vector from one tile to the next. Rescaling the two sums when the maximum changes preserves the relative weight of earlier keys.

Follow the last query through two tiles

The content-only text abcd supplies four constructed feature rows at sequence cells $0,1,2,3$ . There is one query/key feature and two value features per row:

Q=\begin{pmatrix}1\\1\\1\\1\end{pmatrix},\qquad K=\begin{pmatrix}0\\1\\2\\3\end{pmatrix},\qquad V=\begin{pmatrix}1&10\\2&20\\3&30\\4&40\end{pmatrix}.

Row $j$ of $K$ and $V$ belongs to the same text cell. The final query cell is $3$ , whose query vector is $(1)$ . Dividing its dot products with the four keys by $\sqrt{1}=1$ gives scores $(0,1,2,3)$ . All four keys are causally eligible for that query. This raw attention example is separate from the configured RoPE decoder: its key width is $1$ and its value width is $2$ , whereas the decoder’s rotary head width remains even.

Use query and key tile edges of $2$ cells. For the final query, the first key tile contains keys $0,1$ . Subtracting its maximum $1$ gives exponential weights $\exp(-1)$ and $1$ . The state after this tile is

\begin{aligned} m&=1,\qquad \ell=\exp(-1)+1\approx1.367879,\\ u&=\bigl(\exp(-1)+2,\;10\exp(-1)+20\bigr)\\ &\approx(2.367879,23.678794). \end{aligned}

Here $\ell$ is the shifted exponential sum, and the two coordinates of $u$ are the corresponding unnormalized weighted value sums. The second tile contains keys $2,3$ , so the maximum rises to $3$ . Earlier weights were measured relative to $1$ ; they must now be measured relative to $3$ . Multiply both previous sums by $\alpha=\exp(1-3)\approx0.135335$ , then add the new keys. Primes below mark the updated sums:

\begin{aligned} \ell'&=\alpha\ell+\exp(-1)+1\approx1.553002,\\ u'&=\alpha u+\exp(-1)(3,30)+(4,40)\\ &\approx(5.424096,54.240960). \end{aligned}

Only after the final tile do we divide each weighted-sum coordinate by the normalizer. The result is approximately $(3.492653,34.926527)$ . Rust reports tiled output [3.492652734585769,34.9265273458577] and dense output [3.49265273458577,34.92652734585769] for this row. Across all $8$ output coordinates, the maximum absolute difference is $7.105427357601002\times10^{-15}$ , within the declared CPU f64 absolute tolerance $10^{-10}$ .

Construct the four content rows and record two successive key-tile summaries for their final query. rust/demos/practical-ch19-online-tiled-attention/src/lib.rs#online-attention-worked-example

/// Construct four content-only text feature rows; no decoder training is implied.
pub fn worked_example(execution: &TensorExecution, tiles: AttentionTiles) -> ArchiveResult<Value> {
    let text = "abcd";
    let tokenizer = BpeTokenizer::from_merge_pairs(&[])?;
    let ids = tokenizer.encode_utf8(text);
    let q = parameter(&[1, 1, ids.len(), 1], vec![1.; ids.len()])?;
    let k = parameter(
        &[1, 1, ids.len(), 1],
        (0..ids.len()).map(|i| i as f64).collect(),
    )?;
    let v = parameter(
        &[1, 1, ids.len(), 2],
        (0..ids.len())
            .flat_map(|i| [(i + 1) as f64, 10. * (i + 1) as f64])
            .collect(),
    )?;
    let dense = causal_attention_with_ranges(
        AutogradContext::no_grad().with_execution(execution),
        &q.reshape(&[1, 4, 1])?,
        &k.reshape(&[1, 4, 1])?,
        &v.reshape(&[1, 4, 2])?,
        &[4],
        &[0; 4],
    )?;
    let tiled = tiled_causal_attention(
        AutogradContext::no_grad().with_execution(execution),
        &q,
        &k,
        &v,
        AttentionLayout::default(),
        tiles,
        Some(AttentionTraceRow {
            batch: 0,
            head: 0,
            query: 3,
        }),
    )?;
    let maximum = maximum_error(
        tiled.output().value().as_slice(),
        dense.output().value().as_slice(),
    )?;
    Ok(
        json!({"text":text,"query_rows":q.value().as_slice(),"key_rows":k.value().as_slice(),"value_rows":v.value().as_slice(),
        "key_width":1,"value_width":2,"last_scores":dense.scaled_scores().value().as_slice()[12..16],
        "last_reference_probabilities":dense.weights().value().as_slice()[12..16],
        "tile_summaries":tiled.trace(),"last_tiled_output":tiled.output().value().as_slice()[6..8],
        "last_reference_output":dense.output().value().as_slice()[6..8],
        "max_abs_output_error":maximum,"absolute_tolerance":OUTPUT_TOLERANCE,"within_tolerance":maximum<=OUTPUT_TOLERANCE,
        "stats":tiled.stats(),"all_outputs_graph_free":!tiled.output().tracks_gradient() && !dense.output().tracks_gradient()}),
    )
}

Rescale earlier mass before adding a new tile

Fix one active query vector $q$ with $d$ features. Key vector $k_j$ has the same positive feature count $d$ ; $j$ is its zero-based sequence-cell index. Their dimensionless score is $s_j=\langle q,k_j\rangle/\sqrt d$ . Value vector $v_j$ has $d_v$ features, which may differ from $d$ in raw attention.

Let $J$ be the nonempty set of eligible keys in the current tile and $a$ its largest score. Before the tile, $m$ is the maximum score seen, $\ell$ is the sum of exponential weights shifted by $m$ , and $u$ is their weighted value sum. Start with $m=-\infty$ , $\ell=0$ and the zero vector $u$ . A prime marks the updated state. The complete rule is

\begin{aligned}s_j&=\langle q,k_j\rangle/\sqrt d,\\a&=\max_{j\in J}s_j,\quad m'=\max(m,a),\\\alpha&=\exp(m-m'),\quad p_j=\exp(s_j-m'),\\\ell'&=\alpha\ell+\sum_{j\in J}p_j,\\u'&=\alpha u+\sum_{j\in J}p_jv_j,\\o&=u/\ell.\end{aligned}

The factor $\alpha$ converts the previous sums to the new maximum. Each $p_j$ is an unnormalized shifted exponential, not a final probability. $\ell'$ is dimensionless and positive after a contributing tile; $u'$ and the final output $o$ have $d_v$ value coordinates. The last line uses the completed sums after all eligible tiles. Rust explicitly uses $\alpha=0$ for the first contributing tile.

For any already consumed key, the identity $\exp(s_j-m)\exp(m-m')=\exp(s_j-m')$ explains why rescaling preserves its relative mass. Applying the same factor to $u$ preserves its contribution to every output coordinate. Independently normalizing each tile and averaging its output would lose the relative normalization mass of the tiles.

A tile with no eligible keys leaves all three state values unchanged. It does not compute $-\infty-(-\infty)$ . Existing valid causal layouts include self for each active query, so its final normalizer is positive. Inactive padded query rows remain zero and undergo no division.

Rescale the normalizer and weighted value vector before adding the current tile's eligible contributions. rust/crates/llm-from-scratch-practical/src/attention/online.rs#online-normalizer-update

                    for (local_query, query) in (query_start..query_end).enumerate() {
                        let row = lane * s.tokens + query;
                        let lower = layout.starts.map_or(0, |a| a[batch * s.tokens + query]);
                        let mut tile_max = f64::NEG_INFINITY;
                        for (local_key, key) in (key_start..key_end).enumerate() {
                            let index = local_query * columns + local_key;
                            let score = probabilities.as_slice()[index] * scale;
                            finite(score, index)?;
                            probabilities.as_mut_slice()[index] = score;
                            if key >= lower && key <= query {
                                tile_max = tile_max.max(score);
                            }
                        }
                        if tile_max == f64::NEG_INFINITY {
                            probabilities.as_mut_slice()
                                [local_query * columns..(local_query + 1) * columns]
                                .fill(0.0);
                            continue; // An empty tile leaves m, l and u untouched.
                        }
                        let maximum = maxima[row].max(tile_max);
                        let alpha = if normalizers[row] == 0.0 {
                            0.0
                        } else {
                            (maxima[row] - maximum).exp()
                        };
                        normalizers[row] *= alpha;
                        for x in &mut output[row * s.value..(row + 1) * s.value] {
                            *x *= alpha;
                        }
                        let mut eligible = 0;
                        for (local_key, key) in (key_start..key_end).enumerate() {
                            let index = local_query * columns + local_key;
                            let p = if key >= lower && key <= query {
                                eligible += 1;
                                (probabilities.as_slice()[index] - maximum).exp()
                            } else {
                                0.0
                            };
                            probabilities.as_mut_slice()[index] = p;
                            normalizers[row] += p;
                        }
                        maxima[row] = maximum;
                        stats.eligible_scores += eligible;
                        if trace_row == Some(AttentionTraceRow { batch, head, query }) {
                            selected = Some((row, eligible, alpha));
                        }
                    }
                    let weighted_values =
                        execution.matmul(&probabilities.view(), &value_tile.view())?;
                    for (local_query, query) in (query_start..query_end).enumerate() {
                        let row = lane * s.tokens + query;
                        for feature in 0..s.value {
                            output[row * s.value + feature] +=
                                weighted_values.as_slice()[local_query * s.value + feature];
                            finite(output[row * s.value + feature], row * s.value + feature)?;
                        }
                    }

In ideal real arithmetic this evaluates the same attention rule as the dense reference. f64 reductions occur in a different order, so small rounding differences are expected. Very small exponentials can underflow to zero. Finite inputs also do not guarantee representable dot products or weighted sums; the operation refuses nonfinite results rather than returning a partial attention result.

Reconstruct probabilities when a gradient needs them

For zero-based query-row index $i$ , retain its output vector $o_i$ , final maximum $m_i$ and positive normalizer $\ell_i$ . During backward, let $g_i$ be the upstream gradient with respect to $o_i$ ; it has the same $d_v$ coordinates. Recompute scores from the saved query and key vectors, then reconstruct the eligible probabilities:

P_{ij}=\frac{\exp(s_{ij}-m_i)}{\ell_i}.

$s_{ij}$ is the scaled score of query $i$ and key $j$ ; $P_{ij}$ is its dimensionless normalized probability. Masked coordinates have probability zero. For a scalar objective $L$ , use the prefix $\mathrm d$ to denote a gradient of $L$ with respect to the following quantity. Define

\begin{aligned} D_i&=\langle g_i,o_i\rangle,\\ \mathrm dP_{ij}&=\langle g_i,v_j\rangle,\\ \mathrm dS_{ij}&=P_{ij}(\mathrm dP_{ij}-D_i). \end{aligned}

$D_i$ is one scalar per query row. $\mathrm dP_{ij}$ measures how the objective changes with that probability; $\mathrm dS_{ij}$ measures how it changes with the score. The softmax correction subtracts the row-wide weighted mean $D_i=\sum_jP_{ij}\mathrm dP_{ij}$ , which can be obtained from the already saved output instead of storing a probability row.

For a query/key tile, $Q$ and $K$ hold its row vectors, $V$ holds its value rows, and $G$ holds its upstream vectors. Matrices $P$ , $\mathrm dP$ and $\mathrm dS$ have query rows and key columns. Accumulate

\begin{aligned} \mathrm dQ&\mathrel{+}=\mathrm dS K/\sqrt d,\\ \mathrm dK&\mathrel{+}=\mathrm dS^{\mathsf T}Q/\sqrt d,\\ \mathrm dV&\mathrel{+}=P^{\mathsf T}G. \end{aligned}

The superscript $\mathsf T$ transposes the tile. These gradients have the same shapes as their corresponding $Q$ , $K$ and $V$ inputs. Excluded keys contribute zero. When several query heads read one shared K/V head, their key and value gradients add into that same shared head.

Reconstruct each eligible probability tile and accumulate the requested input gradient with the existing matrix executor. rust/crates/llm-from-scratch-practical/src/attention/online.rs#tiled-attention-backward

/// Each parent pullback reconstructs only one score tile and the requested VJP.
pub(crate) fn backward(
    saved: &OnlineAttentionSaved,
    operand: AttentionOperand,
    upstream: &Tensor,
    execution: &TensorExecution,
) -> Result<Tensor, TensorAutodiffError> {
    let execution = execution
        .with_matmul_precision(MatmulPrecision::F64)
        .map_err(crate::tensor::matmul::MatmulError::Execution)?;
    let s = shape(&saved.query, &saved.key, &saved.value)?;
    let target = match operand {
        AttentionOperand::Query => &saved.query,
        AttentionOperand::Key => &saved.key,
        AttentionOperand::Value => &saved.value,
    };
    let mut gradient = Tensor::from_vec(target.shape().to_vec(), zeros(target.len())?)?;
    let scale = 1.0 / (s.key as f64).sqrt();
    for batch in 0..s.batch {
        let length = saved.lengths.as_ref().map_or(s.tokens, |l| l[batch]);
        for head in 0..s.heads {
            let lane = batch * s.heads + head;
            let kv_lane = batch * s.kv_heads
                + kv_head_for_query(head, s.heads, s.kv_heads).expect("validated grouping");
            for query_start in (0..length).step_by(saved.tiles.queries) {
                let query_end = query_start.saturating_add(saved.tiles.queries).min(length);
                let q = rows(&saved.query, lane, query_start, query_end)?;
                let g = rows(upstream, lane, query_start, query_end)?;
                for key_start in (0..query_end).step_by(saved.tiles.keys) {
                    let key_end = key_start.saturating_add(saved.tiles.keys).min(query_end);
                    let k = rows(&saved.key, kv_lane, key_start, key_end)?;
                    let v = rows(&saved.value, kv_lane, key_start, key_end)?;
                    let mut p =
                        execution.matmul_with_transpose(&q.view(), &k.view(), false, true)?;
                    let columns = key_end - key_start;
                    for (i, query) in (query_start..query_end).enumerate() {
                        let row = lane * s.tokens + query;
                        let lower = saved
                            .starts
                            .as_ref()
                            .map_or(0, |a| a[batch * s.tokens + query]);
                        for (j, key) in (key_start..key_end).enumerate() {
                            let index = i * columns + j;
                            p.as_mut_slice()[index] = if key >= lower && key <= query {
                                (p.as_slice()[index] * scale - saved.maxima[row]).exp()
                                    / saved.normalizers[row]
                            } else {
                                0.0
                            };
                        }
                    }
                    if operand == AttentionOperand::Value {
                        let dv =
                            execution.matmul_with_transpose(&p.view(), &g.view(), true, false)?;
                        add_rows(&mut gradient, kv_lane, key_start, &dv);
                        continue;
                    }
                    let mut ds =
                        execution.matmul_with_transpose(&g.view(), &v.view(), false, true)?;
                    for (i, query) in (query_start..query_end).enumerate() {
                        let row = lane * s.tokens + query;
                        let delta = (0..s.value)
                            .map(|f| {
                                g.as_slice()[i * s.value + f]
                                    * saved.output.as_slice()[row * s.value + f]
                            })
                            .sum::<f64>();
                        for j in 0..columns {
                            let index = i * columns + j;
                            ds.as_mut_slice()[index] =
                                p.as_slice()[index] * (ds.as_slice()[index] - delta) * scale;
                        }
                    }
                    match operand {
                        AttentionOperand::Query => {
                            let dq = execution.matmul(&ds.view(), &k.view())?;
                            add_rows(&mut gradient, lane, query_start, &dq);
                        }
                        AttentionOperand::Key => {
                            let dk = execution.matmul_with_transpose(
                                &ds.view(),
                                &q.view(),
                                true,
                                false,
                            )?;
                            add_rows(&mut gradient, kv_lane, key_start, &dk);
                        }
                        AttentionOperand::Value => unreachable!("value pullback handled above"),
                    }
                }
            }
        }
    }
    Ok(gradient)
}

The matrix $O$ contains the output rows $o_i$ . The custom operation joins the existing autodiff tape. Its three parent edges share one immutable snapshot of $Q,K,V,O$ , row maxima and normalizers, layout and tile settings through an owned Rc. No thread-local state is needed. Each operand edge currently reconstructs tiles separately; this repeats work in exchange for not saving quadratic probabilities. Ordinary model execution records no recurrence trace; the worked example explicitly selects just one query for tracing. Existing stale-parent and released-graph refusals still apply.

Attach three operand edges to the existing tape while sharing one immutable saved attention snapshot. rust/crates/llm-from-scratch-practical/src/attention/online.rs#tiled-attention-autograd

/// Q:[B,Hq,T,d], K:[B,Hkv,T,d], V:[B,Hkv,T,dv]. No K/V expansion.
/// Default layout is full causal. Trace is an explicit, single-query diagnostic.
#[allow(clippy::too_many_arguments)]
pub fn tiled_causal_attention(
    context: AutogradContext,
    query: &TensorValue,
    key: &TensorValue,
    value: &TensorValue,
    layout: AttentionLayout<'_>,
    tiles: AttentionTiles,
    trace_row: Option<AttentionTraceRow>,
) -> Result<TiledAttentionForward, TensorAutodiffError> {
    let execution = context.execution();
    f64_execution(&execution)?;
    let mut evidence = None;
    let output = TensorValue::model_operation_with_context(
        context,
        TensorOperation::OnlineAttention,
        [query, key, value],
        |[q, k, v]| {
            let result = forward(q, k, v, layout, tiles, trace_row, &execution)?;
            let saved = Rc::new(OnlineAttentionSaved {
                query: q.clone(),
                key: k.clone(),
                value: v.clone(),
                output: result.output.clone(),
                maxima: result.maxima,
                normalizers: result.normalizers,
                lengths: layout.lengths.map(<[usize]>::to_vec),
                starts: layout.starts.map(<[usize]>::to_vec),
                tiles,
            });
            evidence = Some((result.stats, result.trace));
            Ok((
                result.output,
                [
                    ModelSavedContext::OnlineAttention {
                        saved: Rc::clone(&saved),
                        operand: AttentionOperand::Query,
                    },
                    ModelSavedContext::OnlineAttention {
                        saved: Rc::clone(&saved),
                        operand: AttentionOperand::Key,
                    },
                    ModelSavedContext::OnlineAttention {
                        saved,
                        operand: AttentionOperand::Value,
                    },
                ],
            ))
        },
    )?;
    let (stats, trace) = evidence.expect("successful operation supplies evidence");
    Ok(TiledAttentionForward {
        output,
        stats,
        trace,
    })
}

Separate running summaries from GPU memory movement

Materialized attention stores a score and a probability for every query-key coordinate, so these intermediate grids grow quadratically with sequence length even though the output is a vector per query. This familiar reference makes every weight inspectable, but retaining those grids is an implementation choice rather than a requirement of the output.

In Online normalizer calculation for softmax, version 1 dated 8 May 2018, Milakov and Gimelshein’s Algorithm 3 updates a running maximum and rescales the previous exponential normalizer when that maximum changes. That paper’s normalizer supplies the scalar part of the recurrence; our attention calculation also carries a weighted value vector.

Rabe and Staats’s attention paper was initially released in 2021; the cited version 3 is dated 10 October 2022. Rabe and Staats describe stable weighted attention summaries, query/key chunking and recomputation of chunk intermediates during differentiation. Saving every intermediate summary for ordinary automatic differentiation would lose the intended storage saving, so their account makes recomputation part of the solution.

In FlashAttention, version 1 dated 27 May 2022, Dao and colleagues combine tiling, recomputation and fused GPU execution to reduce attention reads and writes between high-bandwidth memory and on-chip SRAM. That memory hierarchy and fused kernel are additional implementation choices; the course’s host-backed primitive calls do not reproduce them. The cited Rabe/Staats revision is later than this FlashAttention edition.

Stable running summaries allow attention to consume key/value tiles; reconstructing probabilities during backward avoids retaining their complete grid. Keep the existing dense attention as an inspectable reference and add output-oriented tiled evaluation through the same projections, layouts, tensor executor and autodiff tape. Evaluate the same decoder attention with smaller explicit intermediate buffers while keeping logical storage, GPU memory traffic and measured speed as separate claims. The course implements host-backed tiled attention and reports bounded CPU f64 comparisons and logical payload counts. It does not reproduce a fused FlashAttention kernel or measure GPU traffic, speed or language quality.

Change attention evaluation inside the existing model

The raw operation accepts $Q:[B,H_q,T,d]$ , $K:[B,H_{kv},T,d]$ and $V:[B,H_{kv},T,d_v]$ , producing $[B,H_q,T,d_v]$ . Here $B$ counts batch rows, $T$ counts stored sequence cells, $H_q$ counts query heads and $H_{kv}$ counts shared key heads and, separately, value heads. Feature counts $d,d_v$ are positive, as are the head counts, and $H_q$ must be divisible by $H_{kv}$ . Chapter 18’s adjacent group map selects the same shared head in the dense layer, tiled forward and tiled backward. Compact K/V stay unexpanded.

Common preparation still performs the learned projections, rotates $Q,K$ with RoPE, leaves $V$ unrotated and intersects the causal range with active lengths, packed-segment boundaries and any local window. The resulting interval determines which entries contribute to each query’s attention sum; the implementation still reads and scores rectangular prefix tiles that can include excluded entries, then sets their weights to zero before the weighted value product. Sequence-row indices remain distinct from RoPE labels. Tiles change neither the direct causal context nor cache capacity or eviction. Head outputs still merge and pass through the existing output projection.

Reuse prepared projections, rotary positions and causal layout while passing compact shared K/V heads to tiled attention. rust/crates/llm-from-scratch-practical/src/attention/multi_head.rs#compact-grouped-attention

    pub(crate) fn forward_output_with_layout(
        &self,
        context: AutogradContext,
        input: &TensorValue,
        position_offset: usize,
        lengths: Option<&[usize]>,
        starts: Option<&[usize]>,
        positions: Option<&[usize]>,
    ) -> Result<TensorValue, MultiHeadAttentionError> {
        let Some(tiles) = context.execution().attention_tiles() else {
            return self
                .forward_with_layout(context, input, position_offset, lengths, starts, positions)
                .map(MultiHeadAttentionForward::into_output);
        };
        super::online::f64_execution(&context.execution())
            .map_err(multi_autodiff(MultiHeadStage::RestoreHeadOutputs))?;
        let prepared = self.prepare_heads_with_layout(
            context,
            input,
            position_offset,
            lengths,
            starts,
            positions,
        )?;
        let shape = input.shape();
        // The existing preparation repeats identical metadata for every head.
        // Restore one batch copy; K/V tensors themselves remain unexpanded.
        let lengths = prepared.lane_lengths.as_ref().map(|values| {
            (0..shape[0])
                .map(|b| values[b * self.heads])
                .collect::<Vec<_>>()
        });
        let starts = prepared.lane_starts.as_ref().map(|values| {
            (0..shape[0])
                .flat_map(|b| {
                    values[b * self.heads * shape[1]..(b * self.heads + 1) * shape[1]]
                        .iter()
                        .copied()
                })
                .collect::<Vec<_>>()
        });
        let attended = tiled_causal_attention(
            context,
            &prepared.rotated_query_heads,
            &prepared.rotated_key_heads,
            &prepared.projected_value_heads,
            AttentionLayout {
                lengths: lengths.as_deref(),
                starts: starts.as_deref(),
            },
            tiles,
            None,
        )
        .map_err(multi_autodiff(MultiHeadStage::RestoreHeadOutputs))?;
        self.project_head_outputs(context, attended.output())
            .map(|(_, output)| output)
    }

TensorExecution::with_tiled_attention selects this output-oriented path. The older inspectable attention APIs retain their dense reports; compact decoder reports omit block intermediates instead of inventing dense matrices. The initial tiled path requires explicit f64 matrix arithmetic and refuses other precision policies before dispatch. Tile products reuse the existing CPU/CUDA primitive boundary. All results here are CPU results; CPU loading and use require no CUDA. Requesting unavailable CUDA is an error, with no silent CPU substitution.

Select tiled output-oriented attention in the execution context while keeping learned-model settings unchanged. rust/crates/llm-from-scratch-practical/src/tensor/execution.rs#explicit-attention-execution

    /// Select compact, tiled full attention for output-oriented decoder paths.
    /// Inspectable attention APIs retain their dense reference reports.
    pub fn with_tiled_attention(mut self, tiles: AttentionTiles) -> Self {
        self.attention_tiles = Some(tiles);
        self
    }

    pub const fn attention_tiles(&self) -> Option<AttentionTiles> {
        self.attention_tiles
    }

Count numerical payload at the executed shapes

The storage report runs $B=1$ , $H_q=4$ , $H_{kv}=2$ , $d=d_v=2$ , with query/key tile edges $2/2$ , at three sequence lengths. Each f64 value occupies $8$ bytes. The reference stores $BH_qT^2$ score values and the same number of probability values, one for each batch/head/query/key coordinate; each column below counts one such grid separately.

Sequence cells $T$	Dense score values	Dense probability values	Largest score-tile values
$8$	$256$	$256$	$4$
$16$	$1024$	$1024$	$4$
$32$	$4096$	$4096$	$4$

The maximum score tile has shape $[2,2]$ : $4$ values or $32$ bytes. The five simultaneously counted owned forward buffers are the query tile, key tile, value tile, score/weight tile and tile value-product. At these widths each contains $4$ values, for $20$ values or $160$ bytes altogether. That count describes these five buffers only.

The operation also retains linear-size payloads. Its input count combines $Q$ , $K$ and $V$ ; its output has shape $[B,H_q,T,d_v]$ . Maxima and normalizers are two scalar arrays, each with shape $[B,H_q,T]$ .

Sequence cells $T$	Q/K/V input values	Output values	Maximum and normalizer values together
$8$	$128$	$64$	$64$
$16$	$256$	$128$	$128$
$32$	$512$	$256$	$256$

These are logical numerical payloads, not peak-memory measurements. They exclude allocator, object and index metadata, temporary snapshots and tape alias accounting, primitive-internal or device workspaces, and model-wide activations. Retained-tape diagnostics conservatively count the shared saved snapshot once per operand edge, three times; that is not three physical copies or a process-RAM measurement. The columns above must not be added into a supposedly complete peak total.

Avoiding quadratic intermediate grids does not remove the quadratic number of eligible query-key pairs in full causal attention. No timing, GPU traffic or throughput measurement follows from these payload counts.

Compare derivatives, then the same restored model

The default report uses the nonuniform upstream seed shown in the next excerpt. It compares all $16$ input-gradient coordinates: $4$ query, $4$ key and $8$ value coordinates. Their observed maximum absolute errors are $2.220446049250313\times10^{-15}$ , $5.551115123125783\times10^{-16}$ and $1.3877787807814457\times10^{-17}$ respectively, each within $10^{-10}$ .

Apply the same nonuniform upstream seed and compare every Q/K/V gradient coordinate with dense causal attention. rust/demos/practical-ch19-online-tiled-attention/src/lib.rs#online-attention-gradient-comparison

/// Compare all Q/K/V gradients against the existing dense causal computation.
pub fn compare_gradients(
    execution: &TensorExecution,
    tiles: AttentionTiles,
) -> ArchiveResult<Value> {
    let make = || -> ArchiveResult<[TensorValue; 3]> {
        Ok([
            parameter(&[1, 1, 4, 1], vec![1.; 4])?,
            parameter(&[1, 1, 4, 1], vec![0., 1., 2., 3.])?,
            parameter(&[1, 1, 4, 2], vec![1., 10., 2., 20., 3., 30., 4., 40.])?,
        ])
    };
    let reference = make()?;
    let online = make()?;
    let context = AutogradContext::recording().with_execution(execution);
    let dense = causal_attention_with_ranges(
        context,
        &reference[0].reshape_with_context(context, &[1, 4, 1])?,
        &reference[1].reshape_with_context(context, &[1, 4, 1])?,
        &reference[2].reshape_with_context(context, &[1, 4, 2])?,
        &[4],
        &[0; 4],
    )?;
    let tiled = tiled_causal_attention(
        context,
        &online[0],
        &online[1],
        &online[2],
        AttentionLayout::default(),
        tiles,
        None,
    )?;
    let seed = vec![1., 0.05, -0.2, 0.15, 0.3, -0.01, 1.2, 0.08];
    dense.output().backward_with_seed_and_execution(
        &tensor(&[1, 4, 2], seed.clone())?.view(),
        GraphRetention::Release,
        execution,
    )?;
    tiled.output().backward_with_seed_and_execution(
        &tensor(&[1, 1, 4, 2], seed)?.view(),
        GraphRetention::Release,
        execution,
    )?;
    let mut comparisons = Vec::new();
    for (i, name) in ["query", "key", "value"].iter().enumerate() {
        let actual = online[i]
            .gradient_snapshot()
            .ok_or("missing tiled gradient")?;
        let expected = reference[i]
            .gradient_snapshot()
            .ok_or("missing reference gradient")?;
        let maximum = maximum_error(actual.as_slice(), expected.as_slice())?;
        comparisons.push(json!({"operand":name,"compared_coordinates":actual.len(),"tiled_gradient":actual.as_slice(),
            "reference_gradient":expected.as_slice(),"max_abs_error":maximum,"absolute_tolerance":GRADIENT_TOLERANCE,
            "within_tolerance":maximum<=GRADIENT_TOLERANCE}));
    }
    Ok(json!(comparisons))
}

Separate small-kernel tests finite-difference every $Q,K,V$ coordinate with shared heads, using step $10^{-6}$ and absolute tolerance $4\times10^{-6}$ . They also cover MHA/GQA/MQA, uneven and wider-than-sequence tiles, large and equal scores, fully masked tiles, inactive queries and empty batches, invalid layouts or selectors, and stale/released parents. Nonfinite failure tests check that gradients and results are not partly committed. These are finite fixtures, not universal floating-point equivalence guarantees.

The model comparison reuses Chapter 18’s initialized four-query-head, two-KV-head, window- $2$ model. It saves and restores the existing portable archive, then compares dense and tiled full-prefix logits for all $6$ framed prefix lengths of abcd. Counting every position at each prefix gives $5418$ full-logit comparisons; comparing each final row with the existing incremental cache gives $1548$ cached-logit comparisons. Both maximum errors are $3.885780586188048\times10^{-16}$ , within $10^{-10}$ . No-grad outputs remain graph-free, and the embedding and text head retain one tied table.

Restore one portable text model and compare tiled full-prefix logits with its dense and cached computations. rust/demos/practical-ch19-online-tiled-attention/src/lib.rs#tiled-model-comparison

/// Load the same portable model, compare every full-prefix logit, then the
/// existing incremental cache. Tile selection never changes model identity.
pub fn compare_model(execution: &TensorExecution, tiles: AttentionTiles) -> ArchiveResult<Value> {
    let config = practical_ch18_grouped_query_attention_context::settings(None)?;
    let (model, tokenizer) =
        practical_ch18_grouped_query_attention_context::configured_model(&config)?;
    let archive = PortableModel::from_text(&model, &tokenizer)?;
    let loaded = PortableModel::from_bytes(&archive.to_bytes()?)?;
    let (restored, tokenizer) = loaded.restore_text()?;
    let tokens = tokenizer.encode_utf8_document("abcd");
    let tiled_execution = execution.clone().with_tiled_attention(tiles);
    let mut cache = DecoderKvCache::new(&restored)?;
    let mut session = cache.bind(&restored)?;
    let mut full_error = 0.0_f64;
    let mut cached_error = 0.0_f64;
    let mut compared_full = 0;
    let mut compared_cached = 0;
    let mut graph_free = true;
    let mut empty_reports = true;
    for count in 1..=tokens.len() {
        let dense = restored.forward_with_context(
            AutogradContext::no_grad().with_execution(execution),
            &tokens[..count],
            &[1, count],
        )?;
        let tiled = restored.forward_with_context(
            AutogradContext::no_grad().with_execution(&tiled_execution),
            &tokens[..count],
            &[1, count],
        )?;
        let cached = if count == 1 {
            session.prefill(&tokens[..1])?
        } else {
            session.decode(tokens[count - 1])?
        };
        full_error = full_error.max(maximum_error(
            tiled.logits().value().as_slice(),
            dense.logits().value().as_slice(),
        )?);
        compared_full += tiled.logits().value().len();
        let vocabulary = model.config().vocabulary_size();
        cached_error = cached_error.max(maximum_error(
            cached.logits().value().as_slice(),
            &tiled.logits().value().as_slice()[(count - 1) * vocabulary..count * vocabulary],
        )?);
        compared_cached += vocabulary;
        graph_free &= !tiled.logits().tracks_gradient() && !cached.logits().tracks_gradient();
        empty_reports &= tiled.blocks().is_empty();
    }
    Ok(
        json!({"model_id_before":archive.model_id()?,"model_id_after":loaded.model_id()?,"completed_training_updates":0,
        "query_heads":model.config().heads(),"kv_heads":model.config().kv_heads(),"window":model.config().context_window(),
        "prefix_rows":tokens.len(),"compared_full_logits":compared_full,"compared_cached_logits":compared_cached,
        "max_abs_full_logit_error":full_error,"max_abs_cached_logit_error":cached_error,"absolute_tolerance":OUTPUT_TOLERANCE,
        "within_tolerance":full_error<=OUTPUT_TOLERANCE && cached_error<=OUTPUT_TOLERANCE,
        "all_outputs_graph_free":graph_free,"compact_reports_omit_dense_intermediates":empty_reports,
        "one_tied_table":restored.embedding().table().tensor().is_same_node(restored.text_output_head().table())}),
    )
}

The model ID before and after restoration is 234d9d80f5d63091842fbf7b3cef7ac867aefd20732853a1a463c5b594f3f1f2. Device and tile choices are execution state, so they do not alter the portable learned-model identity. The source-independent backbone still accepts text-derived encoded vectors without requiring a tokenizer; the complete text model retains its tokenizer and tied head. This default example performs zero training updates and establishes no language-quality result.

Continue actual training under the same execution choices

The second demo library test uses a separate seeded training fixture and Chapter 15’s existing training loop through advance_with_execution. It trains a four-query-head, two-KV-head, window- $2$ model with dropout, two-microbatch accumulation, clipping, schedule, loss scaling and block recomputation. After update $1$ it saves a native checkpoint, restores it and advances both branches through update $2$ . Their state_record values agree exactly.

Resume accumulated dropout training with the same tiled execution and reject a mismatched tile-execution label. rust/demos/practical-ch19-online-tiled-attention/src/tests.rs#tiled-accumulated-resume

#[test]
fn existing_accumulated_dropout_training_resumes_with_tiled_recomputation() -> ArchiveResult<()> {
    use llm_from_scratch_practical::checkpoint::resume::TrainingCheckpoint;
    use practical_ch15_save_load_resume_model::{
        advance_with_execution, checkpoint, data_id, executable_envelope, fresh_job, state_record,
    };
    let execution = TensorExecution::cpu().with_tiled_attention(AttentionTiles::new(2, 3)?);
    let mut job = fresh_job()?;
    let config = practical_ch18_grouped_query_attention_context::settings(None)?;
    job.model = llm_from_scratch_practical::models::decoder::DecoderModel::new(
        practical_ch18_grouped_query_attention_context::model_config(
            &config,
            job.tokenizer.layout().vocabulary_size(),
            2,
            Some(2),
        )?,
        &mut llm_from_scratch_practical::nn::init::SplitMix64::from_seed(19),
    )?;
    job.optimizer = llm_from_scratch_practical::training::adamw::AdamW::new(job.optimizer.config());
    advance_with_execution(&mut job, 1, &execution)?;
    let mut envelope = executable_envelope()?;
    envelope.numerical_execution = "same-binary-scalar-cpu-f64-tiled-q2-k3-serial-v1".into();
    let saved = checkpoint(&job, &envelope)?;
    let restored = TrainingCheckpoint::from_bytes(&saved.to_bytes()?)?;
    let mut different_execution = executable_envelope()?;
    different_execution.numerical_execution =
        "same-binary-scalar-cpu-f64-tiled-q3-k2-serial-v1".into();
    assert!(restored.restore(&different_execution, &data_id()?).is_err());
    let mut resumed = restored.restore(&envelope, &data_id()?)?;
    assert_eq!(state_record(&job)?, state_record(&resumed)?);
    advance_with_execution(&mut job, 2, &execution)?;
    advance_with_execution(&mut resumed, 2, &execution)?;
    assert_eq!(state_record(&job)?, state_record(&resumed)?);
    Ok(())
}

This exact continuation binds the same executable and CPU f64 execution label same-binary-scalar-cpu-f64-tiled-q2-k3-serial-v1. The test refuses same-binary-scalar-cpu-f64-tiled-q3-k2-serial-v1. That narrow training envelope is separate from portable inference and from the tolerance-based dense comparisons above. The default report uses tile edges $2/2$ ; this continuation test deliberately uses $2/3$ .

Follow one state into the next

The figure keeps the two summaries attached to the same final query. The second summary includes the first tile’s rescaled contributions; it is not an independently normalized answer to average with the first. Divide each coordinate of the final weighted sum by the same final normalizer.

One query, two tiles, one final normalization

For the final query of the four-row example, the first card summarizes keys zero and one. The second card continues that same state: before adding keys two and three, it multiplies the previous normalizer and both weighted value sums by its reported rescaling factor. Dividing the completed weighted value sums by the final normalizer produces the tiled output. The coordinate cards compare that output with the dense reference for the same query.

Continue the final query across two key tiles

Each card records the running state after its key tile has contributed. Key intervals include their start and exclude their end.

After the first key tile

Key interval: start included, end excluded: $[0,2)$
Running score maximum $m$: $1$
Factor applied to the previous state $\alpha$: $0$
Shifted exponential sum $\ell$: $1.3678794411714423$
Unnormalized weighted value sum, coordinate $0$: $2.3678794411714423$
Unnormalized weighted value sum, coordinate $1$: $23.678794411714424$

After the second key tile

Key interval: start included, end excluded: $[2,4)$
Running score maximum $m$: $3$
Factor applied to the previous state $\alpha$: $0.1353352832366127$
Shifted exponential sum $\ell$: $1.553001792775919$
Unnormalized weighted value sum, coordinate $0$: $5.424095958355416$
Unnormalized weighted value sum, coordinate $1$: $54.240959583554165$

The second card carries the first summary forward: rescale both the previous normalizer and weighted value vector, then add the new tile's contributions. The first card has no previous mass, so its rescaling factor is zero.

Divide the final sums and compare the same output coordinates

Each tiled coordinate is its completed weighted value sum divided by the final normalizer. The dense reference uses all eligible keys in one stable row calculation.

Output coordinate $0$

Tiled output: $3.492652734585769$
Dense reference output: $3.49265273458577$

Output coordinate $1$

Tiled output: $34.9265273458577$
Dense reference output: $34.92652734585769$

Maximum absolute error over all output rows and coordinates: $7.105427357601002\times10^{-15}$
Declared absolute comparison tolerance: $1\times10^{-10}$

Constructed content-only rows from abcd; one query/key feature and two value features. Values come from the CPU f64 report. These comparisons measure numerical agreement for the fixture, not GPU performance or model quality.

Reproduce the report and inspect its checks

Run from the repository root with Docker and the existing course Docker setup available. All configuration, tokenizer and text fixtures are included; no downloaded model or external corpus is needed. The default CPU command is

COURSE_CORPUS=false ./course run cargo run --quiet --offline --locked -p practical-ch19-online-tiled-attention

It prints JSON with worked_example, gradient_comparison, storage_comparisons and model_comparison. It reports Cpu, f64, tile edges $2/2$ and within_tolerance: true for every comparison. The two successive summaries are in worked_example.tile_summaries.

From the same working directory and Docker setup, run the two demo tests:

COURSE_CORPUS=false ./course run cargo test --offline --locked -p practical-ch19-online-tiled-attention --lib

Exactly $2$ tests pass: the report checks and the actual accumulated-training continuation through update $2$ . To run the small attention-kernel checks described above, use the same setup and working directory:

COURSE_CORPUS=false ./course run cargo test --offline --locked -p llm-from-scratch-practical attention::online::tests --lib

Exactly $8$ tests pass, including the exhaustive small-coordinate finite differences and masked/inactive-row cases.

Optional inspection and explanation practice

Use the default command’s tile_summaries to explain why the second tile rescales both earlier sums. Identify the final division and explain why independently normalized tile outputs cannot simply be averaged.
Inspect storage_comparisons from that report. Compare each dense grid with the largest score tile, then identify the input, output and row-statistic payloads excluded by a score-tile-only count.
Inspect all three gradient_comparison entries and run the kernel-test command above. State the coordinate counts, numerical tolerances and the difference between a dense-gradient comparison and finite differences.
Inspect model_comparison and run the two-demo-test command above. Explain the two logit counts, the preserved model identity and the narrower execution condition for exact training continuation.

Checked answers: rescaling, payload counts, derivatives and model continuation

The maximum changes from $1$ to $3$ , so $\exp(-2)$ converts the previous normalizer and weighted value vector to the new score scale. The completed sums are divided only once. Tile outputs normalized separately omit their different relative masses; an unweighted average is not the full attention result.
At $T=8,16,32$ , each dense grid contains $256,1024,4096$ values, while the largest score tile has $4$ values, or $32$ bytes. Five counted temporary forward buffers contain $20$ values, or $160$ bytes. Inputs contain $128,256,512$ values; output and combined row statistics each contain $64,128,256$ . Each value occupies $8$ bytes. These selected payloads exclude overhead and other workspaces, so they establish neither peak process RAM nor speed.
The report compares $4$ query, $4$ key and $8$ value gradient coordinates with dense attention; all errors are within $10^{-10}$ . This checks two analytic paths under the same upstream seed. The separate finite-difference test perturbs every small $Q,K,V$ coordinate, with step $10^{-6}$ and absolute tolerance $4\times10^{-6}$ . The kernel command selects and passes $8$ tests.
The six framed prefixes contribute $5418$ full-logit comparisons, including earlier positions again at later prefixes, and $1548$ last-row cached comparisons. Both maximum errors are $3.885780586188048\times10^{-16}$ within $10^{-10}$ . The portable model ID and one tied table remain unchanged under tile selection. The default is untrained. Of the $2$ passing demo tests, the continuation test actually trains and restores through update $2$ under the same executable and CPU f64 tile- $2/3$ label; it refuses tile- $3/2$ . That exact-resume restriction does not make portable inference tile-dependent.

Keep the model while changing its intermediate work

The existing text encoder, source-independent backbone, tied head and trainer can now evaluate attention with compact shared heads and reconstructed gradients. Running summaries replace complete intermediate score/probability grids in the output-oriented path while the dense reference remains available for inspection. Eligible keys and learned projections keep their existing meaning.

Planned Chapter 20, Paged KV layout, will study logical-to-physical cache pages and partial-page waste while preserving cached logits. That future storage layout is separate from this chapter’s attention tiles; no page allocator, serving scheduler or multi-tenant cache is introduced here.