← All chapters

Chapter 07 · Content revision 4

Configure one decoder core before construction

Validate model and run records, plan resources before constructing the existing CPU decoder, then preserve full and cached text through one shared decoder path.

The problem: dimension choices must agree

Changing a decoder’s width changes its embeddings, attention projections, feed-forward matrices and normalization gains. Changing its context capacity changes the storage needed for cached tokens. If those choices are made independently, a model can have incompatible shapes or reserve more memory than its execution environment allows. Constructing arrays first discovers such problems after the expensive work has begun.

We will describe the model once, calculate its shapes and resource needs, and allow construction only after the selected execution path passes its checks. The decoder family stays the same: causal text attention with pre-RMSNorm, RoPE, SwiGLU, no biases, and a vocabulary head tied to the token embedding table. Dimensions are runtime data. Reference and bridge records use the existing CPU implementation; larger records can produce useful plans while execution remains refused.

Solve the tiny parameter census

Take vocabulary size V=266V=266, residual width D=4D=4, one complete block L=1L=1, one query head and one key/value head, and feed-forward width F=4F=4. The model can represent positions up to its configured capacity C=4C=4.

The embedding table contains 266⋅4=1064266\cdot4=1064 learned scalars. Its transpose also supplies the vocabulary head, so that head owns no second table. The four attention matrices contain 4⋅42=644\cdot4^2=64 scalars. The three SwiGLU matrices contain 3⋅4⋅4=483\cdot4\cdot4=48. Two block RMSNorm gains and one final gain contain 3⋅4=123\cdot4=12. Together they give

1064+64+48+12=11881064+64+48+12=1188

unique learned parameters. This census describes storage ownership, so using the embedding twice does not count its scalars twice. A separate untied head would add another 10641064, producing 22522252 parameters and a different model architecture.

Derive shapes before counting parameters

Let HqH_q count query heads and HkvH_{kv} count key/value heads. Each head has width dhd_h. The total key or value projection width is KK:

dh=D/Hq,K=Hkvdh.d_h=D/H_q,\qquad K=H_{kv}d_h.

The divisions must be exact: HqH_q divides DD, and HkvH_{kv} divides HqH_q, with Hkv≤HqH_{kv}\le H_q. The selected RoPE policy also requires even dhd_h. Truncating a quotient would describe incomplete head groups. An unequal-head shape can be valid for planning even though its execution kernel is not yet available here.

For a bias-free block, the query and output matrices contribute 2D22D^2, key and value contribute 2DK2DK, the three feed-forward matrices contribute 3DF3DF, and two norm gains contribute 2D2D. Call their sum pp. The embedding and final gain sit outside the LL blocks:

p=2D2+2DK+3DF+2D,P=VD+Lp+D.\begin{aligned} p&=2D^2+2DK+3DF+2D,\\ P&=VD+Lp+D. \end{aligned}

PP counts unique learned scalar values. RoPE adds no learned position table, and the tied head adds no independent vocabulary matrix. This is why changing CC alone leaves PP unchanged for this family.

Four checked model records: literal dimensions, unique learned scalar counts and current execution scope
Model quantityReferenceBridgeLaptop planProduction plan
Vocabulary entries VV26626616384128000
Residual features DD4165128192
Complete blocks LL12880
Query heads HqH_q12864
Key/value heads HkvH_{kv}1228
SwiGLU hidden features FF420153628672
Position capacity CC41651232768
Unique learned scalars PP118883043251456069500936192
Current executionCPU supportedCPU supportedRefusedRefused

Keep dimensions, capacity and token budgets distinct

VV counts vocabulary entries; DD counts residual features per token; LL counts complete blocks; HqH_q and HkvH_{kv} count heads; and FF counts SwiGLU hidden features. The derived widths dhd_h and KK count features, while pp and PP count learned scalars per block and per model.

CC counts available token positions per row. BB is the nominal row count in one full-context forward accounting unit. A particular call may use fewer admitted rows or positions. NN is a declared training-token budget, separate from these tensor axes. Reference uses B=16,N=2048B=16,N=2048, bridge B=8,N=65536B=8,N=65536, laptop B=1,N=20000000B=1,N=20000000, and production B=1,N=15000000000000B=1,N=15000000000000. These budgets do not mean the chapter trains on those tokens.

Later resource formulas use bkvb_{kv} for bytes per stored key/value element. The examples explicitly choose either two-byte accounting or eight-byte f64 CPU storage. A byte count requires both the element count and this width.

Let one checked plan control construction

The semantic model record owns dimensions and capacity, vocabulary/tokenizer bindings, numerical policies, dtype and Chapter 6’s initialization policy. The run record owns nominal rows, token budget, microbatch and accumulation, named initialization/shuffle/sampling seeds, topology and backend choices. Changing a run seed changes the run identity while the semantic identity stays the same. Changing capacity from 512512 to 128128 changes semantic identity while the parameter count stays 3251456032514560.

The two identities use SHA-256 over the shared ArtifactIdentity framing. Its model domain is decoder-model-v1, its run domain decoder-run-v1, and its identity schema version is one. A configuration string holds the exact canonical model or run record, including its final line feed; the framed identity is itself canonically encoded before hashing.

The accepted record encoding is compact UTF-8 JSON with recursively sorted object keys, preserved array order and one final line feed. Parsing is bounded to one MiB; strict typed objects reject duplicate or unknown fields. The parser also requires the decoded record to encode back to the exact input bytes, so it does not silently insert defaults or normalize a bound record.

Dimensions and budgets use canonical unsigned integers. Resource arithmetic uses checked u128 products, sums and alignment, followed by range checks for each output representation. RoPE base and RMSNorm epsilon are separate finite f64 semantics whose exact bit patterns belong to identity; they do not supply floating-point approximations of resource counts. Dimensions must be positive, RoPE base positive and epsilon nonnegative. Zero seeds and zero epsilon remain allowed. No overflow, inexact division or resource limit is repaired by clamping a model.

Validate head grouping and derive the checked named tensor census shared by decoder planning and construction. rust/crates/llm-from-scratch-practical/src/config/model.rs#configurable-model-shapes
/// Zero blocks remain legal for the existing legacy API, not the typed profile.
#[derive(Clone, Copy, Debug)]
pub struct DecoderShape {
    pub v: u128,
    pub d: u128,
    pub l: u128,
    pub hq: u128,
    pub hkv: u128,
    pub f: u128,
}
#[derive(Clone, Debug, Serialize)]
pub struct TensorGroup {
    pub name: &'static str,
    pub category: &'static str,
    pub rows: u64,
    pub columns: u64,
    pub repetitions: u64,
    pub partition_axis: &'static str,
    pub shape: Vec<u64>,
}
impl TensorGroup {
    pub fn elements_per_tensor(&self) -> Result<u128, ConfigError> {
        mul(self.rows.into(), self.columns.into(), "tensor shape")
    }
    pub fn elements(&self) -> Result<u128, ConfigError> {
        mul(
            self.elements_per_tensor()?,
            self.repetitions.into(),
            "tensor repetitions",
        )
    }
}
#[derive(Clone, Debug, Serialize)]
pub struct ParameterCensus {
    pub embedding: u64,
    pub attention: u64,
    pub feed_forward: u64,
    pub norms: u64,
    pub total: u64,
    pub tensors: Vec<TensorGroup>,
    pub tied_head_alias: &'static str,
}
impl DecoderShape {
    pub fn validate(self) -> Result<(), ConfigError> {
        if [self.v, self.d, self.hq, self.hkv, self.f].contains(&0)
            || self.hkv > self.hq
            || !self.d.is_multiple_of(self.hq)
            || !self.hq.is_multiple_of(self.hkv)
        {
            return Err(ConfigError::Invalid("head grouping or dimensions"));
        }
        if !(self.d / self.hq).is_multiple_of(2) {
            return Err(ConfigError::Invalid("even RoPE head width"));
        }
        Ok(())
    }
    pub fn head_width(self) -> Result<u128, ConfigError> {
        self.validate()?;
        Ok(self.d / self.hq)
    }
    pub fn kv_width(self) -> Result<u128, ConfigError> {
        mul(self.hkv, self.head_width()?, "KV width")
    }
    pub fn census(self) -> Result<ParameterCensus, ConfigError> {
        self.validate()?;
        let k = self.kv_width()?;
        let mut tensors = Vec::with_capacity(11);
        for (name, category, rows, cols, repetitions, axis) in [
            (
                "token_embedding.weight",
                "embedding",
                self.v,
                self.d,
                1,
                "rows",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[0],
                "norms",
                self.d,
                1,
                self.l,
                "replicated",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[1],
                "attention",
                self.d,
                self.d,
                self.l,
                "columns",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[2],
                "attention",
                self.d,
                k,
                self.l,
                "columns",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[3],
                "attention",
                self.d,
                k,
                self.l,
                "columns",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[4],
                "attention",
                self.d,
                self.d,
                self.l,
                "rows",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[5],
                "norms",
                self.d,
                1,
                self.l,
                "replicated",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[6],
                "feed_forward",
                self.d,
                self.f,
                self.l,
                "columns",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[7],
                "feed_forward",
                self.d,
                self.f,
                self.l,
                "columns",
            ),
            (
                BLOCK_PARAMETER_SUFFIXES[8],
                "feed_forward",
                self.f,
                self.d,
                self.l,
                "rows",
            ),
            ("final_norm.gain", "norms", self.d, 1, 1, "replicated"),
        ] {
            tensors.push(TensorGroup {
                name,
                category,
                rows: checked_u64(rows)?,
                columns: checked_u64(cols)?,
                repetitions: checked_u64(repetitions)?,
                partition_axis: axis,
                shape: if category == "norms" {
                    vec![checked_u64(rows)?]
                } else {
                    vec![checked_u64(rows)?, checked_u64(cols)?]
                },
            });
        }
        let mut totals = [0u128; 4];
        for t in &tensors {
            let i = match t.category {
                "embedding" => 0,
                "attention" => 1,
                "feed_forward" => 2,
                _ => 3,
            };
            totals[i] = add(totals[i], t.elements()?, "parameter group")?;
        }
        let total = totals
            .iter()
            .try_fold(0, |a, &b| add(a, b, "parameter census"))?;
        Ok(ParameterCensus {
            embedding: checked_u64(totals[0])?,
            attention: checked_u64(totals[1])?,
            feed_forward: checked_u64(totals[2])?,
            norms: checked_u64(totals[3])?,
            total: checked_u64(total)?,
            tensors,
            tied_head_alias: "token_embedding.weight",
        })
    }
}

Construction follows this order:

  1. Parse and validate the two records.
  2. Calculate the named tensor, state, activation, KV and communication plans.
  3. Compare execution support and every applicable profile limit.
  4. Admit the checked plan, then construct from that same plan.

Only bounded parser and report metadata may be allocated before admission. Model, cache, optimizer and device arrays, and initialization draws, come after it. An invalid shape rejects during validation. A valid unequal-head model can be planned but refuses execution because GQA support is absent. A plan can also refuse because its demand exceeds a resource partition or required workspace is unknown. These outcomes answer different questions.

Reference and bridge are the only executable profiles here. They initialize through the existing fresh-depth owner and use actual f64 computation. The smoke record explicitly changes laptop capacity to 128128 and binds its own identity; it does not clamp the laptop record. Smoke, seed-sensitivity, core, adapter and production execution remain refused. Adapter additionally lacks a selected compatible artifact.

The planner’s plan_decoder function produces the complete resource report. The excerpt below shows admit_decoder returning its private admission token and construct consuming that token through the existing fresh initializer.

Admit the checked decoder plan, then delegate construction and bind its configuration and actual initialization identity. rust/crates/llm-from-scratch-practical/src/config/planner.rs#decoder-resource-plan
/// Private fields prevent callers fabricating a positive allocation decision.
pub struct DecoderAdmission {
    plan: DecoderPlan,
}
pub fn admit_decoder(
    model: &ModelConfig,
    run: &RunConfig,
) -> Result<DecoderAdmission, ConfigError> {
    let plan = plan_decoder(model, run, ForwardMode::TrackedFull)?;
    if !plan.admissible {
        return Err(ConfigError::Refused(plan.refusals));
    }
    model.executable_config()?;
    Ok(DecoderAdmission { plan })
}
impl DecoderAdmission {
    pub fn plan(&self) -> &DecoderPlan {
        &self.plan
    }
    /// Delegates construction and the one-time transform to Chapter 6's owner.
    pub fn construct(self, rng: &mut SplitMix64) -> Result<DepthInitializedDecoder, ConfigError> {
        if rng.state() != self.plan.run.seeds.initialization {
            return Err(ConfigError::Invalid(
                "caller RNG differs from declared initialization seed",
            ));
        }
        let cfg = self.plan.model.executable_config()?;
        let fresh = FreshDepthInitialization::prepare(
            cfg,
            DepthInitializationBudget {
                max_parameters: checked_usize(self.plan.profile.max_parameters.into())?,
                max_positions: checked_usize(self.plan.profile.max_context.into())?,
            },
            rng,
        )
        .map_err(|_| ConfigError::Invalid("fresh decoder preparation"))?;
        let initialized = fresh
            .complete()
            .map_err(|_| ConfigError::Invalid("fresh decoder completion"))?;
        let record = initialized.record();
        let initial = record.initial_rng_state().to_string();
        let final_state = record.final_rng_state().to_string();
        let identity = ArtifactIdentity::new(
            "decoder-construction-v1",
            [
                ("semantic-sha256", self.plan.semantic_sha256.as_str()),
                ("run-sha256", self.plan.run_sha256.as_str()),
                ("initial-rng", initial.as_str()),
                ("final-rng", final_state.as_str()),
                ("policy", record.policy_id()),
                ("base-initializer", record.base_initializer_id()),
            ],
        )
        .and_then(|v| v.digest())
        .map_err(|_| ConfigError::Invalid("construction identity"))?;
        initialized
            .model()
            .bind_execution_identity(
                self.plan.semantic_sha256,
                self.plan.run_sha256,
                identity.to_string(),
                checked_usize(self.plan.run.microbatch.into())?,
            )
            .map_err(|_| ConfigError::Invalid("execution identity binding"))?;
        Ok(initialized)
    }
}

Preserve the text seam as it evolves

A checked configuration keeps model construction consistent. Chapter 7 originally preserved text behavior with two private prepared-input producers. Ordinary full, padded and packed input performed text-token embedding lookup through one model-owned producer. Cached text performed its one-new-token lookup through the second. Both wrapped the result in private PreparedDecoderInput and sent it to one causal decoder core, which then contained the block loop, final RMSNorm and tied vocabulary head.

In a full prepared tensor of shape [b,T,D][b,T,D], bb is the actual number of rows and TT the number of supplied positions per row; they fit the admitted row and position bounds. That original carrier kept checked batch metadata by reference for padded or packed input; ordinary and cached input carried their position and causal policies. Together these bound positions, implicit causal ranges, segment IDs, validity and loss eligibility, source provenance, semantic/run digests, dtype, device and construction identity without allocating a dense attention mask. The original core checked these fields against the model before running any block. It also checked the original embedding handle’s node and revision, then the actual parameter node identities and revisions, including the tied table.

A matching digest describes matching bound bytes. It does not prove that a hidden tensor came from the model’s lookup or that live parameters remain current. A same-shaped, equal-valued model with different nodes cannot use the first model’s prepared input or cache. The original private carrier offered no public raw embedding or modality entry point.

Chapter 10 evolves this seam into a checked vector input and a borrowing DecoderBackbone that returns hidden features. Text source checks stay in TextInputEncoder, targets and masks stay with text loss, and TextOutputHead applies the same live tied table outside the backbone. The live excerpts below show that evolved boundary; the graph, parameter and cache preservation rules continue to apply.

Check prepared vector shape, layout, graph policy and live backbone and learned-producer bindings before consumption. rust/crates/llm-from-scratch-practical/src/models/prepared_decoder_input.rs#prepared-text-input
// region:prepared-decoder-input
/// A checked [B,T,D] vector sequence. No token IDs, vocabulary, tokenizer,
/// targets or output table are required by this carrier.
///
/// Learned encoders must list every live parameter used to produce the vectors.
/// Handles are retained without copying learned values. A changed parameter
/// invalidates a prepared result before the transformer performs any work.
pub struct PreparedDecoderInput<'execution> {
    embedding: TensorValue,
    embedding_binding: TensorValueBinding,
    context: AutogradContext<'execution>,
    layout: SequenceLayout,
    provenance: String,
    config: DecoderBlockConfig,
    layers: usize,
    identity: Option<CoreIdentity>,
    parameters: Vec<TensorValueBinding>,
    producer_parameters: Vec<TensorValueBinding>,
}
impl<'execution> PreparedDecoderInput<'execution> {
    pub fn new(
        backbone: &DecoderBackbone<'_>,
        context: AutogradContext<'execution>,
        embedding: TensorValue,
        layout: SequenceLayout,
        provenance: impl Into<String>,
        producer_parameters: &[NamedParameter],
    ) -> Result<Self, PreparedInputError> {
        let input = Self {
            embedding_binding: TensorValueBinding::capture(&embedding),
            embedding,
            context,
            layout,
            provenance: provenance.into(),
            config: backbone.config(),
            layers: backbone.blocks().len(),
            identity: backbone.identity().cloned(),
            parameters: backbone
                .parameters()
                .map(|p| TensorValueBinding::capture(p.tensor()))
                .collect(),
            producer_parameters: producer_parameters
                .iter()
                .map(|p| TensorValueBinding::capture(p.tensor()))
                .collect(),
        };
        input.validate(backbone)?;
        Ok(input)
    }
    /// Recheck exact live nodes/revisions at consumption, not just preparation.
    pub fn validate(&self, backbone: &DecoderBackbone<'_>) -> Result<(), PreparedInputError> {
        let [b, t] = self.layout.shape();
        if self.embedding.shape() != [b, t, backbone.config().model_width()]
            || t > backbone.config().max_positions()
            || self.identity.as_ref().is_some_and(|id| b > id.max_rows)
        {
            return Err(PreparedInputError("B,T,D or admitted row count"));
        }
        if self
            .layout
            .positions()
            .iter()
            .any(|&p| p >= backbone.config().max_positions())
        {
            return Err(PreparedInputError("position capacity"));
        }
        if self.config != backbone.config()
            || self.layers != backbone.blocks().len()
            || self.identity.as_ref() != backbone.identity()
        {
            return Err(PreparedInputError(
                "backbone configuration or execution identity",
            ));
        }
        if self.provenance.is_empty() {
            return Err(PreparedInputError("empty producer provenance"));
        }
        if !self.embedding_binding.node_matches(&self.embedding)
            || !self.embedding_binding.revision_matches(&self.embedding)
        {
            return Err(PreparedInputError(
                "prepared hidden tensor identity or revision",
            ));
        }
        if !self.context.records_graph() && self.embedding.tracks_gradient() {
            return Err(PreparedInputError("graph-recording mode"));
        }
        let current: Vec<_> = backbone.parameters().collect();
        if self.parameters.len() != current.len()
            || self.parameters.iter().zip(current).any(|(binding, p)| {
                !binding.node_matches(p.tensor()) || !binding.revision_matches(p.tensor())
            })
        {
            return Err(PreparedInputError("backbone parameter node or revision"));
        }
        if self
            .producer_parameters
            .iter()
            .any(|binding| !binding.revision_matches(binding.value_handle()))
        {
            return Err(PreparedInputError("producer parameter revision"));
        }
        Ok(())
    }
    pub fn embeddings(&self) -> &TensorValue {
        &self.embedding
    }
    pub fn context(&self) -> AutogradContext<'execution> {
        self.context
    }
    pub fn layout(&self) -> &SequenceLayout {
        &self.layout
    }
    /// Descriptive provenance is never used in place of a live node binding.
    pub fn provenance(&self) -> &str {
        &self.provenance
    }
    pub(crate) fn producer_parameters(&self) -> &[TensorValueBinding] {
        &self.producer_parameters
    }
}
// endregion:prepared-decoder-input

Full execution propagates the caller’s explicit AutogradContext. With recording enabled, it keeps the original lookup handle and graph, so repeated token IDs accumulate lookup and tied-head contributions into the same learned table. Cached inference selects AutogradContext::no_grad() before embedding lookup and passes that context through every child operation. Its logits are graph-free and existing training gradients stay unchanged. Sharing the core does not turn detached cached history into full-prefix training.

Full packed input retains occurrence-local positions and causal ranges. Cached input retains its existing growing-prefix, one-new-token policy; it does not gain packed segment resets or ragged serving. Cache sessions keep their parameter guards, and every layer ticket plus the final output is checked before a token commits. A rejected token retains the preceding cache and work counters. If a later token fails during prefill, the whole prefill returns to empty logical cache and zero work while retaining physical capacity. Reset does not authorize stale parameter binding.

Run full and cached vector inputs through one decoder block loop and return final-normalized hidden features. rust/crates/llm-from-scratch-practical/src/models/causal_decoder_core.rs#shared-causal-decoder
/// Mode chooses execution, independently of the source that produced vectors.
fn consume_with_dropout(
    backbone: &DecoderBackbone<'_>,
    input: PreparedDecoderInput,
    dropout: &mut DropoutExecution<'_>,
    recompute_blocks: bool,
) -> Result<BackboneForward, CoreError> {
    // The layout borrows the carrier; keep it alive through the shared loop.
    let layout = FullLayout {
        lengths: Some(input.layout().lengths()),
        starts: Some(input.layout().range_starts()),
        positions: Some(input.layout().positions()),
    };
    consume(
        backbone,
        &input,
        Mode::Full(layout),
        None,
        dropout,
        recompute_blocks,
    )
}
fn consume(
    backbone: &DecoderBackbone<'_>,
    input: &PreparedDecoderInput,
    mode: Mode<'_>,
    prefix: Option<usize>,
    dropout: &mut DropoutExecution<'_>,
    recompute_blocks: bool,
) -> Result<BackboneForward, CoreError> {
    let _backbone_guards: Vec<_> = backbone.parameters().map(|p| p.tensor().value()).collect();
    let _producer_guards: Vec<_> = input
        .producer_parameters()
        .iter()
        .map(|p| p.value_handle().value())
        .collect();
    let _embedding_guard = input.embeddings().value();
    input
        .validate(backbone)
        .map_err(|e| CoreError::Prepared(e.reason()))?;
    if let Mode::Cached(layers) = &mode {
        let prefix = prefix.expect("cached execution supplies its prefix");
        if input.context().records_graph()
            || input.layout().shape() != [1, 1]
            || input.layout().lengths() != [1]
            || input.layout().range_starts() != [0]
            || input.layout().positions() != [prefix]
            || layers.len() != backbone.blocks.len()
        {
            return Err(CoreError::Prepared("cached mode, position or layers"));
        }
        for (cache, block) in layers.iter().zip(backbone.blocks) {
            if cache.len() != prefix
                || cache.batch_size() != 1
                || cache.capacity() != backbone.config.max_positions()
            {
                return Err(CoreError::Prepared("cache prefix/batch/capacity"));
            }
            block
                .attention()
                .validate_incremental_cache_binding(cache)
                .map_err(|_| CoreError::Prepared("cache parameter binding"))?;
        }
    }
    run(
        backbone,
        input.context(),
        input.embeddings().clone(),
        0,
        mode,
        dropout,
        recompute_blocks,
    )
}

/// Residual diagnostics enter the same internal engine.
#[cfg(test)]
pub(super) fn residual_tail(
    backbone: &DecoderBackbone<'_>,
    context: AutogradContext,
    current: TensorValue,
    first_site: usize,
    layout: FullLayout<'_>,
) -> Result<BackboneForward, DecoderModelError> {
    run(
        backbone,
        context,
        current,
        first_site,
        Mode::Full(layout),
        &mut DropoutExecution::Identity,
        false,
    )
    .map_err(full_error)
}

fn run(
    backbone: &DecoderBackbone<'_>,
    context: AutogradContext,
    mut current: TensorValue,
    first_site: usize,
    mode: Mode<'_>,
    dropout: &mut DropoutExecution<'_>,
    recompute_blocks: bool,
) -> Result<BackboneForward, CoreError> {
    let cached = matches!(mode, Mode::Cached(_));
    let site_count = 2 * backbone.blocks().len() + 1;
    if first_site >= site_count {
        return Err(CoreError::Full(
            DecoderModelError::ResidualSiteOutOfBounds {
                site: first_site,
                site_count,
            },
        ));
    }
    let mut first_layer = first_site / 2;
    if first_site % 2 == 1 {
        current = backbone.blocks()[first_layer]
            .forward_feed_forward_path_and_dropout(context, &current, first_layer, dropout)
            .map_err(|source| {
                CoreError::Full(DecoderModelError::Block {
                    layer: first_layer,
                    source,
                })
            })?
            .2;
        first_layer += 1;
    }
    let mut blocks = Vec::new();
    let mut tickets = Vec::new();
    if cached {
        tickets
            .try_reserve_exact(backbone.blocks().len())
            .map_err(|_| {
                CoreError::Cached(DecoderKvCacheError::LayerAllocationFailed {
                    layers: backbone.blocks().len(),
                })
            })?;
    } else {
        blocks
            .try_reserve_exact(backbone.blocks().len() - first_layer)
            .map_err(|_| {
                CoreError::Full(DecoderModelError::LayerAllocationFailed {
                    layers: backbone.blocks().len(),
                })
            })?;
    }
    let mut score_values = 0usize;
    for (layer, block) in backbone.blocks().iter().enumerate().skip(first_layer) {
        match &mode {
            Mode::Full(layout) => {
                let checkpoint = recompute_blocks && context.records_graph();
                let execution = context.execution();
                let block_context = if checkpoint {
                    AutogradContext::no_grad().with_execution(&execution)
                } else {
                    context
                };
                // Compact execution exposes hidden/output tensors, not a
                // fabricated dense-probability report. All other block math is
                // shared with the inspectable reference and checkpoint replay.
                let (block_output, f) = if execution.attention_tiles().is_some() {
                    (
                        block
                            .forward_output_with_layout_and_dropout(
                                block_context,
                                &current,
                                0,
                                layout.lengths,
                                layout.starts,
                                layout.positions,
                                layer,
                                dropout,
                            )
                            .map_err(|source| {
                                CoreError::Full(DecoderModelError::Block { layer, source })
                            })?,
                        None,
                    )
                } else {
                    let f = block
                        .forward_with_layout_and_dropout(
                            block_context,
                            &current,
                            0,
                            layout.lengths,
                            layout.starts,
                            layout.positions,
                            layer,
                            dropout,
                        )
                        .map_err(|source| {
                            CoreError::Full(DecoderModelError::Block { layer, source })
                        })?;
                    (f.output().clone(), Some(f))
                };
                if checkpoint {
                    // Drop all inspectable intermediates. The recipe captures
                    // shared parameters, fixed RoPE tables, layout and tickets.
                    let output = block_output.value_snapshot();
                    drop(f);
                    let replay = dropout
                        .last_block_replay()
                        .map_err(|source| CoreError::Full(DecoderModelError::Dropout(source)))?;
                    let parameters = block
                        .parameters()
                        .iter()
                        .map(|p| p.tensor().clone())
                        .collect::<Vec<_>>();
                    let block = block.clone();
                    let rope = block.attention().rope();
                    let recipe_tensor_values = rope
                        .inverse_frequencies()
                        .len()
                        .checked_add(rope.cosines().len())
                        .and_then(|n| n.checked_add(rope.sines().len()))
                        .ok_or_else(|| {
                            CoreError::Full(DecoderModelError::Autodiff {
                                stage: super::decoder::DecoderModelStage::ActivationCheckpoint,
                                source: crate::tensor::storage::TensorError::ShapeOverflow.into(),
                            })
                        })?;
                    let lengths = layout.lengths.map(<[usize]>::to_vec);
                    let starts = layout.starts.map(<[usize]>::to_vec);
                    let positions = layout.positions.map(<[usize]>::to_vec);
                    current = TensorValue::checkpoint_with_context(context, &current, &parameters, output, recipe_tensor_values, move |boundary, context| {
                        let mut dropout = replay.execution();
                        let output = block.forward_output_with_layout_and_dropout(context, boundary, 0, lengths.as_deref(), starts.as_deref(), positions.as_deref(), layer, &mut dropout)
                            .map_err(|_| crate::autograd::tensor_core::TensorAutodiffError::Checkpoint { reason: "block forward or site replay failed during recomputation" })?;
                        dropout.finish_compact(None).map_err(|_| crate::autograd::tensor_core::TensorAutodiffError::Checkpoint { reason: "block replay did not consume its exact tickets" })?;
                        Ok(output)
                    }).map_err(|source| CoreError::Full(DecoderModelError::Autodiff { stage: super::decoder::DecoderModelStage::ActivationCheckpoint, source }))?;
                } else {
                    current = block_output;
                    if let Some(f) = f {
                        blocks.push(f);
                    }
                }
            }
            Mode::Cached(caches) => {
                let norm = block
                    .attention_norm()
                    .forward_with_context(context, &current)
                    .map_err(|source| {
                        CoreError::Cached(DecoderKvCacheError::AttentionNorm { layer, source })
                    })?;
                let ticket = block
                    .attention()
                    .prepare_incremental_bound_with_context(context, &norm, &caches[layer])
                    .map_err(|source| {
                        CoreError::Cached(DecoderKvCacheError::IncrementalAttention {
                            layer,
                            source,
                        })
                    })?;
                score_values = score_values
                    .checked_add(ticket.attention_score_values())
                    .ok_or(CoreError::Cached(DecoderKvCacheError::WorkOverflow {
                        counter: DecoderKvCacheCounter::AttentionScoreValues,
                    }))?;
                let after = residual_add_with_context(context, &current, ticket.output()).map_err(
                    |source| {
                        CoreError::Cached(DecoderKvCacheError::AttentionResidual { layer, source })
                    },
                )?;
                current = block
                    .forward_feed_forward_path(context, &after)
                    .map_err(|source| {
                        CoreError::Cached(match source {
                            DecoderBlockError::FeedForwardNorm(source) => {
                                DecoderKvCacheError::FeedForwardNorm { layer, source }
                            }
                            DecoderBlockError::FeedForward(source) => {
                                DecoderKvCacheError::FeedForward { layer, source }
                            }
                            DecoderBlockError::FeedForwardResidual(source) => {
                                DecoderKvCacheError::FeedForwardResidual { layer, source }
                            }
                            _ => unreachable!("only the FFN path ran"),
                        })
                    })?
                    .2;
                tickets.push(ticket);
            }
        }
    }
    let final_norm = backbone
        .final_norm()
        .forward_with_intermediates_and_context(context, &current)
        .map_err(|source| {
            if cached {
                CoreError::Cached(DecoderKvCacheError::FinalNorm(source))
            } else {
                CoreError::Full(DecoderModelError::FinalNorm(source))
            }
        })?;
    Ok(BackboneForward {
        blocks,
        final_norm,
        tickets,
        score_values,
    })
}

The configured demo supplies full inputs [0,99,100] with targets [99,100,1] and cached prefill [0,99] followed by 100. It initializes both supported profiles with seed 4701. Its raw negative-log-likelihood sum adds the three selected target losses before division; the mean divides by 33. Both reductions use the existing indexed-loss traversal. A nonfinite raw sum is refused, while the existing scaled mean path can still produce a finite mean. The earlier batch loss API, which multiplies the mean by valid-target count, stays unchanged.

Separate compatibility comparisons retain the legacy reference seed 39 and bridge seed 4701, including repeated inputs [0,99,99] with targets [99,99,1]. These are controlled valid token IDs, not a claim about a tokenizer’s segmentation. Each path’s before/after comparison covers its own logits, loss or cache, and work. Full-prefix versus cached-prefix comparison is forward only: it compares corresponding last-position logits, with absolute tolerance 10−1010^{-10} plus relative tolerance 10−910^{-9} when operation order differs. It supplies no cached-gradient parity claim.

Keep counts, simultaneous storage and measurements separate

KV storage uses physical capacity

A full-capacity KV cache reserves a key and a value in each block for every token position and row. With bkvb_{kv} bytes per element, the payload capacity is

KV⁡bytes=2BLCKbkv.\operatorname{KV}_{bytes}=2BLC Kb_{kv}.

The factor two counts keys and values. The formula uses key/value heads through KK, rather than query heads. A prefix with fewer tokens still owns its reserved capacity, and a logical reset does not make its allocated bytes zero.

A MAC convention identifies the counted work

A MAC is one multiply-accumulate pair. For one nominal full-context forward, matrix products contribute

Mlinear=BC [L(2D2+2DK+3DF)+DV].M_{\mathrm{linear}}=BC\,[L(2D^2+2DK+3DF)+DV].

For ordinary unsegmented causal attention, permitted triangular score and value-mix entries add BLDC(C+1)BLDC(C+1). Materializing the full dense score and mix products adds 2BLDC22BLDC^2 instead:

Mcausal=Mlinear+BLDC(C+1),Mdense=Mlinear+2BLDC2.\begin{aligned} M_{\mathrm{causal}}&=M_{\mathrm{linear}}+BLDC(C+1),\\ M_{\mathrm{dense}}&=M_{\mathrm{linear}}+2BLDC^2. \end{aligned}

Neither convention counts RMSNorm, RoPE, softmax, activation, masking or other scalar operations. Twice a MAC count covers multiply/add arithmetic only, not a complete FLOP total or hardware instruction count. Packed segments change permitted edges, so this unsegmented causal formula is not their exact edge count. Counts also establish no execution time.

Calculated full-context forward counts at nominal rows: two-byte KV comparison and two MAC conventions; these are not measured memory or timing
ScaleNominal rowsKV payload bytes at two bytes per elementCausal MAC pairsDense MAC pairs
Reference1610247654477312
Bridge81638411223041183744
Laptop plan120971521771883724818790481920
Production plan11073741824029810723756441603684738342584320

The two-byte column is a common accounting comparison. Replacing its element width with eight-byte f64 gives 40964096 and 6553665536 bytes for reference and bridge at the same nominal row counts. Those are comparison calculations. The existing cache supports one row; the actual reference and bridge cache payload capacities are 256256 and 81928192 bytes at capacities four and sixteen, before overhead. The small executed comparisons also use one row. Production’s two-byte KV alone is 10,737,418,24010{,}737{,}418{,}240 bytes, already greater than the laptop device envelope of 6,710,886,4006{,}710{,}886{,}400 bytes. That decisive refusal needs no model allocation.

State and lifetimes need allocation ownership

Persistent state lists each unique buffer’s element count, bytes per element and alignment. Count a master-weight copy only when it is a separate allocation, and round up each allocated buffer individually. Aliases such as the tied head own no second allocation. An illustrative state layout uses two-byte weights and four-byte master weights, gradients and two moment buffers: 18P18P bytes before alignment. The production subtotal is 1,251,016,851,4561{,}251{,}016{,}851{,}456 bytes. These five named allocations are an accounting example, not the actual CPU footprint of every model or optimizer. Persistent storage also includes nonlearned RoPE inverse frequencies, cosine and sine tables. They change stored bytes without changing the learned parameter census or the illustrative 18P18P subtotal. A tensor-parallel state ledger partitions declared matrix axes and keeps norm gains replicated; dividing the whole subtotal by rank count would miss that ownership. Batch layout metadata has its own byte ledger, separate from tensor payloads and runtime overhead. A profile’s slack partition is a ceiling: the selected schedule’s zero slack_reserve_bytes means no named allocation is reserved there. It does not turn unknown workspace into a zero demand.

Peak storage depends on what is live at the same event. In a synthetic schedule, keep 6464 bytes of state, allocate A with 6464 bytes, then B with 128128 bytes while B needs 3232 bytes of workspace. Their joint peak is 288288 bytes. Release A before allocating B and the peak falls to 224224 bytes. An alias of B adds no owned bytes. These results test lifetime accounting; they do not measure the decoder’s peak.

Tracked full execution keeps forward tensors needed for backward and their adjoints. Untracked full execution retains inspectable forward tensors until its result is dropped, while omitting the recording graph and adjoints. Cached execution reserves one row, physical KV capacity and staged layer tickets until commit. The selected symbolic schedule enumerates these buffers and computes their simultaneous resident peak. Its declared workspace contributes at its own event to the combined symbolic peak. If an actual backend’s required workspace is unknown, a complete execution peak cannot be asserted. The selected schedule remains distinct from an actual backend allocator bound. Unavailable GPU workspace stays Unknown and prevents execution admission; it is not a zero-byte allocation.

The CPU host plan additionally reserves a named 3232 MiB runtime/metadata envelope and explicit tensor and scratch bounds. Compiled-child process RSS includes more than tensor payloads and must fit that host plan and the profile’s 256256 MiB cap. These are different quantities from compiler memory, symbolic GPU storage and the toy liveness peak.

Communication counts sent bytes

The local production TP8 forward fixture partitions compatible matrix axes across eight ranks and keeps norm gains replicated. A full residual payload has 32768⋅8192=26843545632768\cdot8192=268435456 elements. Eight ring chunks contain 3355443233554432 elements each. Reduce-scatter plus all-gather sends 2(8−1)=142(8-1)=14 chunks per rank per event. Two events in each of 8080 blocks give 160160 events. At two bytes per element, this produces 939,524,096939{,}524{,}096 sent bytes per rank per event, 150,323,855,360150{,}323{,}855{,}360 per rank, and 1,202,590,842,8801{,}202{,}590{,}842{,}880 summed across all eight ranks.

This is calculated forward traffic. It excludes backward collectives, optimizer communication and protocol headers, and establishes no bandwidth or distributed execution. Traffic is separate from resident state. Only communication buffers with their actual planned lifetimes enter the memory ledger. Single-device profiles have no communication events; an unsupported topology receives an error rather than a zero-traffic fallback.

History: model storage and training budgets are different choices

A model-size description alone neither solves a device’s storage constraint nor chooses how many training tokens to process.

The Megatron-LM abstract describes an intra-layer model-parallel approach to large-Transformer memory constraints using communication operations; it does not validate the course’s TP8 ledger. It motivates separating partitioned storage from communication. Our named-matrix TP8 ledger is a local arithmetic fixture; it does not reproduce Megatron-LM’s implementation, cluster measurements or execution support.

The Chinchilla abstract studies model size and training-token count together under a compute budget, keeping those choices distinct. This study does not establish optimality for the course’s four scales. In our Rust comparison, changing only NN leaves the model’s PP unchanged and changes the run identity. This demonstrates distinct configuration responsibilities without transferring the paper’s empirical training result to an untrained toy decoder. A per-forward MAC count alone does not determine total training compute or quality.

A decoder configuration needs separate semantic dimensions, run budgets and topology-dependent resource accounting, with supported execution checked independently.

Compare local named-tensor partitions and sent traffic, then change a training-token budget while keeping the parameter census fixed. rust/demos/practical-ch07-configurable-decoder-core/src/lib.rs#configurable-budget-history
/// P and N are separate; partition traffic is a derived fixture, not speed.
pub fn budget_history() -> Result<Value, Box<dyn Error>> {
    let model = ModelConfig::scale("production")?;
    let mut run = RunConfig::fixture("production-plan-only")?;
    let original = planner::plan_decoder(&model, &run, ForwardMode::TrackedFull)?;
    run.n = 7_500_000_000_000;
    let fewer_tokens = planner::plan_decoder(&model, &run, ForwardMode::TrackedFull)?;
    run.world_size = 8;
    run.topology = "tensor-parallel-ring".into();
    run.sharding = "matrix-axes-v1".into();
    run.collectives = "two-allreduces-per-block-forward-v1".into();
    let ring = planner::plan_decoder(&model, &run, ForwardMode::TrackedFull)?;
    Ok(
        json!({"same_model_semantic_sha256":original.semantic_sha256,
        "token_budget_before":original.run.n,"token_budget_after":fewer_tokens.run.n,
        "parameters_before":original.census.total,"parameters_after":fewer_tokens.census.total,
        "run_digest_changed":original.run_sha256!=fewer_tokens.run_sha256,
        "single_device_traffic":original.traffic,"traffic_tp8":ring.traffic,
        "execution":"none; neither a Megatron-LM implementation nor a compute-optimality claim"}),
    )
}

Read the plan before either text path runs

Follow the model and run records into one checked plan. The plan admits construction only when execution support and resource limits pass. Once admitted, ordinary and cached text feed the shared decoder path. Chapter 10 separates its encoding, backbone and head while preserving that path. The reference census and separate synthetic lifetime schedules explain what the accounting counts.

One checked plan, two text paths, one decoder core

The semantic model record and run record enter shape validation and resource planning together. Construction occurs only after execution support and resource gates pass: reference and bridge use the CPU path, while laptop and production retain plans with execution refused. Ordinary full, padded or packed text lookup and cached one-token lookup converge on the same backbone, followed by the tied text head. Parameter counts are calculated from the reference record; the two lifetime peaks are synthetic accounting results, not measured decoder memory.

Two records enter planning

Semantic model record

Dimensions, capacity, numerical policies, vocabulary, dtype and initialization identity

Reference semantic model SHA-256
cb053bc16f5fafa2e5843ad30e3e43e26ba74e6d1d853823c30a3706f102bdfc
Run record

Rows, token budget, named seeds, topology, backend and resource profile

Reference run SHA-256
e2c50615797b78a2f95f7ff95b8e0684e80243c41f9647733aa4f705189f4985

The semantic-model and run-record identities shown identify the reference plan. Matching record digests do not replace actual parameter-node and revision checks at the input boundary.

Validate, plan, then admit

One checked shape descriptor determines the tensor census and resource ledgers. Execution support and every applicable resource limit are checked before construction.

Admitted CPU construction

Reference and bridge pass the supported CPU gates. The admitted plan supplies construction before ordinary or cached text runs.

Plan with execution refused

Laptop and production retain known accounting. Missing GQA or GPU support, exceeded limits or Unknown actual workspace prevent construction.

Reference parameter census

Calculated unique learned scalars; the tied vocabulary head shares the embedding allocation.

Embedding and tied head
1064
All attention matrices
64
All feed-forward matrices
48
All norm gains
12
Total unique learned scalars
1188
Separate synthetic lifetime schedules

Both schedules keep 64 bytes of state and use B with 128 bytes plus 32 bytes of workspace. Aliasing B adds no owned allocation.

A's 64 bytes remain live while B runs: joint peak in bytes
288
Release A's 64 bytes before B runs: joint peak in bytes
224

Ordinary and cached text share the backbone

Ordinary text producer

Full, padded and packed token lookup preserves the caller's context; recording keeps the original embedding graph.

Cached text producer

No-gradient mode is selected before one-token lookup; the cached producer preserves the growing-prefix binding.

One backbone, then the tied text head

Check prepared vectors and live bindings, run the shared block loop and final RMSNorm, then project through the wrapper's tied text head.

Optional practice: inspect the same plans and boundaries

Run the chapter demo from the repository root with Docker running. The launcher builds its workspace image as setup, then runs this CPU demo with locked offline Cargo and container networking disabled. Supplied fixtures require no corpus or model download. A successful default run emits a concise JSON report of totals, configuration records, profiles and refusal reasons, including the two supported CPU observations. The counts are calculated; the CPU observations come from the supplied forward, backward and cached calls.

./course run cargo run --offline --locked -p practical-ch07-configurable-decoder-core

To inspect the same fixtures’ named tensor, allocation and lifetime-event ledgers, use the trace option from the same repository root and Docker setup. It prints the complete JSON ledger rather than the concise report:

./course run cargo run --offline --locked -p practical-ch07-configurable-decoder-core -- --trace
Print decoder totals, refusals and supported CPU observations; use --trace for the same fixtures' complete accounting ledgers. rust/demos/practical-ch07-configurable-decoder-core/src/main.rs#configurable-decoder-example
fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args = std::env::args().skip(1).collect::<Vec<_>>();
    let output = match args.as_slice() {
        [] => practical_ch07_configurable_decoder_core::demo_report()?,
        [flag] if flag == "--trace" => practical_ch07_configurable_decoder_core::demo_trace()?,
        _ => return Err("usage: practical-ch07-configurable-decoder-core [--trace]".into()),
    };
    println!("{}", serde_json::to_string_pretty(&output)?);
    Ok(())
}
  1. Inspect the reference census. Explain why its tied head gives 1,188 learned scalars and how a hypothetical separate head would change that count.
  2. For the laptop plan, hold every model field fixed except key/value heads, changing two to one. Reproduce the changed key/value width, parameter count and two-byte KV capacity. Identify the categories that stay unchanged.
  3. Reproduce the byte excess when production KV alone is compared with the laptop device envelope. Explain why this is enough to refuse that allocation.
  4. Inspect the two synthetic lifetime schedules. Explain the 288-byte and 224-byte joint peaks and the effect of aliasing B.
  5. Reproduce TP8’s 14 sent chunks per rank/event and 160 forward events. Distinguish the sent-byte totals for one rank and all ranks.
  6. Inspect a run-seed or token-budget change, then the explicit laptop capacity change from 512 to 128. Explain which identity changes and why the learned parameter census stays fixed.
  7. Explain admission for a 100-byte limit with 64 bytes of known state and Unknown required workspace. Relate this case to the distinction between complete symbolic accounting and actual GPU readiness.
Checked explanations for tied storage, KV changes, refusal, lifetimes, traffic and configuration identities
  1. The reference embedding contains 10641064 learned scalars, attention matrices 6464, feed-forward matrices 4848 and norm gains 1212: 1064+64+48+12=11881064+64+48+12=1188 in total. The transpose head aliases the embedding. An independent head would add 10641064, giving 22522252 and changing the architecture.
  2. For the laptop plan, head width stays 512/8=64512/8=64 features; key/value width changes from 2⋅64=1282\cdot64=128 to 6464. K/V projection scalars fall by 2⋅8⋅512⋅64=5242882\cdot8\cdot512\cdot64=524288, giving 3199027231990272 unique parameters. Two-byte KV capacity halves from 20971522097152 to 10485761048576 bytes. Q/output, embedding, feed-forward and norm scalars stay fixed; the whole model and all MAC terms do not halve.
  3. Production KV payload minus the laptop device envelope is 10737418240−6710886400=402653184010737418240-6710886400=4026531840 bytes. KV capacity alone exceeds the envelope, before weights or activations. Refusal precedes allocation; no out-of-memory attempt or assumption about unknown workspace is needed.
  4. The synthetic overlap schedule keeps state with 6464 bytes, A with 6464, B with 128128 and B’s workspace with 3232: 64+64+128+32=28864+64+128+32=288 bytes together. Releasing A before B runs gives 64+128+32=22464+128+32=224 bytes. A B alias adds zero owned bytes. Peak accounting takes simultaneous lifetimes and workspace; adding independently maximized categories can misstate that relationship.
  5. In this TP8 forward ring fixture, reduce-scatter and all-gather each send 8−18-1 chunks per rank, giving 1414 per event. Two events per block over 8080 blocks give 160160. Each message has 3355443233554432 elements at two bytes each. The totals are 150,323,855,360150{,}323{,}855{,}360 sent bytes per rank and 1,202,590,842,8801{,}202{,}590{,}842{,}880 summed across eight ranks; they are calculated traffic, not resident memory or measured bandwidth.
  6. A run seed or training-token budget NN belongs to the run record: changing it changes that digest, while the semantic digest and learned census PP stay fixed. Position capacity belongs to the semantic record: changing 512512 to 128128 token positions per row changes its digest. RoPE owns no learned position table here, so the change leaves P=32514560P=32514560. The smoke record is explicit rather than a clamped laptop input.
  7. A 100100-byte limit with 6464 bytes of state leaves 3636 bytes, but that remainder does not tell us the required workspace. Preserve Unknown and refuse until an actual backend descriptor supplies a defensible demand or enforced bound. A complete symbolic schedule can supply a calculation under its assumptions while actual GPU readiness remains unestablished.

Carry the checked core into stochastic execution

One checked configuration now connects dimension choices to planning and admitted construction. The shared text path preserves the occurrence-local attention, positions and valid-target loss from Chapter 5 and the construction policy from Chapter 6. Chapter 8 adds optional dropout with its own reproducibility streams. Unequal-head GQA execution, GPU and distributed backends, adapter selection and trained-checkpoint serialization remain later capabilities.