30 · Content revision 2

Run causal attention in several heads, then mix

Learn how full-width query, key, and value projections become separate rotary causal attention heads before concatenation and one learned output projection.

Predict the boundaries before calculating probabilities

Use one batch with three token rows, model width $d_{\mathrm{model}}=4$ , and head count $h=2$ . The per-head width is

d_h=\frac{d_{\mathrm{model}}}{h}=2.

This chapter applies RoPE to every complete head, so $d_h$ must also be even. That is why $h=4$ is rejected here even though it divides a model width of $4$ : it would leave one unpaired coordinate per head.

The frozen input is

X= \begin{bmatrix} 1&0&1&0\\ \cos(1)&-\sin(1)&0&1\\ \cos(2)&-\sin(2)&1&1 \end{bmatrix}.

For this fixture only, $W_Q=W_K=W_V=I$ . That makes the projected head blocks easy to inspect. It is not a restriction of multi-head attention: in a trained layer, the columns assigned to any head can read every input feature.

Before looking at the Rust output, predict this shape sequence:

\begin{aligned} X:[1,3,4] &\to(Q,K,V):[1,3,4] \to\operatorname{split}(Q,K,V):[1,2,3,2],\\ \operatorname{split}(Q,K,V):[1,2,3,2] &\to A:[1,2,3,3] \to H:[1,2,3,2],\\ H:[1,2,3,2] &\to\operatorname{Concat}(H):[1,3,4] \xrightarrow{W_O}\operatorname{MHA}(X):[1,3,4]. \end{aligned}

The first two input coordinates deliberately counter-rotate position. At token position $t$ , they equal $R(-t)[1,0]$ , so Chapter 29’s rotation returns

R(t)R(-t)[1,0]=[1,0].

Head $0$ therefore compares three identical query/key rows. Predict its causal probabilities before running anything:

A^{(0)}= \begin{bmatrix} 1&0&0\\ \tfrac12&\tfrac12&0\\ \tfrac13&\tfrac13&\tfrac13 \end{bmatrix}.

Head $1$ receives $[1,0]$ , $[0,1]$ , and $[1,1]$ before rotation, so its allowed scores are not equal. You can still predict that both heads have exactly zero probability above the diagonal.

Now carry token position $1$ through both heads. Head $0$ averages its first two unrotated value rows, not its rotated query rows:

H^{(0)}_1 =\tfrac12[1,0]+\tfrac12[0.540302,-0.841471] =[0.770151,-0.420735].

For head $1$ , the two visible scaled scores and their softmax are

[s^{(1)}_{1,0},s^{(1)}_{1,1}] =[-0.595010,0.707107] \xrightarrow{\operatorname{softmax}} [0.213809,0.786191].

Its unrotated value rows are $[1,0]$ and $[0,1]$ , so $H^{(1)}_1=[0.213809,0.786191]$ . Concatenating the two head rows gives $[0.770151,-0.420735,0.213809,0.786191]$ .

Finally, predict where features from the two heads can first influence the same learned output coordinate. Splitting and both attention operations keep the lanes separate; concatenation merely places their coordinates side by side. After the separate attention calculations, only multiplication by $W_O$ permits learned cross-head mixing. The frozen $W_O$ swaps the blocks, which demonstrates learned routing; exercise 5 adds an off-block entry to demonstrate an actual sum across heads. The earlier dense query/key/value projections may already use every coordinate of the shared input row.

Project into head views before scoring

The implementation stores three packed, bias-free matrices:

Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V,

with $Q$ , $K$ , and $V$ shaped $[B,T,d_{\mathrm{model}}]$ . A reshape and transpose then produce $[B,h,T,d_h]$ .

The packed form is just an efficient organization of the paper’s per-head columns. For queries,

W_Q=[W_1^Q\;W_2^Q\;\cdots\;W_h^Q], \qquad Q_i=XW_i^Q,

and the same relationship holds for keys and values. This order matters: splitting raw $X$ first would confine a head to a fixed input slice. Projecting first lets every $W_i^Q$ , $W_i^K$ , and $W_i^V$ read the whole input row before their outputs enter separate head lanes.

For head $i$ , apply RoPE only to its query and key rows:

\widetilde Q_i=\operatorname{RoPE}(Q_i),\qquad \widetilde K_i=\operatorname{RoPE}(K_i).

Then reuse the causal attention operation from Chapter 28:

A_i=\operatorname{softmax}_{j}\!\left( \frac{\widetilde Q_i\widetilde K_i^\top}{\sqrt{d_h}}+M \right), \qquad H_i=A_iV_i.

$M$ blocks future keys, and the softmax runs separately over each head’s key positions. A single full-width attention calculation would produce one normalized $[T,T]$ table shared by every value feature. Multi-head attention instead produces $h$ separately normalized $[T,T]$ tables, which may differ because their projected queries and keys may differ.

Each $A_i$ has shape $[B,T,T]$ . Each $H_i$ has shape $[B,T,d_h]$ . Restore the model width by concatenating on the final feature axis, then apply one learned output matrix:

\operatorname{MHA}(X)=\operatorname{Concat}(H_1,\ldots,H_h)W_O

where

W_O\in\mathbb{R}^{d_{\mathrm{model}}\times d_{\mathrm{model}}}.

Concatenation has no parameters and performs no averaging. $W_O$ can preserve, swap, add, or otherwise recombine coordinates produced by different heads. The target implementation has four bias-free model-width matrices, so

N_\theta=4d_{\mathrm{model}}^2.

For $d_{\mathrm{model}}=4$ , that is $64$ scalar parameters.

Keep the model, head, and token axes distinct

$\operatorname{MHA}$ is the complete multi-head causal self-attention layer.
$X$ is its batch of input hidden-state rows.
$\operatorname{Concat}$ joins the head outputs along their feature axis at each token position; it does not normalize or mix them.
$B$ is the batch size.
$T$ is the number of token positions in each sequence.
$d_{\mathrm{model}}$ is the input and output width of the complete layer.
$h$ is the number of independently normalized attention heads.
$d_h=d_{\mathrm{model}}/h$ is the width of one head; it is even here because RoPE rotates adjacent coordinate pairs.
$t$ is a zero-based token position in the worked example, and $R(t)$ is the two-dimensional rotation by $t$ radians used by its first head.
$W_Q$ , $W_K$ , and $W_V$ are packed model-width projection matrices. The columns belonging to head $i$ form $W_i^Q$ , $W_i^K$ , and $W_i^V$ .
$Q_i$ , $K_i$ , and $V_i$ are head $i$ ‘s projected query, key, and value rows.
$\widetilde Q_i$ and $\widetilde K_i$ are those query and key rows after RoPE; values are not rotated.
$M$ is the causal mask that blocks keys to the right of a query position.
$j$ indexes key positions and is the axis normalized by each row softmax.
$A_i$ is head $i$ ‘s row-normalized causal attention table.
$H_i=A_iV_i$ is head $i$ ‘s weighted value mixture.
$W_O$ is the learned output matrix applied after all $H_i$ rows have been concatenated.
$N_\theta$ is the number of learned scalar parameters. $\delta$ is the central-difference perturbation, $\varepsilon_g$ is its accepted error, and $\varepsilon$ is the tolerance for the assembly and invariance checks.

The equations follow the paper-style convention $i\in\{1,\ldots,h\}$ . Rust and the frozen trace use zero-based head IDs $0,\ldots,h-1$ , so trace heads $0$ and $1$ correspond to $H_1$ and $H_2$ in this two-head example.

The head axis and token axis serve different roles. Each head gets its own probability distribution over token positions, while concatenation joins head features at the same token position. No probability normalization occurs across heads.

From one recurrent alignment to parallel projected attention

In their neural machine translation paper, Bahdanau et al. replace one fixed source-sentence vector with a target-dependent context vector, but the recurrent decoder still forms one alignment distribution and context at a time rather than several parallel projected self-attention heads. Each context is a weighted sum of encoder annotations whose scores depend on the previous decoder state. This is an important learned-attention predecessor, but it is not Transformer self-attention: the paper does not define queries, keys, values, causal masks, or several parallel heads.

One scaled dot-product head, as built in Chapter 27, gives each query row one learned comparison space, one probability distribution over key positions, and one value mixture. In Attention Is All You Need, Vaswani et al. project queries, keys, and values several times, perform the projected attention functions in parallel, concatenate their outputs, and project the concatenation. The authors motivate this as an opportunity to attend jointly to information from different representation subspaces and positions.

That motivation is not a specialization guarantee. Heads share $X$ , are trained together, and are coupled again by $W_O$ . A particular trained model may show repeated, mixed, or difficult-to-interpret head behavior. This fixture proves only that one valid parameter setting produces two different causal tables and keeps their lanes separate through concatenation and until $W_O$ .

The LLaMA paper provides a concrete modern causal language-model example. Touvron et al. document a Transformer-based model with multiple attention heads, RoPE at every layer, and an optimized causal multi-head attention implementation. That optimized implementation avoids storing attention weights. The dense calculation below keeps those weights visible for inspection; it does not reproduce that optimized kernel.

The executable comparison does not pretend to reproduce an entire recurrent decoder. It computes the defining weighted source-context sum for one fixed alignment distribution, runs the existing single-head causal operation on the same three-token input, and compares its one $[1,3,3]$ probability tensor with the multi-head fixture’s $[1,2,3,3]$ tensor. Both have normalized rows, but the latter contains two independently normalized tables. The comparison does not attribute the fixture’s packed-matrix layout, even-width rule, bias policy, names, or values to the papers:

Compute one weighted source context, one full-width causal table, and two head-local causal tables rust/demos/ch30-multi-head-attention/src/lib.rs#historical-multi-head-contrast

#[derive(Clone, Debug, PartialEq)]
pub struct HistoryEvidence {
    pub earlier_weighted_context: [f64; 2],
    pub earlier_distributions_per_target: usize,
    pub single_head_weight_shape: Vec<usize>,
    pub multi_head_weight_shape: Vec<usize>,
    pub single_head_tables: usize,
    pub multi_head_tables: usize,
    pub all_rows_normalized: bool,
    pub mixing_stage: &'static str,
    pub modern_example: &'static str,
    pub weight_api: &'static str,
}

fn weighted_source_context(weights: [f64; 2], annotations: [[f64; 2]; 2]) -> [f64; 2] {
    [
        weights[0] * annotations[0][0] + weights[1] * annotations[1][0],
        weights[0] * annotations[0][1] + weights[1] * annotations[1][1],
    ]
}

fn rows_are_normalized(weights: &Tensor, tokens: usize) -> bool {
    weights
        .as_slice()
        .chunks_exact(tokens)
        .all(|row| (row.iter().sum::<f64>() - 1.0).abs() <= INVARIANT_TOLERANCE)
}

pub fn historical_attention_contrast(
    primary: &PrimaryEvidence,
) -> Result<HistoryEvidence, FixtureError> {
    let earlier_weighted_context = weighted_source_context([0.25, 0.75], [[1.0, 0.0], [0.0, 1.0]]);

    let full_width = constant(&[BATCH, TOKENS, MODEL_WIDTH], primary.input.as_slice())?;
    let single_head =
        causal_scaled_dot_product_self_attention(&full_width, &full_width, &full_width)?;
    let single_head_weights = single_head.weights().value();
    let single_head_weight_shape = single_head_weights.shape().to_vec();
    let multi_head_weight_shape = primary.attention_weights.shape().to_vec();
    let single_head_tables = single_head_weights.len() / (TOKENS * TOKENS);
    let multi_head_tables = primary.attention_weights.len() / (TOKENS * TOKENS);

    Ok(HistoryEvidence {
        earlier_weighted_context,
        earlier_distributions_per_target: 1,
        single_head_weight_shape,
        multi_head_weight_shape,
        single_head_tables,
        multi_head_tables,
        all_rows_normalized: rows_are_normalized(&single_head_weights, TOKENS)
            && rows_are_normalized(&primary.attention_weights, TOKENS),
        mixing_stage: "after-concatenation",
        modern_example: "llama-causal-heads-plus-rope",
        weight_api: "dense-teaching-evidence",
    })
}

Keep the complete Rust path differentiable and inspectable

The public split_heads helper uses a reshape followed by a transpose:

[B,T,d_{\mathrm{model}}] \to[B,T,h,d_h] \to[B,h,T,d_h].

merge_heads applies the inverse transpose and reshape. Both operations remain on the autodiff tape; there is no second hand-written copying algorithm.

Split and merge the head axis with taped reshape and transpose operations rust/crates/llm-from-scratch/src/attention/multi_head.rs#head-layout

/// A taped reshape/transpose stage in the public split and merge helpers.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum HeadLayoutStage {
    SplitReshape,
    SplitTranspose,
    MergeTranspose,
    MergeReshape,
}

impl fmt::Display for HeadLayoutStage {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::SplitReshape => "split reshape",
            Self::SplitTranspose => "split transpose",
            Self::MergeTranspose => "merge transpose",
            Self::MergeReshape => "merge reshape",
        })
    }
}

/// A rejected rank or feature partition in `split_heads` or `merge_heads`.
#[derive(Clone, Debug, PartialEq)]
pub enum HeadLayoutError {
    SplitRank {
        rank: usize,
    },
    MergeRank {
        rank: usize,
    },
    ZeroHeadCount,
    ZeroHeadWidth,
    WidthNotDivisible {
        width: usize,
        heads: usize,
    },
    ModelWidthOverflow {
        heads: usize,
        head_width: usize,
    },
    Autodiff {
        stage: HeadLayoutStage,
        source: TensorAutodiffError,
    },
}

impl fmt::Display for HeadLayoutError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::SplitRank { rank } => write!(
                formatter,
                "head split input must have rank three [batch, tokens, model_width], got rank {rank}"
            ),
            Self::MergeRank { rank } => write!(
                formatter,
                "head merge input must have rank four [batch, heads, tokens, head_width], got rank {rank}"
            ),
            Self::ZeroHeadCount => formatter.write_str("head count must be nonzero"),
            Self::ZeroHeadWidth => formatter.write_str("head width must be nonzero"),
            Self::WidthNotDivisible { width, heads } => write!(
                formatter,
                "model width {width} must be divisible by head count {heads}"
            ),
            Self::ModelWidthOverflow { heads, head_width } => write!(
                formatter,
                "merged model width overflows for {heads} heads of width {head_width}"
            ),
            Self::Autodiff { stage, source } => {
                write!(formatter, "multi-head {stage}: {source}")
            }
        }
    }
}

impl Error for HeadLayoutError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Autodiff { source, .. } => Some(source),
            _ => None,
        }
    }
}

fn layout_autodiff(stage: HeadLayoutStage) -> impl FnOnce(TensorAutodiffError) -> HeadLayoutError {
    move |source| HeadLayoutError::Autodiff { stage, source }
}

/// Converts `[batch, tokens, model_width]` to `[batch, heads, tokens, head_width]`.
pub fn split_heads(input: &TensorValue, heads: usize) -> Result<TensorValue, HeadLayoutError> {
    let shape = input.shape();
    if shape.len() != 3 {
        return Err(HeadLayoutError::SplitRank { rank: shape.len() });
    }
    if heads == 0 {
        return Err(HeadLayoutError::ZeroHeadCount);
    }
    let width = shape[2];
    if width == 0 {
        return Err(HeadLayoutError::ZeroHeadWidth);
    }
    if !width.is_multiple_of(heads) {
        return Err(HeadLayoutError::WidthNotDivisible { width, heads });
    }
    let head_width = width / heads;
    let reshaped = input
        .reshape(&[shape[0], shape[1], heads, head_width])
        .map_err(layout_autodiff(HeadLayoutStage::SplitReshape))?;
    reshaped
        .transpose(1, 2)
        .map_err(layout_autodiff(HeadLayoutStage::SplitTranspose))
}

/// Converts `[batch, heads, tokens, head_width]` back to model-width rows.
pub fn merge_heads(input: &TensorValue) -> Result<TensorValue, HeadLayoutError> {
    let shape = input.shape();
    if shape.len() != 4 {
        return Err(HeadLayoutError::MergeRank { rank: shape.len() });
    }
    if shape[1] == 0 {
        return Err(HeadLayoutError::ZeroHeadCount);
    }
    if shape[3] == 0 {
        return Err(HeadLayoutError::ZeroHeadWidth);
    }
    let model_width =
        shape[1]
            .checked_mul(shape[3])
            .ok_or(HeadLayoutError::ModelWidthOverflow {
                heads: shape[1],
                head_width: shape[3],
            })?;
    let transposed = input
        .transpose(1, 2)
        .map_err(layout_autodiff(HeadLayoutStage::MergeTranspose))?;
    transposed
        .reshape(&[shape[0], shape[2], model_width])
        .map_err(layout_autodiff(HeadLayoutStage::MergeReshape))
}

MultiHeadAttention::new checks nonzero model width, nonzero head count, exact divisibility, even head width, position capacity, and RoPE base before committing random state. It initializes query, key, value, and output matrices in stable order through a trial generator. MultiHeadAttention::from_parameters gives the fixture and future checkpoint loader the same checked layer with exact matrices.

Input validation follows an intentional order: rank, final width, nonempty token axis, then the complete absolute position interval. Position bounds are checked before any tape-backed operation, so an invalid interval is reported consistently even when the input has also been released. Delegated layout, projection, RoPE, and causal errors retain the stage that failed.

Keep configuration, input, position, and delegated stage failures typed rust/crates/llm-from-scratch/src/attention/multi_head.rs#multi-head-errors

/// A Q/K/V branch whose head-layout or rotary operation failed.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MultiHeadInput {
    Query,
    Key,
    Value,
}

impl fmt::Display for MultiHeadInput {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::Query => "query",
            Self::Key => "key",
            Self::Value => "value",
        })
    }
}

/// A cumulative tensor stage owned by the multi-head assembly.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MultiHeadStage {
    QueryLanes,
    KeyLanes,
    ValueLanes,
    RestoreWeights,
    RestoreHeadOutputs,
}

impl fmt::Display for MultiHeadStage {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::QueryLanes => "query lane flatten",
            Self::KeyLanes => "key lane flatten",
            Self::ValueLanes => "value lane flatten",
            Self::RestoreWeights => "attention-weight head restore",
            Self::RestoreHeadOutputs => "head-output restore",
        })
    }
}

/// A rejected configuration, parameter set, input, or cumulative forward stage.
#[derive(Clone, Debug, PartialEq)]
pub enum MultiHeadAttentionError {
    ZeroModelWidth,
    ZeroHeadCount,
    ModelWidthNotDivisible {
        model_width: usize,
        heads: usize,
    },
    OddHeadWidth {
        head_width: usize,
    },
    QkvProjection(QkvError),
    QkvOutputWidthMismatch {
        model_width: usize,
        projected_width: usize,
    },
    OutputProjection(LinearError),
    OutputInputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    OutputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    Initialization(InitializationError),
    InputRank {
        rank: usize,
    },
    InputWidthMismatch {
        expected: usize,
        actual: usize,
    },
    EmptyTokens,
    PositionOffsetOverflow {
        offset: usize,
        tokens: usize,
    },
    PositionRangeExceeded {
        offset: usize,
        tokens: usize,
        max_positions: usize,
    },
    BatchHeadOverflow {
        batch: usize,
        heads: usize,
    },
    HeadLayout {
        input: MultiHeadInput,
        source: HeadLayoutError,
    },
    MergeLayout(HeadLayoutError),
    RotaryConfiguration(RopeError),
    Rotary {
        input: MultiHeadInput,
        source: RopeError,
    },
    CausalAttention(CausalMaskingError),
    Autodiff {
        stage: MultiHeadStage,
        source: TensorAutodiffError,
    },
}

impl fmt::Display for MultiHeadAttentionError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::ZeroModelWidth => formatter.write_str("multi-head model width must be nonzero"),
            Self::ZeroHeadCount => formatter.write_str("multi-head head count must be nonzero"),
            Self::ModelWidthNotDivisible { model_width, heads } => write!(
                formatter,
                "multi-head model width {model_width} must be divisible by head count {heads}"
            ),
            Self::OddHeadWidth { head_width } => write!(
                formatter,
                "multi-head per-head width must be even for RoPE, got {head_width}"
            ),
            Self::QkvProjection(source) => source.fmt(formatter),
            Self::QkvOutputWidthMismatch {
                model_width,
                projected_width,
            } => write!(
                formatter,
                "Q/K/V projected width must equal model width {model_width}, got {projected_width}"
            ),
            Self::OutputProjection(source) => write!(formatter, "output projection: {source}"),
            Self::OutputInputWidthMismatch { expected, actual } => write!(
                formatter,
                "output projection input width must equal model width {expected}, got {actual}"
            ),
            Self::OutputWidthMismatch { expected, actual } => write!(
                formatter,
                "output projection width must equal model width {expected}, got {actual}"
            ),
            Self::Initialization(source) => source.fmt(formatter),
            Self::InputRank { rank } => write!(
                formatter,
                "multi-head input must have rank three [batch, tokens, model_width], got rank {rank}"
            ),
            Self::InputWidthMismatch { expected, actual } => write!(
                formatter,
                "multi-head input final width must equal model width {expected}, got {actual}"
            ),
            Self::EmptyTokens => formatter.write_str(
                "multi-head attention needs at least one token so every causal row has a key",
            ),
            Self::PositionOffsetOverflow { offset, tokens } => write!(
                formatter,
                "multi-head position interval overflows: offset {offset} plus {tokens} tokens"
            ),
            Self::PositionRangeExceeded {
                offset,
                tokens,
                max_positions,
            } => write!(
                formatter,
                "multi-head position interval [{offset}, {}) exceeds capacity {max_positions}",
                offset.saturating_add(*tokens)
            ),
            Self::BatchHeadOverflow { batch, heads } => write!(
                formatter,
                "multi-head lane count overflows for batch {batch} and {heads} heads"
            ),
            Self::HeadLayout { input, source } => {
                write!(formatter, "{input} head layout: {source}")
            }
            Self::MergeLayout(source) => write!(formatter, "head output merge: {source}"),
            Self::RotaryConfiguration(source) => {
                write!(formatter, "multi-head RoPE configuration: {source}")
            }
            Self::Rotary { input, source } => write!(formatter, "{input} RoPE: {source}"),
            Self::CausalAttention(source) => source.fmt(formatter),
            Self::Autodiff { stage, source } => {
                write!(formatter, "multi-head {stage}: {source}")
            }
        }
    }
}

impl Error for MultiHeadAttentionError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::QkvProjection(source) => Some(source),
            Self::OutputProjection(source) => Some(source),
            Self::Initialization(source) => Some(source),
            Self::HeadLayout { source, .. } => Some(source),
            Self::MergeLayout(source) => Some(source),
            Self::RotaryConfiguration(source) => Some(source),
            Self::Rotary { source, .. } => Some(source),
            Self::CausalAttention(source) => Some(source),
            Self::Autodiff { source, .. } => Some(source),
            _ => None,
        }
    }
}

impl From<InitializationError> for MultiHeadAttentionError {
    fn from(source: InitializationError) -> Self {
        Self::Initialization(source)
    }
}

fn multi_autodiff(
    stage: MultiHeadStage,
) -> impl FnOnce(TensorAutodiffError) -> MultiHeadAttentionError {
    move |source| MultiHeadAttentionError::Autodiff { stage, source }
}

fn head_layout(input: MultiHeadInput) -> impl FnOnce(HeadLayoutError) -> MultiHeadAttentionError {
    move |source| MultiHeadAttentionError::HeadLayout { input, source }
}

fn rotary_error(input: MultiHeadInput) -> impl FnOnce(RopeError) -> MultiHeadAttentionError {
    move |source| MultiHeadAttentionError::Rotary { input, source }
}

fn validate_configuration(
    model_width: usize,
    heads: usize,
) -> Result<usize, MultiHeadAttentionError> {
    if model_width == 0 {
        return Err(MultiHeadAttentionError::ZeroModelWidth);
    }
    if heads == 0 {
        return Err(MultiHeadAttentionError::ZeroHeadCount);
    }
    if !model_width.is_multiple_of(heads) {
        return Err(MultiHeadAttentionError::ModelWidthNotDivisible { model_width, heads });
    }
    let head_width = model_width / heads;
    if !head_width.is_multiple_of(2) {
        return Err(MultiHeadAttentionError::OddHeadWidth { head_width });
    }
    Ok(head_width)
}

The forward path projects to full model width, splits all three tensors, rotates only $Q$ and $K$ , flattens batch and head into independent lanes, reuses the tested single-head causal operation, restores the head axis, merges, and applies the output projection. The returned teaching evidence exposes every boundary while the final output stays shaped $[B,T,d_{\mathrm{model}}]$ .

Compose full-width projections, per-head RoPE and causal attention, merge, and output projection rust/crates/llm-from-scratch/src/attention/multi_head.rs#multi-head-layer

/// Inspectable intermediate tensors from one complete multi-head forward pass.
#[derive(Clone, Debug)]
pub struct MultiHeadAttentionForward {
    projected_query_heads: TensorValue,
    projected_key_heads: TensorValue,
    projected_value_heads: TensorValue,
    rotated_query_heads: TensorValue,
    rotated_key_heads: TensorValue,
    attention_weights: TensorValue,
    head_outputs: TensorValue,
    merged: TensorValue,
    output: TensorValue,
}

impl MultiHeadAttentionForward {
    pub fn projected_query_heads(&self) -> &TensorValue {
        &self.projected_query_heads
    }

    pub fn projected_key_heads(&self) -> &TensorValue {
        &self.projected_key_heads
    }

    pub fn projected_value_heads(&self) -> &TensorValue {
        &self.projected_value_heads
    }

    pub fn rotated_query_heads(&self) -> &TensorValue {
        &self.rotated_query_heads
    }

    pub fn rotated_key_heads(&self) -> &TensorValue {
        &self.rotated_key_heads
    }

    pub fn attention_weights(&self) -> &TensorValue {
        &self.attention_weights
    }

    pub fn head_outputs(&self) -> &TensorValue {
        &self.head_outputs
    }

    pub fn merged(&self) -> &TensorValue {
        &self.merged
    }

    pub fn output(&self) -> &TensorValue {
        &self.output
    }

    pub fn into_output(self) -> TensorValue {
        self.output
    }
}

/// Full-width Q/K/V projections, per-head RoPE and causal attention, and W_O.
#[derive(Clone, Debug)]
pub struct MultiHeadAttention {
    qkv: QkvProjections,
    output: Linear,
    rope: RotaryEmbedding,
    parameters: NamedParameters,
    model_width: usize,
    heads: usize,
    head_width: usize,
}

impl MultiHeadAttention {
    /// Initializes Q, K, V, and O in order without partially advancing `rng`.
    pub fn new(
        parameter_prefix: impl Into<String>,
        model_width: usize,
        heads: usize,
        max_positions: usize,
        rope_base: f64,
        rng: &mut SplitMix64,
    ) -> Result<Self, MultiHeadAttentionError> {
        let head_width = validate_configuration(model_width, heads)?;
        let rope = RotaryEmbedding::new(head_width, max_positions, rope_base)
            .map_err(MultiHeadAttentionError::RotaryConfiguration)?;
        let parameter_prefix = parameter_prefix.into();
        let mut trial = rng.clone();
        let qkv = QkvProjections::new(&parameter_prefix, model_width, model_width, &mut trial)
            .map_err(MultiHeadAttentionError::QkvProjection)?;
        let output = Linear::new(
            format!("{parameter_prefix}.output"),
            model_width,
            model_width,
            false,
            &mut trial,
        )
        .map_err(MultiHeadAttentionError::OutputProjection)?;
        let layer = Self::from_validated_parts(qkv, output, rope, heads, head_width)?;
        *rng = trial;
        Ok(layer)
    }

    /// Builds an exact deterministic layer from four named full-width matrices.
    pub fn from_parameters(
        query_weight: NamedParameter,
        key_weight: NamedParameter,
        value_weight: NamedParameter,
        output_weight: NamedParameter,
        heads: usize,
        max_positions: usize,
        rope_base: f64,
    ) -> Result<Self, MultiHeadAttentionError> {
        let qkv = QkvProjections::from_weights(query_weight, key_weight, value_weight)
            .map_err(MultiHeadAttentionError::QkvProjection)?;
        let output = Linear::from_parameters(output_weight, None)
            .map_err(MultiHeadAttentionError::OutputProjection)?;
        let model_width = qkv.model_width();
        let head_width = validate_configuration(model_width, heads)?;
        if qkv.head_width() != model_width {
            return Err(MultiHeadAttentionError::QkvOutputWidthMismatch {
                model_width,
                projected_width: qkv.head_width(),
            });
        }
        if output.input_width() != model_width {
            return Err(MultiHeadAttentionError::OutputInputWidthMismatch {
                expected: model_width,
                actual: output.input_width(),
            });
        }
        if output.output_width() != model_width {
            return Err(MultiHeadAttentionError::OutputWidthMismatch {
                expected: model_width,
                actual: output.output_width(),
            });
        }
        let rope = RotaryEmbedding::new(head_width, max_positions, rope_base)
            .map_err(MultiHeadAttentionError::RotaryConfiguration)?;
        Self::from_validated_parts(qkv, output, rope, heads, head_width)
    }

    fn from_validated_parts(
        qkv: QkvProjections,
        output: Linear,
        rope: RotaryEmbedding,
        heads: usize,
        head_width: usize,
    ) -> Result<Self, MultiHeadAttentionError> {
        let model_width = qkv.model_width();
        let mut listed = Vec::with_capacity(4);
        listed.extend(qkv.parameters().iter().cloned());
        listed.push(output.weight().clone());
        let parameters = NamedParameters::try_new(listed)?;
        Ok(Self {
            qkv,
            output,
            rope,
            parameters,
            model_width,
            heads,
            head_width,
        })
    }

    /// Runs position-aware causal attention independently in every head.
    pub fn forward(
        &self,
        input: &TensorValue,
        position_offset: usize,
    ) -> Result<MultiHeadAttentionForward, MultiHeadAttentionError> {
        let shape = input.shape();
        if shape.len() != 3 {
            return Err(MultiHeadAttentionError::InputRank { rank: shape.len() });
        }
        if shape[2] != self.model_width {
            return Err(MultiHeadAttentionError::InputWidthMismatch {
                expected: self.model_width,
                actual: shape[2],
            });
        }
        if shape[1] == 0 {
            return Err(MultiHeadAttentionError::EmptyTokens);
        }
        let position_end = position_offset.checked_add(shape[1]).ok_or(
            MultiHeadAttentionError::PositionOffsetOverflow {
                offset: position_offset,
                tokens: shape[1],
            },
        )?;
        if position_end > self.rope.max_positions() {
            return Err(MultiHeadAttentionError::PositionRangeExceeded {
                offset: position_offset,
                tokens: shape[1],
                max_positions: self.rope.max_positions(),
            });
        }

        let projected = self
            .qkv
            .forward(input)
            .map_err(MultiHeadAttentionError::QkvProjection)?;
        let projected_query_heads = split_heads(projected.query(), self.heads)
            .map_err(head_layout(MultiHeadInput::Query))?;
        let projected_key_heads =
            split_heads(projected.key(), self.heads).map_err(head_layout(MultiHeadInput::Key))?;
        let projected_value_heads = split_heads(projected.value(), self.heads)
            .map_err(head_layout(MultiHeadInput::Value))?;
        let rotated_query_heads = self
            .rope
            .rotate(&projected_query_heads, position_offset)
            .map_err(rotary_error(MultiHeadInput::Query))?;
        let rotated_key_heads = self
            .rope
            .rotate(&projected_key_heads, position_offset)
            .map_err(rotary_error(MultiHeadInput::Key))?;

        let lanes =
            shape[0]
                .checked_mul(self.heads)
                .ok_or(MultiHeadAttentionError::BatchHeadOverflow {
                    batch: shape[0],
                    heads: self.heads,
                })?;
        let lane_shape = [lanes, shape[1], self.head_width];
        let query_lanes = rotated_query_heads
            .reshape(&lane_shape)
            .map_err(multi_autodiff(MultiHeadStage::QueryLanes))?;
        let key_lanes = rotated_key_heads
            .reshape(&lane_shape)
            .map_err(multi_autodiff(MultiHeadStage::KeyLanes))?;
        let value_lanes = projected_value_heads
            .reshape(&lane_shape)
            .map_err(multi_autodiff(MultiHeadStage::ValueLanes))?;
        let attended =
            causal_scaled_dot_product_self_attention(&query_lanes, &key_lanes, &value_lanes)
                .map_err(MultiHeadAttentionError::CausalAttention)?;
        let attention_weights = attended
            .weights()
            .reshape(&[shape[0], self.heads, shape[1], shape[1]])
            .map_err(multi_autodiff(MultiHeadStage::RestoreWeights))?;
        let head_outputs = attended
            .output()
            .reshape(&[shape[0], self.heads, shape[1], self.head_width])
            .map_err(multi_autodiff(MultiHeadStage::RestoreHeadOutputs))?;
        let merged = merge_heads(&head_outputs).map_err(MultiHeadAttentionError::MergeLayout)?;
        let output = self
            .output
            .forward(&merged)
            .map_err(MultiHeadAttentionError::OutputProjection)?;

        Ok(MultiHeadAttentionForward {
            projected_query_heads,
            projected_key_heads,
            projected_value_heads,
            rotated_query_heads,
            rotated_key_heads,
            attention_weights,
            head_outputs,
            merged,
            output,
        })
    }

    pub fn qkv(&self) -> &QkvProjections {
        &self.qkv
    }

    pub fn output_projection(&self) -> &Linear {
        &self.output
    }

    pub fn rope(&self) -> &RotaryEmbedding {
        &self.rope
    }

    pub fn parameters(&self) -> &[NamedParameter] {
        self.parameters.as_slice()
    }

    pub const fn model_width(&self) -> usize {
        self.model_width
    }

    pub const fn heads(&self) -> usize {
        self.heads
    }

    pub const fn head_width(&self) -> usize {
        self.head_width
    }

    pub const fn parameter_count(&self) -> usize {
        4 * self.model_width * self.model_width
    }
}

The nonuniform output seed produces finite gradient tensors for the input and all four matrices. Central differences probe all $12+4\cdot16=76$ coordinates with

\delta=10^{-6},\qquad \varepsilon_g=8\times10^{-6}.

Tests also prove bitwise split/merge inversion, one uniform and one nonuniform causal table, exact future zeros, head isolation before output projection, bitwise prefix invariance under a final-token perturbation, RoPE attention-weight preservation when unchanged query/key content receives one shared position offset, empty-batch differentiation, a single-token boundary, stable parameter names and identity, typed configuration and shape errors, finite saved tape context, and bitwise replay.

The example collects those values and checks in one reproducible report:

Collect the exact fixture, invariants, gradients, boundaries, and handoff rust/demos/ch30-multi-head-attention/src/lib.rs#learner-report

pub fn render_report(evidence: &LearnerEvidence) -> String {
    let primary = &evidence.primary;
    let parameters = &evidence.parameters;
    let shapes = &evidence.shapes;
    let errors = &evidence.errors;
    let gradients = &evidence.gradients;
    let history = &evidence.history;
    [
        "chapter=30-multi-head-attention".to_owned(),
        "prediction=projection creates two learned feature lanes; each lane normalizes its own causal rows; W_O first learns how to mix the concatenated results".to_owned(),
        format!(
            "config=batch:{BATCH} tokens:{TOKENS} d_model:{MODEL_WIDTH} heads:{HEADS} d_h:{HEAD_WIDTH} offset:0 capacity:{MAX_POSITIONS} rope_base:{ROPE_BASE:.6} bias:false"
        ),
        format!(
            "input=shape:{} values:{}",
            format_shape(primary.input.shape()),
            format_vector(primary.input.as_slice())
        ),
        format!(
            "projected_query_heads=shape:{} values:{}",
            format_shape(primary.projected_query_heads.shape()),
            format_vector(primary.projected_query_heads.as_slice())
        ),
        format!(
            "projected_key_heads=shape:{} values:{}",
            format_shape(primary.projected_key_heads.shape()),
            format_vector(primary.projected_key_heads.as_slice())
        ),
        format!(
            "projected_value_heads=shape:{} values:{}",
            format_shape(primary.projected_value_heads.shape()),
            format_vector(primary.projected_value_heads.as_slice())
        ),
        format!(
            "rotated_query_heads=shape:{} values:{}",
            format_shape(primary.rotated_query_heads.shape()),
            format_vector(primary.rotated_query_heads.as_slice())
        ),
        format!(
            "rotated_key_heads=shape:{} values:{}",
            format_shape(primary.rotated_key_heads.shape()),
            format_vector(primary.rotated_key_heads.as_slice())
        ),
        format!(
            "attention_weights=shape:{} values:{}",
            format_shape(primary.attention_weights.shape()),
            format_vector(primary.attention_weights.as_slice())
        ),
        format!(
            "head_outputs=shape:{} values:{}",
            format_shape(primary.head_outputs.shape()),
            format_vector(primary.head_outputs.as_slice())
        ),
        format!(
            "merged=shape:{} values:{}",
            format_shape(primary.merged.shape()),
            format_vector(primary.merged.as_slice())
        ),
        format!(
            "output_weight=shape:{} values:{}",
            format_shape(primary.output_weight.shape()),
            format_vector(primary.output_weight.as_slice())
        ),
        format!(
            "output=shape:{} values:{}",
            format_shape(primary.output.shape()),
            format_vector(primary.output.as_slice())
        ),
        format!(
            "heads=head_0_uniform:{} head_1_distinct:{} future_probabilities_zero:{}",
            primary.uniform_head_zero,
            primary.distinct_head_weights,
            primary.future_probabilities_zero
        ),
        format!(
            "prefix_perturbed_output={} position_0_unchanged:{} position_1_unchanged:{} position_2_changed:{}",
            format_vector(primary.prefix_perturbed_output.as_slice()),
            primary.prefix_zero_unchanged,
            primary.prefix_one_unchanged,
            primary.suffix_changed
        ),
        format!(
            "layout=split_merge_bitwise:{} head_isolation_before_output:{} common_offset_weights_preserved:{} tolerance:{INVARIANT_TOLERANCE:.12}",
            primary.split_merge_bitwise,
            primary.head_isolation_before_output,
            primary.common_offset_weights_preserved
        ),
        format!(
            "parameters=names:{} shapes:{} count:{} bias_free:{} node_distinct:{}",
            format_names(&parameters.names),
            format_shapes(&parameters.shapes),
            parameters.count,
            parameters.bias_free,
            parameters.node_distinct
        ),
        format!(
            "upstream={} loss:{:.6}",
            format_vector(primary.upstream.as_slice()),
            canonical(primary.loss)
        ),
        format!("input_gradient={}", format_vector(primary.input_gradient.as_slice())),
        format!(
            "query_weight_gradient={}",
            format_vector(primary.parameter_gradients[0].as_slice())
        ),
        format!(
            "key_weight_gradient={}",
            format_vector(primary.parameter_gradients[1].as_slice())
        ),
        format!(
            "value_weight_gradient={}",
            format_vector(primary.parameter_gradients[2].as_slice())
        ),
        format!(
            "output_weight_gradient={}",
            format_vector(primary.parameter_gradients[3].as_slice())
        ),
        format!(
            "gradcheck=input:{} query:{} key:{} value:{} output:{} total:{} tolerance:{GRADIENT_TOLERANCE:.6} passed:{}",
            gradients.input_checks,
            gradients.query_checks,
            gradients.key_checks,
            gradients.value_checks,
            gradients.output_checks,
            gradients.input_checks
                + gradients.query_checks
                + gradients.key_checks
                + gradients.value_checks
                + gradients.output_checks,
            gradients.passed
        ),
        format!(
            "shapes=input:{} split:{} rotated:{} weights:{} head_output:{} merged:{} output_weight:{} output:{} empty_batch_weights:{} empty_batch_output:{} single_token_weights:{}",
            format_shape(&shapes.input),
            format_shape(&shapes.split),
            format_shape(&shapes.rotated),
            format_shape(&shapes.weights),
            format_shape(&shapes.head_output),
            format_shape(&shapes.merged),
            format_shape(&shapes.output_weight),
            format_shape(&shapes.output),
            format_shape(&shapes.empty_batch_weights),
            format_shape(&shapes.empty_batch_output),
            format_shape(&shapes.single_token_weights)
        ),
        format!(
            "errors=zero_model_width:{} zero_heads:{} nondivisible:{} odd_head_width:{} input_rank:{} input_width:{} empty_tokens:{} offset_overflow:{} position_range:{} released_input:{}",
            errors.zero_model_width_rejected,
            errors.zero_heads_rejected,
            errors.nondivisible_rejected,
            errors.odd_head_width_rejected,
            errors.input_rank_rejected,
            errors.input_width_rejected,
            errors.empty_tokens_rejected,
            errors.offset_overflow_rejected,
            errors.position_range_rejected,
            errors.released_input_rejected
        ),
        format!(
            "history=earlier_weighted_context:{} earlier_distributions_per_target:{} single_head_shape:{} multi_head_shape:{} single_head_tables:{} multi_head_tables:{} rows_normalized:{} mixing:{} modern_example:{} weight_api:{}",
            format_vector(&history.earlier_weighted_context),
            history.earlier_distributions_per_target,
            format_shape(&history.single_head_weight_shape),
            format_shape(&history.multi_head_weight_shape),
            history.single_head_tables,
            history.multi_head_tables,
            history.all_rows_normalized,
            history.mixing_stage,
            history.modern_example,
            history.weight_api
        ),
        format!(
            "proof=tape_finite:{} replay:{} heads_distinct:{} causal:{} split_merge:{} gradients:{}",
            primary.tape_finite,
            if evidence.replay_bitwise { "bitwise" } else { "mismatch" },
            primary.distinct_head_weights,
            primary.future_probabilities_zero,
            primary.split_merge_bitwise,
            gradients.passed
        ),
        "next=wrap this attention transformation in the first pre-normalized residual path".to_owned(),
    ]
    .join("\n")
        + "\n"
}

The executable prints the report:

Run the complete multi-head attention example rust/demos/ch30-multi-head-attention/src/main.rs

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let evidence = ch30_multi_head_attention::learner_evidence()?;
    print!("{}", ch30_multi_head_attention::render_report(&evidence));
    Ok(())
}

Run cargo run --quiet --locked -p ch30-multi-head-attention. Its standard output matches rust/demos/ch30-multi-head-attention/expected.txt byte for byte, including the final newline.

See exactly where separate heads become one model-width row

The diagram uses values from the same fixture: projected and rotated rows, all six causal probability rows, every head output, all three concatenated rows, the four output-matrix rows, the three final rows, and the prefix check.

Produce the multi-head attention values used in the diagram rust/demos/ch30-multi-head-attention/src/diagram_trace.rs#multi-head-trace

pub fn render_trace(evidence: &LearnerEvidence) -> String {
    let primary = &evidence.primary;
    let shapes = &evidence.shapes;
    let parameters = &evidence.parameters;
    let gradients = &evidence.gradients;
    let gradient_checks = gradients.input_checks
        + gradients.query_checks
        + gradients.key_checks
        + gradients.value_checks
        + gradients.output_checks;
    debug_assert_eq!(gradient_checks, 76);
    let mut lines = vec![
        String::from(
            "CONFIG|batch=1|tokens=3|model_width=4|heads=2|head_width=2|offset=0|max_positions=6|rope_base=100.000000|bias=false|parameter_order=[query.weight,key.weight,value.weight,output.weight]|layout=reshape-transpose",
        ),
        format!("SHAPE|stage=input|value={}", format_shape(&shapes.input)),
        format!("SHAPE|stage=split|value={}", format_shape(&shapes.split)),
        format!(
            "SHAPE|stage=rotated|value={}",
            format_shape(&shapes.rotated)
        ),
        format!(
            "SHAPE|stage=weights|value={}",
            format_shape(&shapes.weights)
        ),
        format!(
            "SHAPE|stage=head-output|value={}",
            format_shape(&shapes.head_output)
        ),
        format!("SHAPE|stage=merged|value={}", format_shape(&shapes.merged)),
        format!(
            "SHAPE|stage=output-weight|value={}",
            format_shape(&shapes.output_weight)
        ),
        format!("SHAPE|stage=output|value={}", format_shape(&shapes.output)),
    ];
    lines.extend((0..HEADS).map(|head| partition_record(evidence, head)));
    lines.extend(
        (0..HEADS)
            .flat_map(|head| (0..TOKENS).map(move |query| (head, query)))
            .map(|(head, query)| weight_record(evidence, head, query)),
    );
    lines.extend(
        (0..HEADS)
            .flat_map(|head| (0..TOKENS).map(move |token| (head, token)))
            .map(|(head, token)| head_output_record(evidence, head, token)),
    );
    lines.extend((0..TOKENS).map(|token| merged_record(evidence, token)));
    lines.extend((0..MODEL_WIDTH).map(|row| output_map_record(evidence, row)));
    lines.extend((0..TOKENS).map(|token| output_record(evidence, token)));
    lines.push(format!(
        "PREFIX_PROOF|position_0={}|position_1={}|position_2={}|split_merge={}|head_isolation={}|future_probabilities={}|common_offset={}|tolerance={INVARIANT_TOLERANCE:.12}|parameters={}|gradchecks={}",
        if primary.prefix_zero_unchanged {
            "bitwise-unchanged"
        } else {
            "changed"
        },
        if primary.prefix_one_unchanged {
            "bitwise-unchanged"
        } else {
            "changed"
        },
        if primary.suffix_changed {
            "changed"
        } else {
            "unchanged"
        },
        if primary.split_merge_bitwise {
            "bitwise"
        } else {
            "mismatch"
        },
        if primary.head_isolation_before_output {
            "before-output"
        } else {
            "failed"
        },
        if primary.future_probabilities_zero {
            "exact-zero"
        } else {
            "nonzero"
        },
        if primary.common_offset_weights_preserved {
            "preserved"
        } else {
            "changed"
        },
        parameters.count,
        gradient_checks,
    ));
    debug_assert_eq!(lines.len(), 34);
    lines.join("\n") + "\n"
}

Keep two causal attention lanes separate until the output projection

Follow projected features through the head split, RoPE, two causal probability tables, value mixtures, concatenation, and the output projection.

Solid border marks Head 0
Dashed border marks Head 1
Allowed key: solid underline
Blocked key: dashed underline
Allowed diagonal: double border

Project first, then split the feature axis into heads

The packed matrices can read every input feature; identity weights are used here only to expose two visible blocks.

$[B,T,d_{\mathrm{model}}]\to[B,h,T,d_h]$

Input: $[1,3,4]$
After the head split: $[1,2,3,2]$
After RoPE: $[1,2,3,2]$
Attention probabilities: $[1,2,3,3]$
Head outputs: $[1,2,3,2]$
After concatenation: $[1,3,4]$
Output matrix: $[4,4]$
Layer output: $[1,3,4]$

The packed matrices can read every input feature; identity weights are used here only to expose two visible blocks.
Head and projected feature coordinates	Head 0 $0,1$	Head 1 $2,3$
Projected $Q$	$[1.000000,0.000000],$ $[0.540302,-0.841471],$ $[-0.416147,-0.909297]$	$[1.000000,0.000000],$ $[0.000000,1.000000],$ $[1.000000,1.000000]$
Projected $K$	$[1.000000,0.000000],$ $[0.540302,-0.841471],$ $[-0.416147,-0.909297]$	$[1.000000,0.000000],$ $[0.000000,1.000000],$ $[1.000000,1.000000]$
Projected $V$	$[1.000000,0.000000],$ $[0.540302,-0.841471],$ $[-0.416147,-0.909297]$	$[1.000000,0.000000],$ $[0.000000,1.000000],$ $[1.000000,1.000000]$
After RoPE $Q$	$[1.000000,0.000000],$ $[1.000000,0.000000],$ $[1.000000,0.000000]$	$[1.000000,0.000000],$ $[-0.841471,0.540302],$ $[-1.325444,0.493151]$
After RoPE $K$	$[1.000000,0.000000],$ $[1.000000,0.000000],$ $[1.000000,0.000000]$	$[1.000000,0.000000],$ $[-0.841471,0.540302],$ $[-1.325444,0.493151]$

Normalize one causal probability table inside each head

Each head has its own score matrix, row softmax, and value mixture while sharing the same causal visibility triangle.

$A_i=\operatorname{softmax}_{j}\!\left(\widetilde Q_i\widetilde K_i^\top/\sqrt{d_h}+M\right)$ $H_i=A_iV_i$

Head 0

Causal attention probabilities: Head 0
Query position	$k=0$	$k=1$	$k=2$	Row sum	Per-head value mixtures: Head output
$q=0$	$1.000000$ Allowed key · Allowed diagonal	$0.000000$ Blocked key	$0.000000$ Blocked key	$1.000000$	$[1.000000,0.000000]$
$q=1$	$0.500000$ Allowed key	$0.500000$ Allowed key · Allowed diagonal	$0.000000$ Blocked key	$1.000000$	$[0.770151,-0.420735]$
$q=2$	$0.333333$ Allowed key	$0.333333$ Allowed key	$0.333333$ Allowed key · Allowed diagonal	$1.000000$	$[0.374718,-0.583589]$

Head 1

Causal attention probabilities: Head 1
Query position	$k=0$	$k=1$	$k=2$	Row sum	Per-head value mixtures: Head output
$q=0$	$1.000000$ Allowed key · Allowed diagonal	$0.000000$ Blocked key	$0.000000$ Blocked key	$1.000000$	$[1.000000,0.000000]$
$q=1$	$0.213809$ Allowed key	$0.786191$ Allowed key · Allowed diagonal	$0.000000$ Blocked key	$1.000000$	$[0.213809,0.786191]$
$q=2$	$0.054696$ Allowed key	$0.370956$ Allowed key	$0.574348$ Allowed key · Allowed diagonal	$1.000000$	$[0.629044,0.945304]$

Concatenate head outputs, then apply the learned output map

Concatenation preserves both feature blocks; the chosen output matrix then swaps them so the two operations cannot be confused.

$\operatorname{MHA}(X)=\operatorname{Concat}(H_1,\ldots,H_h)W_O$

Concatenated head outputs
Token position	Head 0	Head 1	Concatenated row
$t=0$	$[1.000000,0.000000]$	$[1.000000,0.000000]$	$[1.000000,0.000000,1.000000,0.000000]$
$t=1$	$[0.770151,-0.420735]$	$[0.213809,0.786191]$	$[0.770151,-0.420735,0.213809,0.786191]$
$t=2$	$[0.374718,-0.583589]$	$[0.629044,0.945304]$	$[0.374718,-0.583589,0.629044,0.945304]$

Output-projection row
Output-projection row	$W_O$
$r=0$	$[0.000000,0.000000,1.000000,0.000000]$
$r=1$	$[0.000000,0.000000,0.000000,1.000000]$
$r=2$	$[1.000000,0.000000,0.000000,0.000000]$
$r=3$	$[0.000000,1.000000,0.000000,0.000000]$

Rows before and after the output projection
Token position	Before projection	After projection
$t=0$	$[1.000000,0.000000,1.000000,0.000000]$	$[1.000000,0.000000,1.000000,0.000000]$
$t=1$	$[0.770151,-0.420735,0.213809,0.786191]$	$[0.213809,0.786191,0.770151,-0.420735]$
$t=2$	$[0.374718,-0.583589,0.629044,0.945304]$	$[0.629044,0.945304,0.374718,-0.583589]$

Check prefix invariance and assembly invariants

Changing only the last input row leaves output rows zero and one bitwise unchanged in both heads and after output projection.

$t=0$: Unchanged · Checked
$t=1$: Unchanged · Checked
$t=2$: Changed · Checked

Split followed by merge Restores the tensor exactly
Influence between head outputs None before the output projection
Future-key probabilities Exactly zero
Common position shift with unchanged content Preserves attention probabilities
Learned scalar parameters $64$
Coordinates checked by central differences $76$

Suffix perturbation result $\varepsilon=0.000000000001$

Solid and dashed borders distinguish the two heads. Solid and dashed underlines separate visible keys from causally blocked keys, and the double border marks the allowed diagonal. The two probability tables are normalized independently, and the shape and prefix checks show that their outputs remain separate until the output projection.

Read token position $1$ closely. Before $W_O$ , the concatenated row is

[0.770151,-0.420735,0.213809,0.786191].

The chosen output matrix swaps its two two-coordinate blocks, producing

[0.213809,0.786191,0.770151,-0.420735].

That visible change is the reason to keep concatenation and output projection as two distinct stages.

Test shape, visibility, and mixing—not a story about head roles

For $d_{\mathrm{model}}=6$ , classify $h=1$ , $h=2$ , $h=3$ , $h=4$ , and $h=6$ under this chapter’s divisibility-plus-even-RoPE rule.
Trace every shape for $B=2$ , $T=5$ , $d_{\mathrm{model}}=12$ , and $h=3$ .
Fill the complete three-token causal visibility triangle for both heads.
Explain why one full-width attention calculation is not equivalent to two head-local attention calculations.
Apply the fixture’s block-swapping $W_O$ to symbolic trace rows $H^{(0)}=[a,b]$ and $H^{(1)}=[c,d]$ . Then start from identity $W_O$ , also set $(W_O)_{0,2}=1$ , and derive the new output coordinate $y_2$ .
Change only token position $2$ . Predict which output rows may change and which must remain bitwise identical.
Explain where the fixture’s block isolation ends and why a general dense projection does not isolate raw input slices.
Compute the target layer’s no-bias parameter count at $d_{\mathrm{model}}=128$ .

Check the structural answers

$h=1$ gives even $d_h=6$ ; $h=3$ gives even $d_h=2$ . Both pass. $h=2$ gives odd $d_h=3$ ; $h=6$ gives odd $d_h=1$ . Both fail this full-head RoPE rule. $h=4$ does not divide $6$ .
$Q$ , $K$ , and $V$ are $[2,5,12]$ ; their split forms are $[2,3,5,4]$ ; each head’s probabilities are $[2,5,5]$ , or $[2,3,5,5]$ together; head outputs are $[2,3,5,4]$ ; merge and final output are $[2,5,12]$ .
Each head keeps one, then two, then three keys. The two visibility triangles match even though their allowed probabilities may differ.
A single full-width calculation produces one normalized $[5,5]$ score table shared by all value features. Two heads produce two separately normalized $[5,5]$ tables, each with its own denominator for every query row; their numerical values can coincide for some parameter settings.
Concatenation gives $[a,b,c,d]$ ; the frozen output matrix gives $[c,d,a,b]$ . Starting from identity and adding $(W_O)_{0,2}=1$ makes the third output coordinate $y_2=a+c$ , so it receives contributions from both heads.
Rows $0$ and $1$ cannot see token $2$ , so they remain bitwise identical. Row $2$ may change in both heads and after $W_O$ .
Identity query/key/value matrices make the fixture’s blocks visible. General packed projection columns may read every input feature; only their projected outputs occupy separate lanes. $W_O$ can mix those lanes again.
The count is $4\cdot128^2=65{,}536$ scalar parameters.

Misconception: multi-head attention splits the raw input—or a completed single-head result—into chunks.

Correction: learned projections run before the split. Each projected head then computes its own rotated query/key scores, causal softmax, and value mixture. Only the finished head outputs are concatenated, and $W_O$ may mix them. Distinct fixture tables demonstrate one valid parameter setting, not guaranteed semantic specialization in every trained model.

Put this transformation on a pre-normalized residual path

The cumulative decoder now has a complete position-aware causal attention transformation. It accepts and returns $[B,T,d_{\mathrm{model}}]$ , but it does not yet normalize its input or add its result back to the residual stream.

Chapter 31 will place it in the first pre-normalized residual path,

x'=x+\operatorname{MHA}(\operatorname{RMSNorm}(x)),

then add the feed-forward residual path. Key/value cache ownership and incremental decoding remain deferred until Chapter 37.