31 · Content revision 2

Compose one pre-norm Transformer decoder block

Learn how RMSNorm, causal multi-head attention, SwiGLU, and two residual paths compose one shape-preserving Transformer decoder block.

Trace both residual paths before running the block

Compose one differentiable pre-normalized decoder block and verify the exact order of its attention and feed-forward residual paths. Trace one batch of three model-width-four token rows through two RMSNorm gains, two rotary causal attention heads, identity-like SwiGLU projections, and two residual additions.

The fixture uses $B=1$ , $T=3$ , $d_{\mathrm{model}}=4$ , $h=2$ , and $d_{\mathrm{ff}}=4$ . Its input is

X= \begin{bmatrix} 2&0&0&0\\ 0&2&0&0\\ 0&0&2&0 \end{bmatrix}.

Both RMSNorm gain vectors are $[1,1,1,1]$ , and this teaching fixture uses $\varepsilon=0$ . Every input row therefore has root mean square $1$ . The attention query, key, value, and output matrices are identity matrices. The SwiGLU gate, up, and down matrices are also identity matrices.

Those choices reveal operation order; they do not turn the block into an identity function. Causal attention mixes visible token rows. The intermediate residual state is normalized again. SwiGLU applies $\operatorname{SiLU}(z)\odot z$ to each token row.

Before looking at the output, write the six transformations and two bypasses:

X\xrightarrow{\operatorname{RMSNorm}_a}N_a \xrightarrow{\operatorname{MHA}}A, \qquad X'=X+A,

X'\xrightarrow{\operatorname{RMSNorm}_f}N_f \xrightarrow{\operatorname{FFN}}F, \qquad Y=X'+F.

All seven named model-width tensors have shape $[1,3,4]$ . The attention probabilities alone have shape $[1,2,3,3]$ . For token position $1$ , predict which value reaches the first identity path unchanged and which value reaches attention only after $\operatorname{RMSNorm}_a$ .

Add each branch to the stream that entered it

The complete pre-normalized block is

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

The two appearances of $\operatorname{RMSNorm}$ own separate learned gains. The first branch uses token-mixing $\operatorname{MHA}$ ; the second applies the same learned feature map independently at each token position through $\operatorname{FFN}$ . Both transformation outputs must return to model width so their residual additions are defined:

x,x',y\in\mathbb{R}^{B\times T\times d_{\mathrm{model}}}.

With bias-free projections, the two gain vectors, four attention matrices, and three SwiGLU matrices contain

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

For the frozen widths, $N_\theta=8+64+48=120$ . The example backpropagates a nonuniform scalar loss through all $12$ input coordinates and all $120$ parameter coordinates, checking $132$ analytic gradients against central differences.

Keep branch values separate from residual values

$x$ is the model-width residual stream entering the block.
$x'$ is the stream after the causal attention result is added to $x$ .
$y$ is the block output after the feed-forward result is added to $x'$ .
$\operatorname{RMSNorm}_a$ and $\operatorname{RMSNorm}_f$ use the same operation but own distinct gain vectors.
$\operatorname{MHA}$ is the bias-free rotary causal multi-head attention transformation from Chapter 30.
$\operatorname{FFN}$ is the bias-free SwiGLU transformation from Chapter 20.
$B$ is batch size, $T$ is token count, $d_{\mathrm{model}}$ is residual-stream width, and $d_{\mathrm{ff}}$ is the hidden width of SwiGLU.
$N_a$ and $N_f$ are normalized branch inputs. $A$ and $F$ are learned branch outputs. None of those four names denotes an identity path.

The order matters. Moving either normalizer after its residual merge changes the function even when all component weights stay fixed.

From recurrent state and post-norm blocks to pre-norm decoders

Sequential recurrent state and the original Transformer’s post-normalized residual sublayers. These are the two bounded predecessors compared here.

Long Short-Term Memory is the source for the first boundary. Hochreiter and Schmidhuber introduce an explicitly recurrent architecture for long-time-lag learning, giving the chapter its sequential-state predecessor rather than a claim about every later LSTM language model. Recurrent LSTM language models advance a carried state one token step at a time, while the original Transformer placed LayerNorm after each residual merge; neither layout is the pre-normalized causal decoder block assembled here.

Attention Is All You Need defines the next architecture boundary. Vaswani and colleagues define the original Transformer sublayer output as a residual merge followed by LayerNorm and mask the decoder self-attention against future positions. Its first sublayer can be summarized as $\operatorname{LayerNorm}(x+\operatorname{MHA}(x))$ , which is post-normalized.

On Layer Normalization in the Transformer Architecture names and analyzes the change. Xiong and colleagues distinguish Post-LN from Pre-LN and analyze how placing normalization inside residual blocks changes gradient behavior at initialization. Pre-LN moves normalization onto each sublayer input, and LLaMA provides a bounded modern language-model example that combines input pre-normalization with RMSNorm, causal attention, RoPE, and SwiGLU.

LLaMA supplies that bounded modern example. Touvron and colleagues report a causal Transformer language model that normalizes each sublayer input with RMSNorm and uses SwiGLU and RoPE. This block keeps a same-shaped residual stream while alternating token-mixing causal attention with per-token feature transformation, providing the repeatable unit that Chapter 32 will stack into a decoder-only language model.

LSTM improved long-lag learning but retained stepwise recurrent state; the Transformer removed recurrence but originally normalized after residual merges. Pre-normalization moves the normalizer before each transformation, and a modern causal decoder block uses that order to preserve an explicit residual stream around attention and feed-forward work.

The runnable contrast computes a small RNN-style carried state, then sends the same frozen Transformer components through the chapter’s first pre-norm residual stage and through a residual-then-normalize alternative. It demonstrates the operation-order difference without claiming that the helper reproduces an LSTM, a complete historical model, or any paper’s trained weights.

Compare a sequential carried state with fixed post-norm and pre-norm residual orderings rust/demos/ch31-decoder-block/src/lib.rs#historical-block-order-contrast

#[derive(Clone, Debug, PartialEq)]
pub struct HistoryEvidence {
    pub rnn_style_states: Vec<f64>,
    pub sequential_recurrence: bool,
    pub original_post_norm: bool,
    pub modern_pre_norm: bool,
    pub numeric_order_contrast: bool,
}

/// A bounded serial recurrence; this is RNN-style evidence, not an LSTM gate implementation.
pub fn rnn_style_states(inputs: &[f64]) -> Vec<f64> {
    let mut state = 0.0_f64;
    inputs
        .iter()
        .map(|input| {
            state = (0.5 * input + 0.75 * state).tanh();
            state
        })
        .collect()
}

/// Unit-gain, zero-bias LayerNorm rows for the original post-norm ordering contrast.
pub fn layer_norm_rows(input: &Tensor, epsilon: f64) -> Tensor {
    let width = *input
        .shape()
        .last()
        .expect("history input has a feature axis");
    let mut normalized = Vec::with_capacity(input.len());
    for row in input.as_slice().chunks_exact(width) {
        let mean = row.iter().sum::<f64>() / width as f64;
        let variance = row
            .iter()
            .map(|value| {
                let centered = value - mean;
                centered * centered
            })
            .sum::<f64>()
            / width as f64;
        let inverse_standard_deviation = (variance + epsilon).sqrt().recip();
        normalized.extend(
            row.iter()
                .map(|value| (value - mean) * inverse_standard_deviation),
        );
    }
    Tensor::from_vec(input.shape().to_vec(), normalized).expect("history shape is unchanged")
}

Compose tested parts without hiding their boundaries

DecoderBlockError preserves ownership of configuration, component-width, input-shape, position, released-tape, normalization, attention, feed-forward, and residual failures. Construction uses a trial random stream and commits it only after every component and cross-component width check passes.

Attribute every invalid boundary to configuration, input, position, tape, or the stage that failed rust/crates/llm-from-scratch/src/models/decoder_block.rs#decoder-block-errors

/// A component whose feature width is incompatible with the residual stream.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum DecoderBlockComponent {
    AttentionNorm,
    FeedForwardNorm,
    FeedForwardInput,
    FeedForwardOutput,
}

impl fmt::Display for DecoderBlockComponent {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(match self {
            Self::AttentionNorm => "attention RMSNorm",
            Self::FeedForwardNorm => "feed-forward RMSNorm",
            Self::FeedForwardInput => "feed-forward input",
            Self::FeedForwardOutput => "feed-forward output",
        })
    }
}

/// A rejected component assembly or stage of one decoder-block forward pass.
#[derive(Clone, Debug, PartialEq)]
pub enum DecoderBlockError {
    AttentionNorm(RmsNormError),
    Attention(MultiHeadAttentionError),
    AttentionResidual(ResidualError),
    FeedForwardNorm(RmsNormError),
    FeedForward(SwiGluError),
    FeedForwardResidual(ResidualError),
    ComponentWidthMismatch {
        component: DecoderBlockComponent,
        expected: usize,
        actual: usize,
    },
    Initialization(InitializationError),
}

impl fmt::Display for DecoderBlockError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::AttentionNorm(source) => write!(formatter, "attention RMSNorm: {source}"),
            Self::Attention(source) => write!(formatter, "causal multi-head attention: {source}"),
            Self::AttentionResidual(source) => {
                write!(formatter, "attention residual merge: {source}")
            }
            Self::FeedForwardNorm(source) => {
                write!(formatter, "feed-forward RMSNorm: {source}")
            }
            Self::FeedForward(source) => write!(formatter, "SwiGLU feed-forward: {source}"),
            Self::FeedForwardResidual(source) => {
                write!(formatter, "feed-forward residual merge: {source}")
            }
            Self::ComponentWidthMismatch {
                component,
                expected,
                actual,
            } => write!(
                formatter,
                "decoder-block {component} width must be {expected}, got {actual}"
            ),
            Self::Initialization(source) => source.fmt(formatter),
        }
    }
}

impl Error for DecoderBlockError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::AttentionNorm(source) | Self::FeedForwardNorm(source) => Some(source),
            Self::Attention(source) => Some(source),
            Self::AttentionResidual(source) | Self::FeedForwardResidual(source) => Some(source),
            Self::FeedForward(source) => Some(source),
            Self::Initialization(source) => Some(source),
            Self::ComponentWidthMismatch { .. } => None,
        }
    }
}

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

DecoderBlock owns two RmsNorm values, one MultiHeadAttention, and one SwiGlu. Its stable parameter order is the attention-normalizer gain, four attention matrices, the feed-forward-normalizer gain, then the gate, up, and down matrices. The two gains are separate tensors even when their frozen values match.

The public forward path calls normalization, attention, and the first residual addition before it calls the second normalization, SwiGLU, and the second residual addition. Returned evidence keeps all six stage values and the complete attention and feed-forward evidence on one differentiable tape.

Run two independently normalized transformations and add each result to its entering residual stream rust/crates/llm-from-scratch/src/models/decoder_block.rs#decoder-block-layer

/// Every inspectable value produced by the two pre-normalized residual paths.
#[derive(Clone, Debug)]
pub struct DecoderBlockForward {
    attention_norm: RmsNormForward,
    attention: MultiHeadAttentionForward,
    after_attention: TensorValue,
    feed_forward_norm: RmsNormForward,
    feed_forward: SwiGluForward,
    output: TensorValue,
}

impl DecoderBlockForward {
    pub fn attention_norm(&self) -> &RmsNormForward {
        &self.attention_norm
    }

    pub fn attention(&self) -> &MultiHeadAttentionForward {
        &self.attention
    }

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

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

    pub fn feed_forward_norm(&self) -> &RmsNormForward {
        &self.feed_forward_norm
    }

    pub fn feed_forward(&self) -> &SwiGluForward {
        &self.feed_forward
    }

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

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

/// RMSNorm → causal MHA → residual, then RMSNorm → SwiGLU → residual.
#[derive(Clone, Debug)]
pub struct DecoderBlock {
    attention_norm: RmsNorm,
    attention: MultiHeadAttention,
    feed_forward_norm: RmsNorm,
    feed_forward: SwiGlu,
    parameters: NamedParameters,
    model_width: usize,
}

impl DecoderBlock {
    /// Initializes every matrix transactionally from one deterministic stream.
    pub fn new(
        parameter_prefix: impl Into<String>,
        config: DecoderBlockConfig,
        rng: &mut SplitMix64,
    ) -> Result<Self, DecoderBlockError> {
        let parameter_prefix = parameter_prefix.into();
        let attention_norm = RmsNorm::new(
            format!("{parameter_prefix}.attention_norm.gain"),
            config.model_width,
            config.rms_epsilon,
        )
        .map_err(DecoderBlockError::AttentionNorm)?;
        let feed_forward_norm = RmsNorm::new(
            format!("{parameter_prefix}.ffn_norm.gain"),
            config.model_width,
            config.rms_epsilon,
        )
        .map_err(DecoderBlockError::FeedForwardNorm)?;

        let mut trial = rng.clone();
        let attention = MultiHeadAttention::new(
            format!("{parameter_prefix}.attention"),
            config.model_width,
            config.heads,
            config.max_positions,
            config.rope_base,
            &mut trial,
        )
        .map_err(DecoderBlockError::Attention)?;
        let feed_forward = SwiGlu::new(
            format!("{parameter_prefix}.ffn"),
            config.model_width,
            config.feed_forward_width,
            config.model_width,
            &mut trial,
        )
        .map_err(DecoderBlockError::FeedForward)?;
        let block = Self::from_parts(attention_norm, attention, feed_forward_norm, feed_forward)?;
        *rng = trial;
        Ok(block)
    }

    /// Assembles already named deterministic components after cross-width checks.
    pub fn from_parts(
        attention_norm: RmsNorm,
        attention: MultiHeadAttention,
        feed_forward_norm: RmsNorm,
        feed_forward: SwiGlu,
    ) -> Result<Self, DecoderBlockError> {
        let model_width = attention.model_width();
        for (component, actual) in [
            (
                DecoderBlockComponent::AttentionNorm,
                attention_norm.feature_width(),
            ),
            (
                DecoderBlockComponent::FeedForwardNorm,
                feed_forward_norm.feature_width(),
            ),
            (
                DecoderBlockComponent::FeedForwardInput,
                feed_forward.input_width(),
            ),
            (
                DecoderBlockComponent::FeedForwardOutput,
                feed_forward.output_width(),
            ),
        ] {
            if actual != model_width {
                return Err(DecoderBlockError::ComponentWidthMismatch {
                    component,
                    expected: model_width,
                    actual,
                });
            }
        }

        let mut listed = Vec::with_capacity(9);
        listed.extend(attention_norm.parameters().iter().cloned());
        listed.extend(attention.parameters().iter().cloned());
        listed.extend(feed_forward_norm.parameters().iter().cloned());
        listed.extend(feed_forward.parameters().iter().cloned());
        let parameters = NamedParameters::try_new(listed)?;

        Ok(Self {
            attention_norm,
            attention,
            feed_forward_norm,
            feed_forward,
            parameters,
            model_width,
        })
    }

    /// Runs the two transformation branches in exact pre-normalized order.
    pub fn forward(
        &self,
        input: &TensorValue,
        position_offset: usize,
    ) -> Result<DecoderBlockForward, DecoderBlockError> {
        let attention_norm = self
            .attention_norm
            .forward_with_intermediates(input)
            .map_err(DecoderBlockError::AttentionNorm)?;
        let attention = self
            .attention
            .forward(attention_norm.output(), position_offset)
            .map_err(DecoderBlockError::Attention)?;
        let after_attention = residual_add(input, attention.output())
            .map_err(DecoderBlockError::AttentionResidual)?;
        let feed_forward_norm = self
            .feed_forward_norm
            .forward_with_intermediates(&after_attention)
            .map_err(DecoderBlockError::FeedForwardNorm)?;
        let feed_forward = self
            .feed_forward
            .forward_with_intermediates(feed_forward_norm.output())
            .map_err(DecoderBlockError::FeedForward)?;
        let output = residual_add(&after_attention, feed_forward.output())
            .map_err(DecoderBlockError::FeedForwardResidual)?;

        Ok(DecoderBlockForward {
            attention_norm,
            attention,
            after_attention,
            feed_forward_norm,
            feed_forward,
            output,
        })
    }

    pub fn attention_norm(&self) -> &RmsNorm {
        &self.attention_norm
    }

    pub fn attention(&self) -> &MultiHeadAttention {
        &self.attention
    }

    pub fn feed_forward_norm(&self) -> &RmsNorm {
        &self.feed_forward_norm
    }

    pub fn feed_forward(&self) -> &SwiGlu {
        &self.feed_forward
    }

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

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

    pub fn parameter_count(&self) -> usize {
        self.parameters
            .as_slice()
            .iter()
            .map(|parameter| parameter.tensor().value().len())
            .sum()
    }
}

The fixture verifies shape preservation, exact future-mask zeros, bitwise prefix invariance under a final-token perturbation, stable parameter names, distinct parameter identity, deterministic replay, transactional initialization, and all declared invalid boundaries. Its central-difference step and tolerance are

\delta=10^{-6},\qquad \tau=2\times10^{-5}.

Every one of the $132$ checked coordinates has a finite analytic gradient and a finite numerical comparison inside that tolerance. The learner report freezes these conclusions without replacing the tests:

Format the exact fixture, order, causality, parameters, gradients, errors, replay, and next boundary rust/demos/ch31-decoder-block/src/lib.rs#learner-report

pub fn render_report(evidence: &LearnerEvidence) -> String {
    let primary = &evidence.primary;
    let parameters = &evidence.parameters;
    let errors = &evidence.errors;
    let gradients = &evidence.gradients;
    let history = &evidence.history;
    [
        "chapter=31-decoder-block".to_owned(),
        format!(
            "config=batch:{BATCH} tokens:{TOKENS} model_width:{MODEL_WIDTH} heads:{HEADS} head_width:{HEAD_WIDTH} feed_forward_width:{FEED_FORWARD_WIDTH} epsilon:{RMS_EPSILON:.6}"
        ),
        format!(
            "shape=input:{} attention_norm:{} attention_weights:{} attention_branch:{} after_attention:{} feed_forward_norm:{} feed_forward_branch:{} output:{} probe_logits:{}",
            format_shape(&evidence.shapes.input),
            format_shape(&evidence.shapes.attention_norm),
            format_shape(&evidence.shapes.attention_weights),
            format_shape(&evidence.shapes.attention_branch),
            format_shape(&evidence.shapes.after_attention),
            format_shape(&evidence.shapes.feed_forward_norm),
            format_shape(&evidence.shapes.feed_forward_branch),
            format_shape(&evidence.shapes.output),
            format_shape(&evidence.shapes.probe_logits),
        ),
        format!(
            "order=attention_norm->attention->residual->feed_forward_norm->feed_forward->residual pre_norm:{} post_norm_differs:{}",
            primary.pre_norm_order, primary.post_norm_differs
        ),
        format!(
            "causality=prefix_0_bitwise:{} prefix_1_bitwise:{} suffix_changed:{} future_probabilities_zero:{}",
            primary.prefix_zero_unchanged,
            primary.prefix_one_unchanged,
            primary.suffix_changed,
            primary.future_probabilities_zero,
        ),
        format!(
            "parameters=tensors:{} scalars:{} bias_free:{} stable_order:{} distinct:{}",
            parameters.tensors,
            parameters.scalars,
            parameters.bias_free,
            parameters.stable_order,
            parameters.node_distinct,
        ),
        format!(
            "gradcheck=input:{} parameters:{} total:{} tolerance:{GRADIENT_TOLERANCE:.6} passed:{} tape_finite:{}",
            gradients.input_checks,
            gradients.parameter_checks,
            gradients.input_checks + gradients.parameter_checks,
            gradients.passed,
            primary.tape_finite,
        ),
        format!(
            "errors=configuration:{} component_width:{} input_rank:{} input_width:{} empty_tokens:{} position_range:{} released_input:{}",
            errors.configuration_rejected,
            errors.component_width_rejected,
            errors.input_rank_rejected,
            errors.input_width_rejected,
            errors.empty_tokens_rejected,
            errors.position_range_rejected,
            errors.released_input_rejected,
        ),
        format!(
            "history=sequential_recurrence:{} original_post_norm:{} modern_pre_norm:{} numeric_order_contrast:{}",
            history.sequential_recurrence,
            history.original_post_norm,
            history.modern_pre_norm,
            history.numeric_order_contrast,
        ),
        format!(
            "replay={}",
            if evidence.replay_bitwise {
                "bitwise"
            } else {
                "mismatch"
            }
        ),
        "next=stack these blocks between token embeddings and a tied vocabulary head".to_owned(),
    ]
    .join("\n")
        + "\n"
}

The executable entry point prints only that report:

Print the frozen decoder-block learner report rust/demos/ch31-decoder-block/src/main.rs

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

Run cargo run --quiet --locked -p ch31-decoder-block. Its standard output matches rust/demos/ch31-decoder-block/expected.txt byte for byte, including the final newline.

Inspect both bypasses and both transformation branches

The separate trace executable emits shapes, all stage rows, causal probability rows, residual-merge provenance, probe logits, a numeric order contrast, the causality proof, parameter ownership, gradient totals, and the bounded history contrast.

Emit the strict thirty-three-line decoder-block trace rust/demos/ch31-decoder-block/src/diagram_trace.rs#decoder-block-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 history = &evidence.history;
    let mut lines = vec![
        String::from(
            "CONFIG|batch=1|tokens=3|model_width=4|heads=2|head_width=2|feed_forward_width=4|epsilon=0.000000|stage_order=[attention-norm,attention,residual-1,feed-forward-norm,feed-forward,residual-2]",
        ),
        format!("SHAPE|stage=input|value={}", format_shape(&shapes.input)),
        format!(
            "SHAPE|stage=attention-norm|value={}",
            format_shape(&shapes.attention_norm)
        ),
        format!(
            "SHAPE|stage=attention-weights|value={}",
            format_shape(&shapes.attention_weights)
        ),
        format!(
            "SHAPE|stage=attention-branch|value={}",
            format_shape(&shapes.attention_branch)
        ),
        format!(
            "SHAPE|stage=after-attention|value={}",
            format_shape(&shapes.after_attention)
        ),
        format!(
            "SHAPE|stage=feed-forward-norm|value={}",
            format_shape(&shapes.feed_forward_norm)
        ),
        format!(
            "SHAPE|stage=feed-forward-branch|value={}",
            format_shape(&shapes.feed_forward_branch)
        ),
        format!("SHAPE|stage=output|value={}", format_shape(&shapes.output)),
        format!(
            "SHAPE|stage=probe-logits|value={}",
            format_shape(&shapes.probe_logits)
        ),
        stage_record("input", primary.input.as_slice()),
        stage_record("attention-norm", primary.attention_norm.as_slice()),
        stage_record("attention-branch", primary.attention_branch.as_slice()),
        stage_record("after-attention", primary.after_attention.as_slice()),
        stage_record("feed-forward-norm", primary.feed_forward_norm.as_slice()),
        stage_record(
            "feed-forward-branch",
            primary.feed_forward_branch.as_slice(),
        ),
        stage_record("output", primary.output.as_slice()),
    ];
    lines.extend(
        (0..HEADS)
            .flat_map(|head| (0..TOKENS).map(move |query| (head, query)))
            .map(|(head, query)| weight_record(evidence, head, query)),
    );
    lines.push(format!(
        "MERGE|name=attention|identity=input|branch=attention-branch|result=after-attention|exact={}",
        primary.first_residual_exact
    ));
    lines.push(format!(
        "MERGE|name=feed-forward|identity=after-attention|branch=feed-forward-branch|result=output|exact={}",
        primary.second_residual_exact
    ));
    lines.extend((0..TOKENS).map(|token| {
        let start = token * 3;
        format!(
            "PROBE|token={token}|values={}",
            format_vector(&primary.probe_logits.as_slice()[start..start + 3])
        )
    }));
    lines.push(format!(
        "ORDER_PROOF|pre_norm={}|post_norm_differs={}|post_norm_token_1={}|pre_norm_token_1={}",
        primary.pre_norm_order,
        primary.post_norm_differs,
        format_vector(token_row(primary.post_norm_first_stage.as_slice(), 1)),
        format_vector(token_row(primary.after_attention.as_slice(), 1)),
    ));
    lines.push(format!(
        "CAUSAL_PROOF|position_0={}|position_1={}|position_2={}|future_probabilities={}",
        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.future_probabilities_zero {
            "exact-zero"
        } else {
            "nonzero"
        },
    ));
    lines.push(format!(
        "PARAMETERS|tensors={}|scalars={}|bias={}|stable_order={}|distinct={}|names=[{}]",
        parameters.tensors,
        parameters.scalars,
        !parameters.bias_free,
        parameters.stable_order,
        parameters.node_distinct,
        parameters.names.join(",")
    ));
    lines.push(format!(
        "GRADIENTS|input={}|parameters={}|total={}|tolerance={GRADIENT_TOLERANCE:.6}|passed={}|tape_finite={}",
        gradients.input_checks,
        gradients.parameter_checks,
        gradients.input_checks + gradients.parameter_checks,
        gradients.passed,
        primary.tape_finite,
    ));
    lines.push(format!(
        "HISTORY|rnn_style_states={}|sequential={}|original_post_norm={}|modern_pre_norm={}|numeric_order_contrast={}",
        format_vector(&history.rnn_style_states),
        history.sequential_recurrence,
        history.original_post_norm,
        history.modern_pre_norm,
        history.numeric_order_contrast,
    ));
    debug_assert_eq!(lines.len(), 33);
    lines.join("\n") + "\n"
}

Follow two pre-normalized branches around one residual stream

Follow three exact token rows through attention normalization, causal multi-head attention, the first residual merge, feed-forward normalization, SwiGLU, and the second residual merge.

Solid border: unchanged identity path
Dashed border: learned transformation branch
Double border: residual addition
Solid underline: visible key
Dashed underline: masked future key

One shape-preserving block, in source order

$x'=x+\operatorname{MHA}(\operatorname{RMSNorm}(x))$ $y=x'+\operatorname{FFN}(\operatorname{RMSNorm}(x'))$

Input residual stream: $[1,3,4]$
Attention-normalized rows: $[1,3,4]$
Causal attention weights: $[1,2,3,3]$
Attention branch output: $[1,3,4]$
After attention residual: $[1,3,4]$
Feed-forward-normalized rows: $[1,3,4]$
SwiGLU branch output: $[1,3,4]$
Block output: $[1,3,4]$
Probe logits: $[1,3,3]$

Normalize before the causal attention branch

The identity copy bypasses RMSNorm and attention; only the dashed branch mixes visible token positions.

Input residual stream

x=[0.000000,

2.000000,

0.000000,

0.000000]

[1,3,4]

Identity path

x=[0.000000,

2.000000,

0.000000,

0.000000]

Attention RMSNorm input

N_a=[0.000000,

2.000000,

0.000000,

0.000000]

Causal multi-head attention

A=[0.010881,

1.989119,

0.000000,

0.000000]

Double border: residual addition

+

After attention residual

x'=[0.010881,

3.989119,

0.000000,

0.000000]

[1,3,4]

Normalize the intermediate stream before SwiGLU

The second identity copy bypasses the separate RMSNorm and per-token feature transformation.

After attention residual

x'=[0.010881,

3.989119,

0.000000,

0.000000]

[1,3,4]

Identity path

x'=[0.010881,

3.989119,

0.000000,

0.000000]

Feed-forward RMSNorm input

N_f=[0.005455,

1.999993,

0.000000,

0.000000]

SwiGLU feed-forward branch

F=[0.000015,

3.523159,

0.000000,

0.000000]

Double border: residual addition

+

Block output

y=[0.010896,

7.512278,

0.000000,

0.000000]

[1,3,4]

Check exact stage values, causal rows, order, parameters, and gradients

Exact stage rows expose operation order, causal isolation, parameter ownership, and the complete gradient comparison.

Scrollable decoder-block stage evidence table: Normalize before the causal attention branch
Token position	Input residual stream	Attention-normalized rows	Attention branch output	After attention residual
$t=0$	$[2.000000,0.000000,0.000000,0.000000]$	$[2.000000,0.000000,0.000000,0.000000]$	$[2.000000,0.000000,0.000000,0.000000]$	$[4.000000,0.000000,0.000000,0.000000]$
$t=1$	$[0.000000,2.000000,0.000000,0.000000]$	$[0.000000,2.000000,0.000000,0.000000]$	$[0.010881,1.989119,0.000000,0.000000]$	$[0.010881,3.989119,0.000000,0.000000]$
$t=2$	$[0.000000,0.000000,2.000000,0.000000]$	$[0.000000,0.000000,2.000000,0.000000]$	$[0.666667,0.666667,1.788570,0.000000]$	$[0.666667,0.666667,3.788570,0.000000]$

Scrollable decoder-block stage evidence table: Normalize the intermediate stream before SwiGLU
Token position	Feed-forward-normalized rows	SwiGLU branch output	Block output
$t=0$	$[2.000000,0.000000,0.000000,0.000000]$	$[3.523188,0.000000,0.000000,0.000000]$	$[7.523188,0.000000,0.000000,0.000000]$
$t=1$	$[0.005455,1.999993,0.000000,0.000000]$	$[0.000015,3.523159,0.000000,0.000000]$	$[0.010896,7.512278,0.000000,0.000000]$
$t=2$	$[0.341520,0.341520,1.940806,0.000000]$	$[0.068180,0.068180,3.293781,0.000000]$	$[0.734847,0.734847,7.082351,0.000000]$

Scrollable decoder-block stage evidence table: Check exact stage values, causal rows, order, parameters, and gradients
Token position	Probe logits
$t=0$	$[7.523188,0.000000,-7.523188]$
$t=1$	$[0.010896,7.512278,-7.523174]$
$t=2$	$[7.817198,7.817198,-1.469694]$

Scrollable causal attention probability table
$h$	Query position	$k=0$	$k=1$	$k=2$	Row sum
$0$	$q=0$	$1.000000$ Visible	$0.000000$ Masked	$0.000000$ Masked	$1.000000$
$0$	$q=1$	$0.005440$ Visible	$0.994560$ Visible	$0.000000$ Masked	$1.000000$
$0$	$q=2$	$0.333333$ Visible	$0.333333$ Visible	$0.333333$ Visible	$1.000000$
$1$	$q=0$	$1.000000$ Visible	$0.000000$ Masked	$0.000000$ Masked	$1.000000$
$1$	$q=1$	$0.500000$ Visible	$0.500000$ Visible	$0.000000$ Masked	$1.000000$
$1$	$q=2$	$0.052857$ Visible	$0.052857$ Visible	$0.894285$ Visible	$1.000000$

Pre-norm and post-norm differ

Numerically different

$\operatorname{PostNorm}_1=[-0.573144,$ $1.732042,$ $-0.579449,$ $-0.579449]$

$x'=[0.010881,$ $3.989119,$ $0.000000,$ $0.000000]$

Earlier outputs ignore the changed final token

$t=0$ : Bitwise unchanged

$t=1$ : Bitwise unchanged

$t=2$ : Numerically different

Stable parameter ownership

$N_\theta=120$

tensors=9

bias=false

Complete central-difference check

$n=132$

$\tau=0.000020$

Verified from the Rust fixture

The stage rows show which value travels along each identity path and which value returns from each learned branch. The causal table isolates the only cross-token transformation, while the order comparison places the pre-norm and post-norm results side by side. Solid identity paths, dashed transformation branches, and double residual merges make those roles distinguishable without depending on color.

Read token position $1$ across the first branch. The input row is $[0,2,0,0]$ , while the attention result is $[0.010881,1.989119,0,0]$ . Their first residual merge produces

x'_1=[0.010881,3.989119,0,0].

Normalizing after that merge instead produces $[-0.573144,1.732042,-0.579449,-0.579449]$ in the fixed contrast. The mismatch is direct numeric evidence that pre-norm and post-norm orderings are not interchangeable.

Test order, shape, causality, and parameter ownership

Put these operations in order: first residual addition, feed-forward RMSNorm, attention RMSNorm, SwiGLU, causal attention, second residual addition.
For $B=2$ , $T=5$ , and $d_{\mathrm{model}}=12$ , write every residual-stream shape. Which axis may either residual addition change?
If the attention branch returns zero, derive $x'$ . Does that make the whole block an identity function?
Change only the final input token. Which output rows must remain unchanged, and which transformation establishes that boundary?
For $d_{\mathrm{model}}=12$ and $d_{\mathrm{ff}}=32$ , compute the complete bias-free parameter count, including both RMSNorm gains.
Rewrite only the first residual stage in post-norm order and identify where its value can first diverge.

Check the structural answers

Attention RMSNorm, causal attention, first residual addition, feed-forward RMSNorm, SwiGLU, second residual addition.
Every residual-stream value is $[2,5,12]$ . Neither residual addition may change any axis; both operands must have identical shapes.
$x'=x$ . The feed-forward branch can still change $y$ , so the whole block need not be an identity.
Rows $0$ through $3$ must stay unchanged. The causal mask in attention blocks the suffix from reaching earlier rows; the per-token feed-forward map cannot introduce cross-token mixing afterward.
$2(12)+4(12)^2+3(12)(32)=1752$ scalar parameters.
Post-norm forms $\operatorname{LayerNorm}(x+\operatorname{MHA}(x))$ ; pre-norm forms $x+\operatorname{MHA}(\operatorname{RMSNorm}(x))$ . Their attention inputs can already differ before the residual result exists.

Misconception: pre-normalization means normalizing once before the complete block. Correction: each transformation has its own normalizer, and each learned result is added to the unnormalized residual stream that entered that branch.

Repeat the block only at the next model boundary

The cumulative implementation now has one complete depth-one causal decoder block; Chapter 32 will repeat it between token embeddings, a final RMSNorm, and a tied vocabulary projection.

This chapter’s input and output both have shape $[B,T,d_{\mathrm{model}}]$ . Attention mixes visible token history, SwiGLU transforms features independently at each position, and two residual paths keep direct routes around both transformations. Stacking, final normalization, vocabulary projection, checkpoint ownership, and cache state remain outside this block.