Share key/value heads and make context explicit

Share retained keys and values while keeping separate queries

Cached generation retains earlier keys and values so the decoder can extend a prefix without recomputing every earlier row. In ordinary multi-head attention, each query head has its own learned key and value projections and its own cached K/V lane. More heads therefore retain more key/value data at the same per-head width. A group of query heads can instead read shared keys and values while keeping its distinct learned queries.

There is also a separate choice about context: an attention row may read its whole permitted causal prefix or only a recent part of that prefix. Sharing K/V heads changes which learned lanes exist. A local window changes which sequence cells those heads may read. We will make both choices explicit in the existing attention and cache.

Follow four query heads into two shared groups

The worked example starts with the four content tokens in abcd, without document boundary tokens. Its constructed embedding rows have width $8$ and first features $1,2,3,4$; their other features are zero. There are $4$ query heads and $2$ shared key/value heads, with $2$ features per head. Query heads $0$ and $1$ read KV head $0$; query heads $2$ and $3$ read KV head $1$. The zero-based map is $[0,0,1,1]$.

For this example only, set the query and key projection weights to zero. Every permitted scaled dot product is then zero, so softmax gives equal weight to every permitted key. The two value-head sequences, in token order, are

The output projection is the identity, so merging the four head results leaves their coordinates unchanged. At the last query cell, $t=3$, full causal attention permits key cells $0,1,2,3$. Each query head assigns weights $(1/4,1/4,1/4,1/4)$. The resulting head outputs are

For example, the first shared value coordinate becomes $(1+2+3+4)/4=2.5$. The Rust report’s merged output is [2.5,0.0,2.5,0.0,0.0,25.0,0.0,25.0].

Now give the same attention layer a window of $2$ cells, including the current cell. At $t=3$, only key cells $2$ and $3$ contribute. Each query head’s weights become $(0,0,1/2,1/2)$, giving

The merged Rust output is now [3.5,0.0,3.5,0.0,0.0,35.0,0.0,35.0]. Both cases share the same K/V groups; the changed permitted keys explain the changed averages. Equal probabilities across query heads come from the deliberately zero Q/K in this fixture. With learned nonzero queries, heads reading the same keys and values can assign different probabilities.

/// Zero scores expose the head map and the exact eligible-key averages.
pub fn worked_example() -> ArchiveResult<Value> {
    let tokenizer = BpeTokenizer::from_merge_pairs(&[])?;
    let tokens = tokenizer.encode_utf8("abcd");
    let mut table = vec![0.0; tokenizer.layout().vocabulary_size() * 8];
    for (position, &token) in tokens.iter().enumerate() {
        table[token as usize * 8] = (position + 1) as f64;
    }
    let embedding =
        Embedding::from_parameter(parameter("example.table", &[table.len() / 8, 8], table)?)?;
    let input = embedding.forward_with_context(AutogradContext::no_grad(), &tokens, &[1, 4])?;
    let mut rows = Vec::new();
    for window in [None, Some(2)] {
        let mut values = vec![0.0; 32];
        values[..4].copy_from_slice(&[1.0, 0.0, 0.0, 10.0]);
        let mut output = vec![0.0; 64];
        for i in 0..8 {
            output[i * 8 + i] = 1.0;
        }
        let attention = MultiHeadAttention::from_parameters_with_attention(
            parameter("example.query.weight", &[8, 8], vec![0.0; 64])?,
            parameter("example.key.weight", &[8, 4], vec![0.0; 32])?,
            parameter("example.value.weight", &[8, 4], values)?,
            parameter("example.output.weight", &[8, 8], output)?,
            4,
            2,
            8,
            100.0,
            window,
        )?;
        let pass = attention.forward_with_context(AutogradContext::no_grad(), &input, 0)?;
        let weights = pass.attention_weights().value();
        rows.push(json!({
            "window":window,
            "last_query_weights":(0..4).map(|h|weights.as_slice()[(h*4+3)*4..(h*4+4)*4].to_vec()).collect::<Vec<_>>(),
            "last_merged_output":pass.output().value().as_slice()[24..32],
            "key_shape":pass.projected_key_heads().shape(),
            "value_heads":pass.projected_value_heads().value().as_slice(),
            "head_mapping":(0..4).map(|h|attention.kv_head_for_query(h).unwrap()).collect::<Vec<_>>()
        }));
    }
    Ok(
        json!({"text":"abcd","query_heads":4,"kv_heads":2,"head_width":2,
        "input_rows":input.value().as_slice(),"cases":rows}),
    )
}

Give each query head an explicit shared KV head

Let $D$ be model width, $H_q$ the positive query-head count and $H_{kv}$ the positive key/value-head count. The per-head feature width is $d=D/H_q$. Require $D$ to be divisible by $H_q$, $d$ to be positive and even for RoPE, and $H_q$ to be divisible by $H_{kv}$. These conditions give whole feature lanes, complete rotation pairs and equally sized query-head groups.

Each group has $r=H_q/H_{kv}$ adjacent query heads. Query head $h$, indexed from zero, reads shared KV head $g(h)=\lfloor h/r\rfloor$. In the worked example, $r=2$, so the first two queries use KV head $0$ and the next two use KV head $1$.

With row-vector matrix conventions, the bias-free query and output matrices $W_q,W_o$ each have shape $[D,D]$. The key and value matrices $W_k,W_v$ each have shape $[D,H_{kv}d]$. For $B$ batch rows and $T$ stored sequence cells per row, splitting the projections gives

Here $B$ counts independent batch rows, $T$ counts cells including any storage padding, and $d$ counts features per head. The worked example’s projected keys have shape $[1,2,4,2]$. RoPE rotates Q and K using their existing position labels. It leaves V unchanged and does not change which head shares which K/V.

For one batch row, write $q_{h,t}$ and $k_{g(h),j}$ for the rotated query and key vectors. Their score is

For the set $\mathcal A_t$ of permitted key cells, normalize this query’s scores only over that set and mix its shared values:

$j$ and $u$ index key cells; $p$ is a dimensionless attention probability; $v$ is the unrotated value vector; $z$ is the resulting head vector. Excluded cells have probability zero. Each query head performs its own normalization. The formula states the ideal normalization; Rust reuses its existing stable softmax on the permitted scores. Concatenating all $H_q$ head results restores $D$ features, which pass through the existing output projection $W_o$.

Ordinary multi-head attention (MHA) is the endpoint $H_{kv}=H_q$, with one KV head per query head. Multi-query attention (MQA) is $H_{kv}=1$, with all queries sharing one KV head. Intermediate sharing is grouped-query attention (GQA). These are settings of this same attention path. Changing $H_{kv}$ changes learned matrix shapes; it is not a way to convert trained weights by editing a setting alone.

Intersect the window with the context already permitted

For batch row $b$, let $n_b$ count active cells and let $s_{b,t}$ be the first key cell in query $t$‘s existing permitted segment. An active query satisfies $t<n_b$. Without a local window, its causal keys satisfy $s_{b,t}\le j\le t$. A positive window length $w$ counts the current cell and at most $w-1$ earlier cells. The head map and window rule are

$\ell_{b,t}$ is the intersected lower bound and $\mathcal A_{b,t}$ is the permitted key set. The segment bound is nonnegative, so it also prevents a negative mathematical lower bound near the start. Rust uses saturating subtraction before the maximum. An active valid segment always includes its self key, keeping its normalization set nonempty. A zero window is refused. Inactive padded queries keep zero attention probabilities and outputs.

The window counts sequence cells, independently of the numerical RoPE labels. For six active packed cells, the report contains

At cell $4$, the segment starts at $2$, but the window starts at $3$; only keys $3$ and $4$ remain. At cell $5$, a new segment starts at $5$, so only self remains. Reset RoPE labels do not join segments or change this count. A window of $1$ permits only self; a sufficiently wide window recovers the existing causal range. “No local window” still obeys the model’s position and cache capacity.

Count learned weights and retained cache data separately

The four bias-free attention matrices contain

learned scalars per layer: the first term counts Q and output weights, and the second counts K and V. With $H_{kv}=H_q$, this becomes $4D^2$. Other model parts, including norms, feed-forward layers and the single tied text table, retain their own parameter counts.

Each layer’s cache stores K and V with physical shape $[B,H_{kv},C,d]$, where $C$ is sequence capacity. For $L$ layers of f64 data, allocation payload in bytes is

The factor $2$ counts K plus V; $8$ is bytes per f64 scalar. If $u$ prefix cells are occupied, replace $C$ by $u$ to count their logical payload. Both counts exclude Rust objects, vectors and allocator overhead. At fixed other axes and element type, the GQA-to-MHA payload ratio is $H_{kv}/H_q$.

The complete-model report uses $B=1$, $L=2$, $C=8$, $D=8$, $H_q=4$, $d=2$ and $u=6$. Its six rows include two document boundary tokens, unlike the four content-only rows in the worked example.

Sharing halves the GQA cache payload relative to MHA here. Adding window $2$ does not change GQA’s allocation or occupied prefix payload: this cache keeps all earlier rows up to its fixed capacity. It does not evict old entries.

The full CPU reference also materializes repeated K/V lanes and dense attention scores. Smaller learned K/V matrices and persistent cache payload therefore do not establish a smaller full-training peak memory footprint or faster attention. These counts describe numerical payload, not complete process RAM, throughput or model quality.

Relate shared heads to local context

Ordinary multi-head attention keeps a separate learned key/value projection and cached lane for every query head; restricting a row’s direct context is a separate choice from sharing those lanes.

In GQA, published at EMNLP 2023, Ainslie and colleagues describe one key head and one value head shared within each query-head group. Their comparison places MQA at one group and MHA at one group per query. For conversion from an existing multi-head checkpoint, they separately describe averaging K/V projections and further pretraining. Our initialized examples perform neither conversion nor that adaptation procedure; changing the head count alone supplies no quality guarantee.

In the Mistral 7B report, version 1 dated 10 October 2023, Jiang and colleagues describe GQA with sliding-window attention and a rolling buffer that overwrites old KV entries. Stacked layers can transmit older information through more recent hidden states. Our window restricts each layer’s direct keys while retaining fixed-capacity all-prefix storage. It does not implement that rolling buffer or the report’s chunked prefill.

Grouped-query attention shares key/value heads within query-head groups, while sliding-window attention restricts the direct keys read by each layer. Keep head sharing, direct-context eligibility and physical cache storage distinct when accounting for an autoregressive decoder. The adjacent map, inclusive-self window and small CPU fixtures here are course choices. They establish none of either paper’s speed, training or quality results.

Extend the existing full and cached attention paths

The checked attention constructor establishes the divisibility conditions. kv_head_for_query then exposes the same adjacent map to the full pass and the example report. An out-of-range query head has no KV head.

/// Adjacent query heads share one learned key/value projection head.
pub fn kv_head_for_query(query_head: usize, query_heads: usize, kv_heads: usize) -> Option<usize> {
    if kv_heads == 0 || query_heads == 0 || !query_heads.is_multiple_of(kv_heads) {
        return None;
    }
    (query_head < query_heads).then(|| query_head / (query_heads / kv_heads))
}

QkvProjections::new_grouped reuses the existing Linear layers. Q keeps width $D$ while K and V use width $H_{kv}d$. A trial random state is committed only after all three projections have been constructed successfully.

    /// The query projection may be wider than the two shared K/V projections.
    pub fn new_grouped(
        parameter_prefix: impl Into<String>,
        model_width: usize,
        query_width: usize,
        kv_width: usize,
        rng: &mut SplitMix64,
    ) -> Result<Self, QkvError> {
        let parameter_prefix = parameter_prefix.into();
        let mut trial = rng.clone();
        let query = Linear::new(
            format!("{parameter_prefix}.query"),
            model_width,
            query_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(QkvProjection::Query))?;
        let key = Linear::new(
            format!("{parameter_prefix}.key"),
            model_width,
            kv_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(QkvProjection::Key))?;
        let value = Linear::new(
            format!("{parameter_prefix}.value"),
            model_width,
            kv_width,
            false,
            &mut trial,
        )
        .map_err(projection_error(QkvProjection::Value))?;
        let projections = Self::from_layers(query, key, value, true)?;
        *rng = trial;
        Ok(projections)
    }

For the full CPU reference, expand_kv_heads uses existing row gather to repeat each group’s K/V rows into the query-head lanes expected by the existing attention kernel. The equal-head case needs no expansion. Reverse-mode gather adds contributions from repeated lanes back to the original shared row. Thus the shared K/V gradient is the sum of the corresponding gradients in a controlled MHA model whose K/V weights are repeated copies.

Tests compare those controlled forward values and derivatives, and separately check nonzero seeded model K/V gradients and selected finite differences. The zero-Q/K teaching example alone cannot demonstrate useful query/key learning.

    /// The CPU full pass materializes repeated views, reusing gather's summed
    /// pullback. Persistent weights and the incremental cache remain unexpanded.
    fn expand_kv_heads(
        &self,
        context: AutogradContext,
        input: &TensorValue,
        stage: MultiHeadStage,
    ) -> Result<TensorValue, MultiHeadAttentionError> {
        if self.kv_heads == self.heads {
            return Ok(input.clone());
        }
        let shape = input.shape();
        let [batch, _, tokens, width] =
            <[usize; 4]>::try_from(shape.as_slice()).expect("split KV heads have rank four");
        let count = batch
            .checked_mul(self.heads)
            .and_then(|n| n.checked_mul(tokens))
            .ok_or(MultiHeadAttentionError::BatchHeadOverflow {
                batch,
                heads: self.heads,
            })?;
        let mut indices = Vec::new();
        indices.try_reserve_exact(count).map_err(|_| {
            MultiHeadAttentionError::CausalAttention(CausalMaskingError::MaskAllocationFailed {
                elements: count,
            })
        })?;
        for row in 0..batch {
            for head in 0..self.heads {
                let shared = self.kv_head_for_query(head).expect("query head in range");
                for token in 0..tokens {
                    indices.push((row * self.kv_heads + shared) * tokens + token);
                }
            }
        }
        input
            .reshape_with_context(context, &[batch * self.kv_heads * tokens, width])
            .and_then(|table| {
                table.gather_rows_with_context(context, &indices, &[batch, self.heads, tokens])
            })
            .map_err(multi_autodiff(stage))
    }

The window helper intersects the existing segment with a recent suffix. Saturating subtraction preserves the start of short prefixes, while the maximum preserves a later packed-segment boundary.

/// A window counts the self key and at most `window - 1` preceding cells.
/// Intersect it with the existing segment boundary; neither restriction is lost.
pub fn causal_window_start(query: usize, segment_start: usize, window: usize) -> Option<usize> {
    if window == 0 || segment_start > query {
        return None;
    }
    Some(segment_start.max(query.saturating_sub(window - 1)))
}

The full pass applies that bound to each active query in the existing head lanes. Its prior active lengths and segment starts keep padding and unrelated packed examples outside the probability normalization.

        let mut lane_starts = lane_starts;
        if let Some(window) = self.context_window {
            if lane_starts.is_none() {
                lane_starts = Some(repeat_row_metadata(&vec![0; cells], shape[1], self.heads)?);
            }
            let starts = lane_starts.as_mut().expect("window ranges");
            for (lane, &length) in lane_lengths
                .as_ref()
                .expect("window lengths")
                .iter()
                .enumerate()
            {
                for query in 0..length {
                    let index = lane * shape[1] + query;
                    starts[index] = causal_window_start(query, starts[index], window)
                        .expect("validated segment start and positive window");
                }
            }
        }

Incremental inference retains one rotated K and one unrotated V per KV head. It computes $H_q$ queries, reads each query’s shared storage lane, and normalizes scores only over the eligible suffix. The current candidate key is included before the existing all-layer cache transaction commits. Excluded probability cells remain zero. Failed operations, stale model revisions or incompatible attention settings cannot partly advance this cache.

    let start = cache.context_window().map_or(0, |window| {
        causal_window_start(cache.len(), 0, window).expect("positive bound window")
    });
    let eligible = prefix - start;
    let score_elements =
        lanes
            .checked_mul(eligible)
            .ok_or(IncrementalAttentionError::BufferSizeOverflow {
                stage: IncrementalAttentionStage::Scores,
            })?;
    let mut scores = reserved_buffer(score_elements, IncrementalAttentionStage::Scores)?;
    let scale = 1.0 / (cache.head_width() as f64).sqrt();

    for batch in 0..cache.batch_size() {
        for head in 0..cache.query_heads() {
            let lane = batch * cache.query_heads() + head;
            let shared = head / (cache.query_heads() / cache.heads());
            let kv_lane = batch * cache.heads() + shared;
            let query_start = lane * cache.head_width();
            for position in start..prefix {
                let key_start = if position == cache.len() {
                    kv_lane * cache.head_width()
                } else {
                    cache.row_start(kv_lane, position)
                };
                let key_values = if position == cache.len() {
                    candidate_key.as_slice()
                } else {
                    cache.key_storage()
                };
                let mut dot = 0.0;
                for feature in 0..cache.head_width() {
                    dot +=
                        query.as_slice()[query_start + feature] * key_values[key_start + feature];
                }
                scores[lane * eligible + position - start] = dot * scale;
            }
        }
    }

    let score_tensor = Tensor::from_vec(
        vec![cache.batch_size(), cache.query_heads(), 1, eligible],
        scores,
    )
    .map_err(|source| IncrementalAttentionError::Tensor {
        stage: IncrementalAttentionStage::Scores,
        source,
    })?;
    let weights =
        softmax(&score_tensor.view(), 3).map_err(IncrementalAttentionError::Probability)?;

The window leaves original cached RoPE positions and capacity unchanged. Generation still refuses to exceed the existing position/cache bound. In deeper layers, recent hidden states may already carry information from older cells. Re-encoding only the last raw window would discard those states and is a different computation from a full-prefix pass with the same per-layer mask.

Compare each restored model with its own complete-prefix reference

The default report initializes MHA, GQA, MQA and window-two GQA models for their respective shapes. Each is saved and reloaded before comparison. The document abcd becomes $6$ framed rows with a vocabulary of $258$ tokens. At every prefix length, the reference recomputes the entire prefix and the cached path reuses earlier K/V. Both use the same restored weights, positions and window, with dropout inactive and gradient recording disabled.

/// Each variant compares its own saved/reloaded weights and all vocabulary logits.
pub fn compare_model(
    config: &Config,
    kv_heads: usize,
    window: Option<usize>,
) -> ArchiveResult<Value> {
    let tokenizer = BpeTokenizer::from_merge_pairs(&[])?;
    let model = DecoderModel::new(
        model_config(
            config,
            tokenizer.layout().vocabulary_size(),
            kv_heads,
            window,
        )?,
        &mut SplitMix64::from_seed(u64::try_from(number(config, "model.seed")?)?),
    )?;
    let saved = PortableModel::from_text(&model, &tokenizer)?;
    let loaded = PortableModel::from_bytes(&saved.to_bytes()?)?;
    let (restored, _) = loaded.restore_text()?;
    let document = config
        .get("text.document")
        .ok_or("missing text document")?
        .as_string()?;
    let tokens = tokenizer.encode_utf8_document(document);
    let mut cache = DecoderKvCache::new(&restored)?;
    let mut maximum = 0.0_f64;
    let mut compared = 0;
    let mut all_graph_free = true;
    {
        let mut session = cache.bind(&restored)?;
        for count in 1..=tokens.len() {
            let cached = if count == 1 {
                session.prefill(&tokens[..1])?
            } else {
                session.decode(tokens[count - 1])?
            };
            let full = restored.forward_with_context(
                AutogradContext::no_grad(),
                &tokens[..count],
                &[1, count],
            )?;
            let width = restored.config().vocabulary_size();
            for (&actual, &reference) in cached
                .logits()
                .value()
                .as_slice()
                .iter()
                .zip(&full.logits().value().as_slice()[(count - 1) * width..count * width])
            {
                maximum = maximum.max((actual - reference).abs());
                compared += 1;
            }
            all_graph_free &=
                !cached.logits().tracks_gradient() && !full.logits().tracks_gradient();
        }
    }
    let mut allocated = 0;
    let mut logical = 0;
    for layer in 0..cache.layer_count() {
        let layer = cache.layer_cache(layer).ok_or("missing layer cache")?;
        allocated += (layer.key_storage().len() + layer.value_storage().len()) * 8;
        logical += (layer.keys()?.len() + layer.values()?.len()) * 8;
    }
    let cfg = restored.config();
    Ok(
        json!({"query_heads":cfg.heads(),"kv_heads":cfg.kv_heads(),"window":cfg.context_window(),
        "head_width":cfg.head_width(),"model_width":cfg.model_width(),"layers":cfg.layers(),
        "capacity":cfg.max_positions(),"prefix_rows":tokens.len(),"compared_logits":compared,
        "max_abs_logit_error":maximum,"absolute_tolerance":LOGIT_TOLERANCE,
        "within_tolerance":maximum<=LOGIT_TOLERANCE,"all_outputs_graph_free":all_graph_free,
        "allocated_kv_payload_bytes":allocated,"logical_kv_payload_bytes":logical,
        "attention_scalars_per_layer":restored.blocks()[0].attention().parameter_count(),
        "model_scalars":restored.parameter_count(),"model_id":loaded.model_id()?,
        "one_tied_table":restored.text_output_head().table().is_same_node(restored.embedding().table().tensor()),
        "head_mapping":(0..cfg.heads()).map(|h|restored.blocks()[0].attention().kv_head_for_query(h).unwrap()).collect::<Vec<_>>() }),
    )
}

Each variant compares $6\times258=1548$ vocabulary coordinates and reports maximum absolute logit error 0.0, within absolute tolerance $10^{-10}$ in CPU f64. All outputs are graph-free and the text embedding/head use one tied table. This is an observed fixture result. Ideal full/cached equivalence does not promise universal floating-point bitwise equality or cross-backend equality.

The default performs zero training updates. Separately initialized variants are different models; comparing each model with itself does not compare their language quality or assert cross-model equality.

Save the attention semantics with the model

KV-head count and the optional window belong to semantic model settings. The existing native settings and PortableModel preserve them through text model restoration and tokenizer-independent backbone restoration. The text model still restores one shared embedding/head table. Equal-head/no-window defaults preserve earlier archive identities by omitting the optional fields.

Native exact training checkpoints also retain these settings. The older legacy checkpoint encoding refuses unsupported nondefault attention settings instead of silently dropping them. These are extensions of the existing composition and persistence paths.

The example’s thin save command initializes the grouped-window model and saves it for the existing Chapter 15 inference command:

/// Save logical attention settings and weights for the existing inference CLI.
pub fn save_model(config: &Config, path: &Path) -> ArchiveResult<Value> {
    let (model, tokenizer) = configured_model(config)?;
    let archive = PortableModel::from_text(&model, &tokenizer)?;
    archive.save_atomic(path)?;
    Ok(
        json!({"model_saved":true,"model_id":archive.model_id()?,"initialized_updates":0,
        "query_heads":model.config().heads(),"kv_heads":model.config().kv_heads(),
        "window":model.config().context_window(),"execution":"cpu/f64"}),
    )
}

Follow shared values through two permitted contexts

Reproduce the report and use the saved model

From the repository root, with Docker and the existing course Docker setup available, run the default CPU example. Its configuration, tokenizer and text fixtures are included; no downloaded model or external corpus is required.

COURSE_CORPUS=false ./course run cargo run --offline --locked --quiet -p practical-ch18-grouped-query-attention-context

The JSON contains worked_example, packed_window_example, model_comparisons and the initialized CPU f64 scope. It prints the head map, two observed mixtures, packed starts, all four storage rows and four same-model full/cached comparisons. Every comparison reports compared_logits: 1548, max_abs_logit_error: 0.0 and within_tolerance: true.

To save a file and load it in another Rust program, run this entire block from the same repository-root working directory with the same Docker setup. One Docker shell retains the file while each Rust program exits before the next starts.

COURSE_CORPUS=false ./course run sh -c '
set -eu
mkdir -p tmp/ch18-attention-flow
cargo run --quiet --offline --locked -p practical-ch18-grouped-query-attention-context -- save --model tmp/ch18-attention-flow/grouped-window.safetensors
cargo run --quiet --offline --locked -p practical-ch15-save-load-resume-model -- infer --model tmp/ch18-attention-flow/grouped-window.safetensors --prompt ab --tokens 3
'

The save command reports initialized_updates: 0, query_heads: 4, kv_heads: 2 and window: 2. Both programs report the same model_id: 234d9d80f5d63091842fbf7b3cef7ac867aefd20732853a1a463c5b594f3f1f2. Chapter 15’s inference command loads the saved settings and weights and reports device: "cpu", precision: "f64", tied_alias: true, $3256$ parameter scalars and $20$ named parameter tensors. Generated IDs [91,21,170] map to bytes [89,19,168]. These initialized-model outputs demonstrate loading and generation; they establish neither useful prose nor UTF-8 validity. Separate disposable ./course run invocations would lose the saved file.

With the same repository-root working directory and Docker setup, run the example’s two library tests:

COURSE_CORPUS=false ./course run cargo test --offline --locked -p practical-ch18-grouped-query-attention-context --lib

Both tests pass. The first checks the explained mixtures and full/cached comparisons. The second performs actual accumulated training with shared heads and window $2$, saves after one completed update, restores the existing native checkpoint and checks agreement after both jobs reach update $2$. That training test is separate from the zero-update default and saved inference example.

Optional inspection and explanation practice

1. Use the default command’s worked_example to explain map $[0,0,1,1]$ and the two final-row mixtures. State why the four query heads have equal probabilities here and why sharing alone does not require that equality.
2. Inspect packed_window_example in the same report. Explain the effective starts at cells $4$ and $5$, and distinguish sequence cells from the displayed RoPE labels.
3. Compare all four model_comparisons storage rows from the default command. Explain GQA’s allocated-byte ratio to MHA and why adding window $2$ leaves the GQA byte counts unchanged.
4. Inspect each model’s compared_logits, error and tolerance in that report. Explain what is compared at each prefix and why retaining only the latest raw window would change the reference computation.
5. Run the complete save/infer block and the two-library-test command above. Check the shared saved identity, CPU precision and tied table; distinguish the initialized inference example from the native training-continuation test.

Keep attention semantics with the same model

The text encoder, source-independent backbone, tied head and CPU inference path now preserve shared K/V heads and a local context window. Separate query calculations read compact shared cache lanes, while the window intersects the keys each layer may use. The same-weight full/cached comparison checks those choices together within the existing position and cache capacity.

The full reference still materializes repeated K/V lanes and dense scores. Planned Chapter 19, Online and tiled attention, will study running normalization and tiles for those attention intermediates. It is a separate planned extension from the sharing and direct-context masks implemented here.