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 , residual width , one complete block , one query head and one key/value head, and feed-forward width . The model can represent positions up to its configured capacity .
The embedding table contains learned scalars. Its transpose also supplies the vocabulary head, so that head owns no second table. The four attention matrices contain scalars. The three SwiGLU matrices contain . Two block RMSNorm gains and one final gain contain . Together they give
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 , producing parameters and a different model architecture.
Derive shapes before counting parameters
Let count query heads and count key/value heads. Each head has width . The total key or value projection width is :
The divisions must be exact: divides , and divides , with . The selected RoPE policy also requires even . 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 , key and value contribute , the three feed-forward matrices contribute , and two norm gains contribute . Call their sum . The embedding and final gain sit outside the blocks:
counts unique learned scalar values. RoPE adds no learned position table, and the tied head adds no independent vocabulary matrix. This is why changing alone leaves unchanged for this family.
| Model quantity | Reference | Bridge | Laptop plan | Production plan |
|---|---|---|---|---|
| Vocabulary entries | 266 | 266 | 16384 | 128000 |
| Residual features | 4 | 16 | 512 | 8192 |
| Complete blocks | 1 | 2 | 8 | 80 |
| Query heads | 1 | 2 | 8 | 64 |
| Key/value heads | 1 | 2 | 2 | 8 |
| SwiGLU hidden features | 4 | 20 | 1536 | 28672 |
| Position capacity | 4 | 16 | 512 | 32768 |
| Unique learned scalars | 1188 | 8304 | 32514560 | 69500936192 |
| Current execution | CPU supported | CPU supported | Refused | Refused |
Keep dimensions, capacity and token budgets distinct
counts vocabulary entries; counts residual features per token; counts complete blocks; and count heads; and counts SwiGLU hidden features. The derived widths and count features, while and count learned scalars per block and per model.
counts available token positions per row. is the nominal row count in one full-context forward accounting unit. A particular call may use fewer admitted rows or positions. is a declared training-token budget, separate from these tensor axes. Reference uses , bridge , laptop , and production . These budgets do not mean the chapter trains on those tokens.
Later resource formulas use 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 to changes semantic identity while the parameter count stays .
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.
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:
- Parse and validate the two records.
- Calculate the named tensor, state, activation, KV and communication plans.
- Compare execution support and every applicable profile limit.
- 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 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.
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 , is the actual number of rows and 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.
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.
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, ¤t, 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,
¤t,
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,
¤t,
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, ¤t, ¶meters, 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, ¤t)
.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, ¤t, 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, ¤t)
.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 .
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 plus relative tolerance 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 bytes per element, the payload capacity is
The factor two counts keys and values. The formula uses key/value heads through , 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
For ordinary unsegmented causal attention, permitted triangular score and value-mix entries add . Materializing the full dense score and mix products adds instead:
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.
| Scale | Nominal rows | KV payload bytes at two bytes per element | Causal MAC pairs | Dense MAC pairs |
|---|---|---|---|---|
| Reference | 16 | 1024 | 76544 | 77312 |
| Bridge | 8 | 16384 | 1122304 | 1183744 |
| Laptop plan | 1 | 2097152 | 17718837248 | 18790481920 |
| Production plan | 1 | 10737418240 | 2981072375644160 | 3684738342584320 |
The two-byte column is a common accounting comparison. Replacing its element
width with eight-byte f64 gives and 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 and bytes at capacities four and sixteen,
before overhead. The small executed comparisons also use one row. Production’s two-byte KV
alone is bytes, already greater than the laptop device
envelope of 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: bytes
before alignment. The production subtotal is
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
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 bytes of state, allocate A with bytes, then B with bytes while B needs bytes of workspace. Their joint peak is bytes. Release A before allocating B and the peak falls to 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 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 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 elements. Eight ring chunks contain elements each. Reduce-scatter plus all-gather sends chunks per rank per event. Two events in each of blocks give events. At two bytes per element, this produces sent bytes per rank per event, per rank, and 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 leaves the model’s 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.
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-coreTo 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 -- --tracerust/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(())
} - Inspect the reference census. Explain why its tied head gives 1,188 learned scalars and how a hypothetical separate head would change that count.
- 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.
- Reproduce the byte excess when production KV alone is compared with the laptop device envelope. Explain why this is enough to refuse that allocation.
- Inspect the two synthetic lifetime schedules. Explain the 288-byte and 224-byte joint peaks and the effect of aliasing B.
- Reproduce TP8’s 14 sent chunks per rank/event and 160 forward events. Distinguish the sent-byte totals for one rank and all ranks.
- 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.
- 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
- The reference embedding contains learned scalars, attention matrices , feed-forward matrices and norm gains : in total. The transpose head aliases the embedding. An independent head would add , giving and changing the architecture.
- For the laptop plan, head width stays features; key/value width changes from to . K/V projection scalars fall by , giving unique parameters. Two-byte KV capacity halves from to bytes. Q/output, embedding, feed-forward and norm scalars stay fixed; the whole model and all MAC terms do not halve.
- Production KV payload minus the laptop device envelope is 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.
- The synthetic overlap schedule keeps state with bytes, A with , B with and B’s workspace with : bytes together. Releasing A before B runs gives bytes. A B alias adds zero owned bytes. Peak accounting takes simultaneous lifetimes and workspace; adding independently maximized categories can misstate that relationship.
- In this TP8 forward ring fixture, reduce-scatter and all-gather each send chunks per rank, giving per event. Two events per block over blocks give . Each message has elements at two bytes each. The totals are sent bytes per rank and summed across eight ranks; they are calculated traffic, not resident memory or measured bandwidth.
- A run seed or training-token budget belongs to the run record: changing it changes that digest, while the semantic digest and learned census stay fixed. Position capacity belongs to the semantic record: changing to token positions per row changes its digest. RoPE owns no learned position table here, so the change leaves . The smoke record is explicit rather than a clamped laptop input.
- A -byte limit with bytes of state leaves bytes, but that remainder does not tell us the required workspace. Preserve
Unknownand 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.