38 · Content revision 6
Prefill once, then advance one token
Learn how one KV cache per decoder block and one checked model/cache session support prompt prefill followed by coherent one-token decoding, then compare newest-position logits and generation decisions with complete-prefix references.
Predict both cache lengths before running
The fixture has decoder blocks, attention heads per block, model width , head width , and context capacity . It begins with prompt token IDs and an empty model-wide cache.
Before revealing the trace, predict four facts:
- Prefill sends both prompt positions through both blocks, so each block cache reaches logical length and shape .
- The first later token is token ID at absolute position .
- A successful decode advances both caches from length to length ; one block cannot remain behind the other.
- The selected EOS token is returned to the caller but is not decoded when no later logits are needed.
The final prompt logits are . After token , the next logits are . After prompt prefill and after decoding token , the cached and complete-prefix results have an unrounded maximum absolute difference no greater than ; the displayed difference is .
The restored checkpoint has context capacity . It uses prompt and RNG
state 0x9e3779b97f4a7c38. Both generation paths select . Converting
those generated token IDs back to text produces the literal string 44. The
paths consume matching sampling draws, finish with the same RNG state, and
stop at the context limit. Cached prefill reaches length ; the first is
decoded and advances the cache to length ; the second is selected from
those logits and returned, then context-limit stops before decoding it. With
token configured as EOS, both paths stop after selecting and perform
later decode-token forwards.
Count attention-score cells, not total runtime
Hold batch size, layer count, head count, and head width fixed. Over a final retained length , the two score-cell totals grow as
At retained length , complete-prefix replay rebuilds a dense causal score grid with entries for every fixed batch-layer-head lane. A cached step forms only the newest query’s row of scores against all retained keys. Summing those per-step quantities from through gives the two growth classes.
This comparison is deliberately narrow. Cached attention is not constant-time: the newest query still scans a prefix that grows with . The formula does not count projection or MLP work, memory traffic, allocation behavior, total runtime, or wall-clock speedup.
For the exact fixture, batch , blocks, and heads give a fixed factor of . Serial prefill plus decode forms cached attention-score values. The two independent reference calls at lengths and form values.
These fixture schedules are intentionally different. The cached execution visits retained lengths , , and , while the chapter makes only two independent complete-prefix checks, at lengths and . Thus and are measured counts for the actual calls, not both asymptotic sums evaluated at .
Keep retained length and final length separate
- is the current retained prefix length. It is also the number of keys read by the one newest cached query.
- is the final retained length covered by the comparison.
- is the complete-prefix score grid rebuilt at length for each fixed batch-layer-head lane.
- accumulates score-cell work from retained length through retained length .
- is the growth class of repeated complete-prefix score grids when the omitted factors remain fixed.
- is the growth class of cached newest-query score rows under those same fixed factors.
The cache itself still stores one logical K/V prefix per block with shape
. Equal shapes do not make two caches interchangeable: their rows
were produced at different depths of the decoder stack. DecoderKvCache records
the decoder configuration plus every parameter’s node identity and value
revision. Each nested layer cache separately records its attention geometry,
RoPE configuration, and four attention-parameter bindings. These records are
compatibility evidence, not a live relationship by themselves.
Calling cache.bind(&model) checks the evidence and creates a
DecoderKvSession for that exact model/cache pair. The session then keeps the
parameter values borrowed for reading. Shape and node identity alone are not
enough: an in-place update can preserve the nodes while changing the value
revisions, and rows computed by two parameter revisions cannot share one logical
prefix.
From causal stacks to prompt and decode phases
A causal Transformer decoder can generate one token at a time by replaying the complete known prefix, but that stateless interface rebuilds earlier attention score grids and key/value projections on every later call.
Attention Is All You Need establishes the causal stack. Vaswani and colleagues describe a stacked autoregressive Transformer decoder whose masked self-attention prevents a position from reading later positions; their encoder-decoder architecture also includes cross-attention and does not specify a KV-cache API.
Fast Transformer Decoding: One Write-Head is All You Need makes previous state explicit. Shazeer’s incremental self-attention receives previous key and value tensors, appends the current projected key and value, and returns the updated state; the paper’s contribution is multi-query attention, not a claim to have invented KV caching.
Efficient Memory Management for Large Language Model Serving with PagedAttention continues the road into modern LLM inference. Kwon and colleagues separate a prompt phase from sequential generation, describe later iterations reusing cached keys and values while computing only the newest pair, and account for KV-cache state across Transformer layers and heads.
Incremental decoding made previous key/value tensors explicit state, and later LLM serving work separated one prompt phase from sequential generation while retaining key/value state across decoder layers and heads.
Model-wide cached generation prefills one independent cache per decoder block and advances every block only when later logits are needed. In the exact fixtures, newest-position logits agree with complete-prefix references within tolerance, and restored cached generation matches selected tokens, sampling draws, final RNG state, and stopping reason.
The executable connects that progression to measured work. Across four batch-layer-head lanes, serial cached rows at retained lengths contain attention-score values. The two complete-prefix calls at lengths contain , so this exact schedule avoids score values. These are tensor element counts, not a paging or runtime measurement.
To keep that reusable state coherent, this implementation checks the complete decoder/cache relationship when a session is created, retains read-only access to all parameter values for the session, prepares every block before committing any cache row, and makes reset, counters, typed errors, and serial prompt processing explicit. These are requirements of this implementation, not policies stated by the cited papers.
Paging, sharing, eviction, and production memory management remain outside this chapter. The contiguous cache here isolates the correctness boundary needed before those serving concerns can be studied.
rust/demos/ch38-cached-generation/src/lib.rs#historical-cache-contrast /// Measures complete-prefix replay against retained model-wide KV state.
pub fn historical_cache_contrast(
config: DecoderModelConfig,
cached_retained_lengths: &[usize],
complete_prefix_lengths: &[usize],
measured_cached_scores: usize,
measured_complete_prefix_scores: usize,
) -> Result<HistoricalCacheContrast, FixtureError> {
require(
!cached_retained_lengths.is_empty() && !complete_prefix_lengths.is_empty(),
"history evidence needs cached and complete-prefix calls",
)?;
let batch_layer_head_lanes = config
.layers()
.checked_mul(config.heads())
.ok_or(FixtureError::Invariant("history lane count overflowed"))?;
let cached_attention_score_values =
cached_retained_lengths
.iter()
.try_fold(0usize, |total, &length| {
let scores =
batch_layer_head_lanes
.checked_mul(length)
.ok_or(FixtureError::Invariant(
"cached history score count overflowed",
))?;
total.checked_add(scores).ok_or(FixtureError::Invariant(
"cached history score total overflowed",
))
})?;
let complete_prefix_attention_score_values =
complete_prefix_attention_score_values(config, complete_prefix_lengths)?;
require(
cached_attention_score_values == measured_cached_scores
&& complete_prefix_attention_score_values == measured_complete_prefix_scores,
"history score contrast disagrees with measured work",
)?;
let avoided_attention_score_values = complete_prefix_attention_score_values
.checked_sub(cached_attention_score_values)
.ok_or(FixtureError::Invariant(
"cached history exceeds complete-prefix reference",
))?;
Ok(HistoricalCacheContrast {
batch_layer_head_lanes,
cached_retained_lengths: cached_retained_lengths.to_vec(),
cached_attention_score_values,
complete_prefix_lengths: complete_prefix_lengths.to_vec(),
complete_prefix_attention_score_values,
avoided_attention_score_values,
})
}Prepare every block, then commit every cache together
Chapter 37’s layer operation now has two stages. First it calculates the newest output plus two candidate rows—one rotated key row and one unrotated value row—without changing logical state. Its ordinary public call commits that prepared pair immediately. The model-wide path instead retains one prepared result per block and commits them only after every block and the final vocabulary projection succeed.
rust/crates/llm-from-scratch/src/attention/incremental.rs#incremental-attention /// A fallible incremental-attention buffer or tensor stage.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum IncrementalAttentionStage {
Scores,
HeadOutputs,
HeadOutputLeaf,
}
impl fmt::Display for IncrementalAttentionStage {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::Scores => "attention scores",
Self::HeadOutputs => "weighted head outputs",
Self::HeadOutputLeaf => "head-output tensor",
})
}
}
/// A rejected single-token input, cache pairing, or incremental forward stage.
#[derive(Clone, Debug, PartialEq)]
pub enum IncrementalAttentionError {
InputRank {
rank: usize,
},
SingleTokenRequired {
tokens: usize,
},
InputBatchMismatch {
cache: usize,
input: usize,
},
InputWidthMismatch {
expected: usize,
actual: usize,
},
CacheModelWidthMismatch {
layer: usize,
cache: usize,
},
CacheHeadCountMismatch {
layer: usize,
cache: usize,
},
CacheHeadWidthMismatch {
layer: usize,
cache: usize,
},
CacheLayerMismatch,
CacheLayerRevisionMismatch {
parameter: usize,
cache: u64,
layer: u64,
},
CacheRopeMismatch {
cache_feature_width: usize,
layer_feature_width: usize,
cache_max_positions: usize,
layer_max_positions: usize,
cache_base: f64,
layer_base: f64,
},
BatchHeadOverflow {
batch: usize,
heads: usize,
},
BufferSizeOverflow {
stage: IncrementalAttentionStage,
},
BufferAllocationFailed {
stage: IncrementalAttentionStage,
elements: usize,
},
Cache(LayerKvCacheError),
QkvProjection(QkvError),
HeadLayout {
input: MultiHeadInput,
source: HeadLayoutError,
},
Rotary {
input: MultiHeadInput,
source: RopeError,
},
Probability(ProbabilityError),
Tensor {
stage: IncrementalAttentionStage,
source: TensorError,
},
Autodiff {
stage: IncrementalAttentionStage,
source: TensorAutodiffError,
},
MergeLayout(HeadLayoutError),
OutputProjection(LinearError),
}
impl fmt::Display for IncrementalAttentionError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::InputRank { rank } => write!(
formatter,
"incremental attention input must have rank three [batch, 1, model_width], got rank {rank}"
),
Self::SingleTokenRequired { tokens } => write!(
formatter,
"incremental attention needs exactly one new token, got {tokens}"
),
Self::InputBatchMismatch { cache, input } => write!(
formatter,
"incremental attention input batch {input} must match cache batch {cache}"
),
Self::InputWidthMismatch { expected, actual } => write!(
formatter,
"incremental attention input width must equal model width {expected}, got {actual}"
),
Self::CacheModelWidthMismatch { layer, cache } => write!(
formatter,
"KV cache model width {cache} must match attention layer width {layer}"
),
Self::CacheHeadCountMismatch { layer, cache } => write!(
formatter,
"KV cache head count {cache} must match attention layer head count {layer}"
),
Self::CacheHeadWidthMismatch { layer, cache } => write!(
formatter,
"KV cache head width {cache} must match attention layer head width {layer}"
),
Self::CacheLayerMismatch => formatter
.write_str("KV cache parameter identity does not match this attention layer"),
Self::CacheLayerRevisionMismatch {
parameter,
cache,
layer,
} => write!(
formatter,
"KV cache parameter revision {cache} at stable index {parameter} does not match layer revision {layer}"
),
Self::CacheRopeMismatch {
cache_feature_width,
layer_feature_width,
cache_max_positions,
layer_max_positions,
cache_base,
layer_base,
} => write!(
formatter,
"KV cache RoPE configuration ({cache_feature_width} features, {cache_max_positions} positions, base {cache_base:?}) does not match layer configuration ({layer_feature_width} features, {layer_max_positions} positions, base {layer_base:?})"
),
Self::BatchHeadOverflow { batch, heads } => write!(
formatter,
"incremental attention lane count overflows for batch {batch} and {heads} heads"
),
Self::BufferSizeOverflow { stage } => {
write!(formatter, "incremental {stage} element count overflows")
}
Self::BufferAllocationFailed { stage, elements } => write!(
formatter,
"cannot allocate incremental {stage} buffer for {elements} f64 values"
),
Self::Cache(source) => source.fmt(formatter),
Self::QkvProjection(source) => {
write!(formatter, "incremental Q/K/V projection: {source}")
}
Self::HeadLayout { input, source } => {
write!(formatter, "incremental {input} head layout: {source}")
}
Self::Rotary { input, source } => {
write!(formatter, "incremental {input} RoPE: {source}")
}
Self::Probability(source) => write!(formatter, "incremental softmax: {source}"),
Self::Tensor { stage, source } => {
write!(formatter, "incremental {stage}: {source}")
}
Self::Autodiff { stage, source } => {
write!(formatter, "incremental {stage}: {source}")
}
Self::MergeLayout(source) => {
write!(formatter, "incremental head output merge: {source}")
}
Self::OutputProjection(source) => {
write!(formatter, "incremental output projection: {source}")
}
}
}
}
impl Error for IncrementalAttentionError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::Cache(source) => Some(source),
Self::QkvProjection(source) => Some(source),
Self::HeadLayout { source, .. } => Some(source),
Self::Rotary { source, .. } => Some(source),
Self::Probability(source) => Some(source),
Self::Tensor { source, .. } => Some(source),
Self::Autodiff { source, .. } => Some(source),
Self::MergeLayout(source) => Some(source),
Self::OutputProjection(source) => Some(source),
_ => None,
}
}
}
impl From<LayerKvCacheError> for IncrementalAttentionError {
fn from(source: LayerKvCacheError) -> Self {
Self::Cache(source)
}
}
/// Exact row counts for comparing full-prefix and cached projections.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct IncrementalAttentionWork {
position: usize,
full_prefix_rows_per_projection: usize,
incremental_rows_per_projection: usize,
reused_key_value_rows: usize,
}
impl IncrementalAttentionWork {
pub const fn position(&self) -> usize {
self.position
}
pub const fn full_prefix_rows_per_projection(&self) -> usize {
self.full_prefix_rows_per_projection
}
pub const fn incremental_rows_per_projection(&self) -> usize {
self.incremental_rows_per_projection
}
pub const fn reused_key_value_rows(&self) -> usize {
self.reused_key_value_rows
}
}
/// Inspectable graph-free evidence from one committed cache append.
#[derive(Clone, Debug)]
pub struct IncrementalAttentionForward {
projected_query_heads: TensorValue,
projected_key_heads: TensorValue,
projected_value_heads: TensorValue,
rotated_query_heads: TensorValue,
rotated_key_heads: TensorValue,
attention_weights: Tensor,
head_outputs: TensorValue,
merged: TensorValue,
output: TensorValue,
work: IncrementalAttentionWork,
cache_len: usize,
}
impl IncrementalAttentionForward {
pub fn projected_query_heads(&self) -> &TensorValue {
&self.projected_query_heads
}
pub fn projected_key_heads(&self) -> &TensorValue {
&self.projected_key_heads
}
pub fn projected_value_heads(&self) -> &TensorValue {
&self.projected_value_heads
}
pub fn rotated_query_heads(&self) -> &TensorValue {
&self.rotated_query_heads
}
pub fn rotated_key_heads(&self) -> &TensorValue {
&self.rotated_key_heads
}
/// Returns `[batch, heads, 1, cache_len]` probabilities for the new query.
pub fn attention_weights(&self) -> &Tensor {
&self.attention_weights
}
pub fn head_outputs(&self) -> &TensorValue {
&self.head_outputs
}
pub fn merged(&self) -> &TensorValue {
&self.merged
}
pub fn output(&self) -> &TensorValue {
&self.output
}
pub const fn work(&self) -> IncrementalAttentionWork {
self.work
}
pub const fn cache_len(&self) -> usize {
self.cache_len
}
pub fn into_output(self) -> TensorValue {
self.output
}
}
/// A crate-sealed incremental result whose candidate K/V row is not committed.
///
/// Chapter 38 prepares one ticket per decoder block, completes the later blocks
/// and tied vocabulary head, verifies every ticket still targets its original
/// cache, and only then commits the complete stack.
pub(crate) struct PreparedIncrementalAttention {
forward: IncrementalAttentionForward,
candidate_key: Tensor,
candidate_value: Tensor,
expected_len: usize,
key_storage: *const f64,
value_storage: *const f64,
}
impl PreparedIncrementalAttention {
pub(crate) fn output(&self) -> &TensorValue {
self.forward.output()
}
pub(crate) const fn cache_len(&self) -> usize {
self.forward.cache_len()
}
pub(crate) fn attention_score_values(&self) -> usize {
self.forward.attention_weights().len()
}
pub(crate) fn matches_cache(&self, cache: &LayerKvCache) -> bool {
self.expected_len == cache.len()
&& std::ptr::eq(self.key_storage, cache.key_storage().as_ptr())
&& std::ptr::eq(self.value_storage, cache.value_storage().as_ptr())
}
pub(crate) fn commit(self, cache: &mut LayerKvCache) -> IncrementalAttentionForward {
debug_assert!(self.matches_cache(cache));
cache.append_prevalidated(&self.candidate_key, &self.candidate_value);
self.forward
}
}
impl MultiHeadAttention {
/// Projects one new row, attends over retained K/V rows, and commits one append.
///
/// The cache is changed only after projection, RoPE, stable softmax, value
/// mixing, merge, and output projection all succeed.
pub fn forward_incremental(
&self,
input: &TensorValue,
cache: &mut LayerKvCache,
) -> Result<IncrementalAttentionForward, IncrementalAttentionError> {
let prepared = self.prepare_incremental(input, cache)?;
Ok(prepared.commit(cache))
}
/// Computes one incremental row without changing the layer cache.
pub(crate) fn prepare_incremental(
&self,
input: &TensorValue,
cache: &LayerKvCache,
) -> Result<PreparedIncrementalAttention, IncrementalAttentionError> {
self.validate_incremental_request(input, cache)?;
self.prepare_incremental_bound(input, cache)
}
fn validate_incremental_request(
&self,
input: &TensorValue,
cache: &LayerKvCache,
) -> Result<(), IncrementalAttentionError> {
let shape = input.shape();
if shape.len() != 3 {
return Err(IncrementalAttentionError::InputRank { rank: shape.len() });
}
if shape[1] != 1 {
return Err(IncrementalAttentionError::SingleTokenRequired { tokens: shape[1] });
}
if shape[0] != cache.batch_size() {
return Err(IncrementalAttentionError::InputBatchMismatch {
cache: cache.batch_size(),
input: shape[0],
});
}
if shape[2] != self.model_width() {
return Err(IncrementalAttentionError::InputWidthMismatch {
expected: self.model_width(),
actual: shape[2],
});
}
self.validate_incremental_cache_binding(cache)?;
if cache.is_full() {
return Err(LayerKvCacheError::Full {
capacity: cache.capacity(),
}
.into());
}
Ok(())
}
/// Checks the persistent relationship between one layer and one cache.
///
/// A model-wide session calls this once while binding its complete cache.
/// The standalone entry calls it for every arbitrary layer/cache pairing.
pub(crate) fn validate_incremental_cache_binding(
&self,
cache: &LayerKvCache,
) -> Result<(), IncrementalAttentionError> {
if cache.model_width() != self.model_width() {
return Err(IncrementalAttentionError::CacheModelWidthMismatch {
layer: self.model_width(),
cache: cache.model_width(),
});
}
if cache.heads() != self.heads() {
return Err(IncrementalAttentionError::CacheHeadCountMismatch {
layer: self.heads(),
cache: cache.heads(),
});
}
if cache.head_width() != self.head_width() {
return Err(IncrementalAttentionError::CacheHeadWidthMismatch {
layer: self.head_width(),
cache: cache.head_width(),
});
}
if !self
.parameters()
.iter()
.zip(&cache.parameter_bindings)
.all(|(parameter, cached)| cached.node_matches(parameter.tensor()))
{
return Err(IncrementalAttentionError::CacheLayerMismatch);
}
if let Some((parameter, (cached, layer))) = cache
.parameter_bindings
.iter()
.zip(self.parameters())
.enumerate()
.find(|(_, (cached, parameter))| !cached.revision_matches(parameter.tensor()))
{
return Err(IncrementalAttentionError::CacheLayerRevisionMismatch {
parameter,
cache: cached.revision(),
layer: layer.tensor().value_revision(),
});
}
if cache.rope_feature_width != self.rope().feature_width()
|| cache.rope_max_positions != self.rope().max_positions()
|| cache.rope_base_bits != self.rope().base().to_bits()
{
return Err(IncrementalAttentionError::CacheRopeMismatch {
cache_feature_width: cache.rope_feature_width,
layer_feature_width: self.rope().feature_width(),
cache_max_positions: cache.rope_max_positions,
layer_max_positions: self.rope().max_positions(),
cache_base: f64::from_bits(cache.rope_base_bits),
layer_base: self.rope().base(),
});
}
Ok(())
}
/// Prepares one row after its crate-private caller establishes every precondition.
///
/// A model-wide bind establishes the persistent layer/cache relationship.
/// The current session operation separately guarantees the one-row input
/// shape and remaining capacity. The caller must preserve that exact
/// layer/cache pairing until every prepared row either commits or is
/// discarded. Keeping this entry crate-private lets Chapter 38 reuse the one
/// attention implementation without creating an unchecked public path.
pub(crate) fn prepare_incremental_bound(
&self,
input: &TensorValue,
cache: &LayerKvCache,
) -> Result<PreparedIncrementalAttention, IncrementalAttentionError> {
no_grad(|| {
let position = cache.len();
let projected = self
.qkv()
.forward(input)
.map_err(IncrementalAttentionError::QkvProjection)?;
let projected_query_heads =
split_heads(projected.query(), self.heads()).map_err(|source| {
IncrementalAttentionError::HeadLayout {
input: MultiHeadInput::Query,
source,
}
})?;
let projected_key_heads =
split_heads(projected.key(), self.heads()).map_err(|source| {
IncrementalAttentionError::HeadLayout {
input: MultiHeadInput::Key,
source,
}
})?;
let projected_value_heads =
split_heads(projected.value(), self.heads()).map_err(|source| {
IncrementalAttentionError::HeadLayout {
input: MultiHeadInput::Value,
source,
}
})?;
let rotated_query_heads = self
.rope()
.rotate(&projected_query_heads, position)
.map_err(|source| IncrementalAttentionError::Rotary {
input: MultiHeadInput::Query,
source,
})?;
let rotated_key_heads =
self.rope()
.rotate(&projected_key_heads, position)
.map_err(|source| IncrementalAttentionError::Rotary {
input: MultiHeadInput::Key,
source,
})?;
let candidate_key = rotated_key_heads.value_snapshot();
let candidate_value = projected_value_heads.value_snapshot();
let (attention_weights, head_output_tensor) = incremental_mixture(
&rotated_query_heads.value(),
&candidate_key,
&candidate_value,
cache,
)?;
let head_outputs = TensorValue::constant(head_output_tensor).map_err(|source| {
IncrementalAttentionError::Autodiff {
stage: IncrementalAttentionStage::HeadOutputLeaf,
source,
}
})?;
let merged =
merge_heads(&head_outputs).map_err(IncrementalAttentionError::MergeLayout)?;
let output = self
.output_projection()
.forward(&merged)
.map_err(IncrementalAttentionError::OutputProjection)?;
let cache_len = position + 1;
let result = IncrementalAttentionForward {
projected_query_heads,
projected_key_heads,
projected_value_heads,
rotated_query_heads,
rotated_key_heads,
attention_weights,
head_outputs,
merged,
output,
work: IncrementalAttentionWork {
position,
full_prefix_rows_per_projection: cache_len,
incremental_rows_per_projection: 1,
reused_key_value_rows: position,
},
cache_len,
};
cache.validate_append(&candidate_key, &candidate_value)?;
Ok(PreparedIncrementalAttention {
forward: result,
candidate_key,
candidate_value,
expected_len: position,
key_storage: cache.key_storage().as_ptr(),
value_storage: cache.value_storage().as_ptr(),
})
})
}
}
fn incremental_mixture(
query: &Tensor,
candidate_key: &Tensor,
candidate_value: &Tensor,
cache: &LayerKvCache,
) -> Result<(Tensor, Tensor), IncrementalAttentionError> {
let lanes = cache.batch_size().checked_mul(cache.heads()).ok_or(
IncrementalAttentionError::BatchHeadOverflow {
batch: cache.batch_size(),
heads: cache.heads(),
},
)?;
let prefix = cache.len() + 1;
let score_elements =
lanes
.checked_mul(prefix)
.ok_or(IncrementalAttentionError::BufferSizeOverflow {
stage: IncrementalAttentionStage::Scores,
})?;
let mut scores = reserved_buffer(score_elements, IncrementalAttentionStage::Scores)?;
let scale = 1.0 / (cache.head_width() as f64).sqrt();
for batch in 0..cache.batch_size() {
for head in 0..cache.heads() {
let lane = batch * cache.heads() + head;
let query_start = lane * cache.head_width();
for position in 0..prefix {
let key_start = if position == cache.len() {
lane * cache.head_width()
} else {
(lane * cache.capacity() + position) * cache.head_width()
};
let key_values = if position == cache.len() {
candidate_key.as_slice()
} else {
cache.key_storage()
};
let mut dot = 0.0;
for feature in 0..cache.head_width() {
dot +=
query.as_slice()[query_start + feature] * key_values[key_start + feature];
}
scores[lane * prefix + position] = dot * scale;
}
}
}
let score_tensor = Tensor::from_vec(vec![cache.batch_size(), cache.heads(), 1, prefix], scores)
.map_err(|source| IncrementalAttentionError::Tensor {
stage: IncrementalAttentionStage::Scores,
source,
})?;
let weights =
softmax(&score_tensor.view(), 3).map_err(IncrementalAttentionError::Probability)?;
let output_elements = lanes.checked_mul(cache.head_width()).ok_or(
IncrementalAttentionError::BufferSizeOverflow {
stage: IncrementalAttentionStage::HeadOutputs,
},
)?;
let mut outputs = reserved_buffer(output_elements, IncrementalAttentionStage::HeadOutputs)?;
for batch in 0..cache.batch_size() {
for head in 0..cache.heads() {
let lane = batch * cache.heads() + head;
for feature in 0..cache.head_width() {
let mut mixture = 0.0;
for position in 0..prefix {
let value_start = if position == cache.len() {
lane * cache.head_width()
} else {
(lane * cache.capacity() + position) * cache.head_width()
};
let value_values = if position == cache.len() {
candidate_value.as_slice()
} else {
cache.value_storage()
};
mixture += weights.as_slice()[lane * prefix + position]
* value_values[value_start + feature];
}
outputs[lane * cache.head_width() + feature] = mixture;
}
}
}
let outputs = Tensor::from_vec(
vec![cache.batch_size(), cache.heads(), 1, cache.head_width()],
outputs,
)
.map_err(|source| IncrementalAttentionError::Tensor {
stage: IncrementalAttentionStage::HeadOutputs,
source,
})?;
Ok((weights, outputs))
}
fn reserved_buffer(
elements: usize,
stage: IncrementalAttentionStage,
) -> Result<Vec<f64>, IncrementalAttentionError> {
let mut values = Vec::new();
values
.try_reserve_exact(elements)
.map_err(|_| IncrementalAttentionError::BufferAllocationFailed { stage, elements })?;
values.resize(elements, 0.0);
Ok(values)
}DecoderKvCache::new takes the actual decoder. It allocates one fixed-capacity
LayerKvCache per block and records the exact decoder configuration. For every
model parameter, it captures both the node identity and current value revision.
The cache owns reusable K/V storage and compatibility evidence, but it neither
copies weights nor keeps the decoder borrowed.
Before inference, cache.bind(&model) checks that evidence and returns a
DecoderKvSession for one exact model/cache pair. It checks the decoder
configuration; the parameter count and ordered node identities and value revisions;
the number and common logical length of the layer caches; and every layer cache’s
batch size, capacity, attention geometry, attention-parameter bindings, and RoPE
configuration. “Once” here means once for each newly created session, not once
for the cache’s entire lifetime.
After those checks pass, the session retains live read-only borrows of every parameter value. These borrows are not copied snapshots. Decoder operations can keep reading the original values, while AdamW cannot obtain the exclusive write access needed to change them. This closes the gap between checking a value revision and later using the corresponding value.
The session is also the only mutable borrower of the cache. The session’s private
row transition and its public reset method both change every layer length
together, so the length-coherence invariant checked at bind remains true.
DecoderKvSession::prefill and DecoderKvSession::decode do not accept a model
argument because the session already holds the model they must use. They do not
run without a model, and they do not skip all checks. Each operation still checks
the facts that can change from call to call: prompt validity, phase, token domain,
remaining capacity, checked counter arithmetic, and whether each prepared ticket
still names the same K/V storage at the expected logical length. The session avoids rescanning the
model-wide relationship before every operation and rechecking stable layer/cache
facts for every row. Prefill accepts a validated nonempty
prompt only while state is empty. This teaching implementation sends prompt rows
serially through the same one-row path; it is not an optimized parallel prefill
implementation.
For each row, the decoder computes the embedding, then applies each block’s pre-norm attention and feed-forward residual paths in order, followed by final RMSNorm and the tied vocabulary projection. Every block uses Chapter 37’s one shared attention calculation to prepare two candidate rows—one rotated key row and one unrotated value row—without changing the cache. The fully checked Chapter 37 entry handles a standalone layer/cache call. The session uses a crate-private entry only after binding has established the persistent relationship; it is not a second attention algorithm or a public unchecked shortcut. Only after the remaining blocks, final normalization, and vocabulary projection succeed do the layer caches, common length, phase counts, and score-cell counts advance together.
rust/crates/llm-from-scratch/src/generation/kv_cache.rs#decoder-kv-cache /// A checked model-wide work counter that could not be represented as `usize`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DecoderKvCacheCounter {
TokenForwards,
PrefillTokens,
DecodeTokens,
CacheAppends,
QkvProjectionRows,
AttentionScoreValues,
CompletePrefixAttentionScoreValues,
}
impl fmt::Display for DecoderKvCacheCounter {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::TokenForwards => "token forwards",
Self::PrefillTokens => "prefill tokens",
Self::DecodeTokens => "decode tokens",
Self::CacheAppends => "cache appends",
Self::QkvProjectionRows => "Q/K/V projection rows",
Self::AttentionScoreValues => "cached attention score values",
Self::CompletePrefixAttentionScoreValues => "complete-prefix attention score values",
})
}
}
/// A graph-free cached-decoder stage that rejected a request or computation.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CachedDecoderStage {
TiedWeightTranspose,
TiedVocabularyProjection,
}
impl fmt::Display for CachedDecoderStage {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter.write_str(match self {
Self::TiedWeightTranspose => "transpose tied embedding weight",
Self::TiedVocabularyProjection => "project tied vocabulary logits",
})
}
}
/// A rejected model/cache pairing, phase transition, token, or decoder stage.
#[derive(Clone, Debug, PartialEq)]
pub enum DecoderKvCacheError {
LayerAllocationFailed {
layers: usize,
},
ParameterAllocationFailed {
parameters: usize,
},
LayerCache {
layer: usize,
source: LayerKvCacheError,
},
ModelConfigMismatch,
ModelParameterCountMismatch {
cache: usize,
model: usize,
},
ModelParameterMismatch {
index: usize,
},
ModelParameterRevisionMismatch {
index: usize,
cache: u64,
model: u64,
},
LayerCountMismatch {
cache: usize,
model: usize,
},
LayerBatchSizeMismatch {
layer: usize,
expected: usize,
actual: usize,
},
LayerCapacityMismatch {
layer: usize,
expected: usize,
actual: usize,
},
EmptyPrompt,
PromptTooLong {
tokens: usize,
capacity: usize,
},
PromptTokenOutOfBounds {
position: usize,
token_id: u32,
vocabulary_size: usize,
},
PrefillRequiresEmpty {
len: usize,
},
DecodeRequiresPrefill,
DecodeTokenOutOfBounds {
token_id: u32,
vocabulary_size: usize,
},
Full {
capacity: usize,
},
LayerLengthInvariant {
layer: usize,
expected: usize,
actual: usize,
},
PreparedCacheChanged {
layer: usize,
},
WorkOverflow {
counter: DecoderKvCacheCounter,
},
Embedding(EmbeddingError),
AttentionNorm {
layer: usize,
source: RmsNormError,
},
IncrementalAttention {
layer: usize,
source: IncrementalAttentionError,
},
AttentionResidual {
layer: usize,
source: ResidualError,
},
FeedForwardNorm {
layer: usize,
source: RmsNormError,
},
FeedForward {
layer: usize,
source: SwiGluError,
},
FeedForwardResidual {
layer: usize,
source: ResidualError,
},
FinalNorm(RmsNormError),
Autodiff {
stage: CachedDecoderStage,
source: TensorAutodiffError,
},
}
impl fmt::Display for DecoderKvCacheError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::LayerAllocationFailed { layers } => {
write!(formatter, "cannot allocate {layers} decoder layer caches")
}
Self::ParameterAllocationFailed { parameters } => write!(
formatter,
"cannot retain {parameters} decoder parameter bindings"
),
Self::LayerCache { layer, source } => {
write!(formatter, "decoder layer {layer} cache: {source}")
}
Self::ModelConfigMismatch => formatter
.write_str("decoder cache configuration does not match this decoder model exactly"),
Self::ModelParameterCountMismatch { cache, model } => write!(
formatter,
"decoder cache binds {cache} parameter nodes, but the model exposes {model}"
),
Self::ModelParameterMismatch { index } => write!(
formatter,
"decoder cache parameter identity differs at stable index {index}"
),
Self::ModelParameterRevisionMismatch {
index,
cache,
model,
} => write!(
formatter,
"decoder cache parameter revision {cache} at stable index {index} differs from model revision {model}"
),
Self::LayerCountMismatch { cache, model } => write!(
formatter,
"decoder cache owns {cache} layer caches, but the model exposes {model} blocks"
),
Self::LayerBatchSizeMismatch {
layer,
expected,
actual,
} => write!(
formatter,
"decoder layer {layer} cache batch size must be {expected}, got {actual}"
),
Self::LayerCapacityMismatch {
layer,
expected,
actual,
} => write!(
formatter,
"decoder layer {layer} cache capacity must be {expected}, got {actual}"
),
Self::EmptyPrompt => formatter.write_str("cached prefill needs a nonempty prompt"),
Self::PromptTooLong { tokens, capacity } => write!(
formatter,
"cached prefill has {tokens} prompt tokens, exceeding capacity {capacity}"
),
Self::PromptTokenOutOfBounds {
position,
token_id,
vocabulary_size,
} => write!(
formatter,
"cached prefill token {token_id} at position {position} is out of bounds for vocabulary {vocabulary_size}"
),
Self::PrefillRequiresEmpty { len } => write!(
formatter,
"cached prefill requires empty state, but the cache length is {len}"
),
Self::DecodeRequiresPrefill => {
formatter.write_str("cached decode requires one completed nonempty prompt prefill")
}
Self::DecodeTokenOutOfBounds {
token_id,
vocabulary_size,
} => write!(
formatter,
"cached decode token {token_id} is out of bounds for vocabulary {vocabulary_size}"
),
Self::Full { capacity } => {
write!(formatter, "decoder KV cache is full at capacity {capacity}")
}
Self::LayerLengthInvariant {
layer,
expected,
actual,
} => write!(
formatter,
"decoder layer {layer} cache length must be {expected}, got {actual}"
),
Self::PreparedCacheChanged { layer } => write!(
formatter,
"decoder layer {layer} cache changed after its row was prepared"
),
Self::WorkOverflow { counter } => {
write!(formatter, "decoder cache {counter} counter overflows")
}
Self::Embedding(source) => write!(formatter, "cached token embedding: {source}"),
Self::AttentionNorm { layer, source } => {
write!(
formatter,
"decoder layer {layer} attention RMSNorm: {source}"
)
}
Self::IncrementalAttention { layer, source } => write!(
formatter,
"decoder layer {layer} incremental attention: {source}"
),
Self::AttentionResidual { layer, source } => write!(
formatter,
"decoder layer {layer} attention residual merge: {source}"
),
Self::FeedForwardNorm { layer, source } => write!(
formatter,
"decoder layer {layer} feed-forward RMSNorm: {source}"
),
Self::FeedForward { layer, source } => {
write!(formatter, "decoder layer {layer} SwiGLU: {source}")
}
Self::FeedForwardResidual { layer, source } => write!(
formatter,
"decoder layer {layer} feed-forward residual merge: {source}"
),
Self::FinalNorm(source) => write!(formatter, "cached final RMSNorm: {source}"),
Self::Autodiff { stage, source } => write!(formatter, "cached {stage}: {source}"),
}
}
}
impl Error for DecoderKvCacheError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::LayerCache { source, .. } => Some(source),
Self::Embedding(source) => Some(source),
Self::AttentionNorm { source, .. } | Self::FeedForwardNorm { source, .. } => {
Some(source)
}
Self::IncrementalAttention { source, .. } => Some(source),
Self::AttentionResidual { source, .. } | Self::FeedForwardResidual { source, .. } => {
Some(source)
}
Self::FeedForward { source, .. } => Some(source),
Self::FinalNorm(source) => Some(source),
Self::Autodiff { source, .. } => Some(source),
_ => None,
}
}
}
/// Exact model work committed to one decoder KV state.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct DecoderKvCacheWork {
token_forwards: usize,
prefill_tokens: usize,
decode_tokens: usize,
cache_appends: usize,
qkv_projection_rows: usize,
attention_score_values: usize,
}
impl DecoderKvCacheWork {
pub const fn token_forwards(self) -> usize {
self.token_forwards
}
pub const fn prefill_tokens(self) -> usize {
self.prefill_tokens
}
pub const fn decode_tokens(self) -> usize {
self.decode_tokens
}
pub const fn cache_appends(self) -> usize {
self.cache_appends
}
pub const fn qkv_projection_rows(self) -> usize {
self.qkv_projection_rows
}
pub const fn attention_score_values(self) -> usize {
self.attention_score_values
}
}
/// The newest graph-free vocabulary logits after one committed cached row.
#[derive(Clone, Debug)]
pub struct CachedDecoderOutput {
logits: TensorValue,
position: usize,
cache_len: usize,
attention_score_values: usize,
}
impl CachedDecoderOutput {
/// Returns logits shaped `[1, 1, vocabulary_size]`.
pub fn logits(&self) -> &TensorValue {
&self.logits
}
pub const fn position(&self) -> usize {
self.position
}
pub const fn cache_len(&self) -> usize {
self.cache_len
}
pub const fn attention_score_values(&self) -> usize {
self.attention_score_values
}
}
#[derive(Clone, Copy, Debug)]
enum CachedPhase {
Prefill,
Decode,
}
/// Reusable per-block cache storage plus coherent model-wide sequence state.
#[derive(Clone, Debug)]
pub struct DecoderKvCache {
config: DecoderModelConfig,
parameter_bindings: Vec<TensorValueBinding>,
layers: Vec<LayerKvCache>,
len: usize,
prefill_complete: bool,
work: DecoderKvCacheWork,
}
/// One decoder and one mutable model-wide cache bound for one checked session.
///
/// The retained parameter-value borrows prevent an optimizer from changing the
/// decoder while any K/V rows from that decoder may still be used.
pub struct DecoderKvSession<'model, 'cache> {
model: &'model DecoderModel,
cache: &'cache mut DecoderKvCache,
_parameter_value_guards: Vec<Ref<'model, Tensor>>,
}
impl PartialEq for DecoderKvCache {
fn eq(&self, other: &Self) -> bool {
same_config(self.config, other.config)
&& self.layers == other.layers
&& self.len == other.len
&& self.prefill_complete == other.prefill_complete
&& self.work == other.work
&& self.parameter_bindings.len() == other.parameter_bindings.len()
&& self
.parameter_bindings
.iter()
.zip(&other.parameter_bindings)
.all(|(left, right)| left.same_binding(right))
}
}
impl DecoderKvCache {
/// Allocates one fixed-capacity cache per block and captures compatibility evidence.
pub fn new(model: &DecoderModel) -> Result<Self, DecoderKvCacheError> {
let config = model.config();
let mut parameter_bindings = Vec::new();
parameter_bindings
.try_reserve_exact(model.parameters().len())
.map_err(|_| DecoderKvCacheError::ParameterAllocationFailed {
parameters: model.parameters().len(),
})?;
parameter_bindings.extend(
model
.parameters()
.iter()
.map(|parameter| TensorValueBinding::capture(parameter.tensor())),
);
let mut layers = Vec::new();
layers
.try_reserve_exact(model.blocks().len())
.map_err(|_| DecoderKvCacheError::LayerAllocationFailed {
layers: model.blocks().len(),
})?;
for (layer, block) in model.blocks().iter().enumerate() {
layers.push(
LayerKvCache::new(block.attention(), 1, config.max_positions())
.map_err(|source| DecoderKvCacheError::LayerCache { layer, source })?,
);
}
Ok(Self {
config,
parameter_bindings,
layers,
len: 0,
prefill_complete: false,
work: DecoderKvCacheWork::default(),
})
}
/// Validates one exact pairing and retains read borrows of the model's parameter values.
pub fn bind<'model, 'cache>(
&'cache mut self,
model: &'model DecoderModel,
) -> Result<DecoderKvSession<'model, 'cache>, DecoderKvCacheError> {
self.validate_binding(model)?;
let mut parameter_value_guards: Vec<Ref<'model, Tensor>> = Vec::new();
parameter_value_guards
.try_reserve_exact(model.parameters().len())
.map_err(|_| DecoderKvCacheError::ParameterAllocationFailed {
parameters: model.parameters().len(),
})?;
parameter_value_guards.extend(
model
.parameters()
.iter()
.map(|parameter| parameter.tensor().value()),
);
Ok(DecoderKvSession {
model,
cache: self,
_parameter_value_guards: parameter_value_guards,
})
}
fn next_work(
&self,
phase: CachedPhase,
score_values: usize,
) -> Result<DecoderKvCacheWork, DecoderKvCacheError> {
let layers = self.layers.len();
let qkv_rows = checked_mul(3, layers, DecoderKvCacheCounter::QkvProjectionRows)?;
Ok(DecoderKvCacheWork {
token_forwards: checked_add(
self.work.token_forwards,
1,
DecoderKvCacheCounter::TokenForwards,
)?,
prefill_tokens: checked_add(
self.work.prefill_tokens,
usize::from(matches!(phase, CachedPhase::Prefill)),
DecoderKvCacheCounter::PrefillTokens,
)?,
decode_tokens: checked_add(
self.work.decode_tokens,
usize::from(matches!(phase, CachedPhase::Decode)),
DecoderKvCacheCounter::DecodeTokens,
)?,
cache_appends: checked_add(
self.work.cache_appends,
layers,
DecoderKvCacheCounter::CacheAppends,
)?,
qkv_projection_rows: checked_add(
self.work.qkv_projection_rows,
qkv_rows,
DecoderKvCacheCounter::QkvProjectionRows,
)?,
attention_score_values: checked_add(
self.work.attention_score_values,
score_values,
DecoderKvCacheCounter::AttentionScoreValues,
)?,
})
}
fn validate_binding(&self, model: &DecoderModel) -> Result<(), DecoderKvCacheError> {
if !same_config(self.config, model.config()) {
return Err(DecoderKvCacheError::ModelConfigMismatch);
}
if self.parameter_bindings.len() != model.parameters().len() {
return Err(DecoderKvCacheError::ModelParameterCountMismatch {
cache: self.parameter_bindings.len(),
model: model.parameters().len(),
});
}
if let Some(index) = self
.parameter_bindings
.iter()
.zip(model.parameters())
.position(|(cached, parameter)| !cached.node_matches(parameter.tensor()))
{
return Err(DecoderKvCacheError::ModelParameterMismatch { index });
}
if let Some((index, (cached, parameter))) = self
.parameter_bindings
.iter()
.zip(model.parameters())
.enumerate()
.find(|(_, (cached, parameter))| !cached.revision_matches(parameter.tensor()))
{
return Err(DecoderKvCacheError::ModelParameterRevisionMismatch {
index,
cache: cached.revision(),
model: parameter.tensor().value_revision(),
});
}
if self.layers.len() != model.blocks().len() {
return Err(DecoderKvCacheError::LayerCountMismatch {
cache: self.layers.len(),
model: model.blocks().len(),
});
}
self.validate_layer_lengths()?;
for (layer, (block, cache)) in model.blocks().iter().zip(&self.layers).enumerate() {
if cache.batch_size() != 1 {
return Err(DecoderKvCacheError::LayerBatchSizeMismatch {
layer,
expected: 1,
actual: cache.batch_size(),
});
}
if cache.capacity() != self.capacity() {
return Err(DecoderKvCacheError::LayerCapacityMismatch {
layer,
expected: self.capacity(),
actual: cache.capacity(),
});
}
block
.attention()
.validate_incremental_cache_binding(cache)
.map_err(|source| DecoderKvCacheError::IncrementalAttention { layer, source })?;
}
Ok(())
}
fn validate_layer_lengths(&self) -> Result<(), DecoderKvCacheError> {
for (layer, cache) in self.layers.iter().enumerate() {
if cache.len() != self.len {
return Err(DecoderKvCacheError::LayerLengthInvariant {
layer,
expected: self.len,
actual: cache.len(),
});
}
}
Ok(())
}
fn reset(&mut self) {
for cache in &mut self.layers {
cache.reset();
}
self.len = 0;
self.prefill_complete = false;
self.work = DecoderKvCacheWork::default();
}
pub const fn len(&self) -> usize {
self.len
}
pub const fn capacity(&self) -> usize {
self.config.max_positions()
}
pub fn layer_count(&self) -> usize {
self.layers.len()
}
pub fn layer_len(&self, layer: usize) -> Option<usize> {
self.layers.get(layer).map(LayerKvCache::len)
}
pub fn layer_cache(&self, layer: usize) -> Option<&LayerKvCache> {
self.layers.get(layer)
}
pub const fn is_empty(&self) -> bool {
self.len == 0
}
pub const fn is_full(&self) -> bool {
self.len == self.capacity()
}
pub const fn work(&self) -> DecoderKvCacheWork {
self.work
}
}
impl DecoderKvSession<'_, '_> {
/// Fills every block cache from one validated, nonempty prompt.
///
/// Prompt rows advance serially through the shared one-row path. A later
/// internal failure restores empty logical state while retaining allocation
/// and the session's exact-model binding.
pub fn prefill(&mut self, prompt: &[u32]) -> Result<CachedDecoderOutput, DecoderKvCacheError> {
if prompt.is_empty() {
return Err(DecoderKvCacheError::EmptyPrompt);
}
if self.cache.len != 0 || self.cache.prefill_complete {
return Err(DecoderKvCacheError::PrefillRequiresEmpty {
len: self.cache.len,
});
}
if prompt.len() > self.cache.capacity() {
return Err(DecoderKvCacheError::PromptTooLong {
tokens: prompt.len(),
capacity: self.cache.capacity(),
});
}
for (position, &token_id) in prompt.iter().enumerate() {
if !valid_token(token_id, self.cache.config.vocabulary_size()) {
return Err(DecoderKvCacheError::PromptTokenOutOfBounds {
position,
token_id,
vocabulary_size: self.cache.config.vocabulary_size(),
});
}
}
let mut final_output = None;
for &token_id in prompt {
match self.forward_token(token_id, CachedPhase::Prefill) {
Ok(output) => final_output = Some(output),
Err(error) => {
self.cache.reset();
return Err(error);
}
}
}
self.cache.prefill_complete = true;
Ok(final_output.expect("a validated prompt has at least one token"))
}
/// Appends one selected token and returns logits for the following choice.
pub fn decode(&mut self, token_id: u32) -> Result<CachedDecoderOutput, DecoderKvCacheError> {
if !self.cache.prefill_complete || self.cache.len == 0 {
return Err(DecoderKvCacheError::DecodeRequiresPrefill);
}
if !valid_token(token_id, self.cache.config.vocabulary_size()) {
return Err(DecoderKvCacheError::DecodeTokenOutOfBounds {
token_id,
vocabulary_size: self.cache.config.vocabulary_size(),
});
}
if self.cache.is_full() {
return Err(DecoderKvCacheError::Full {
capacity: self.cache.capacity(),
});
}
self.forward_token(token_id, CachedPhase::Decode)
}
fn forward_token(
&mut self,
token_id: u32,
phase: CachedPhase,
) -> Result<CachedDecoderOutput, DecoderKvCacheError> {
let position = self.cache.len;
let layer_count = self.cache.layers.len();
let (logits, prepared, score_values) = no_grad(|| {
let embedding = self
.model
.embedding()
.forward(&[token_id], &[1, 1])
.map_err(DecoderKvCacheError::Embedding)?;
let mut current = embedding;
let mut prepared = Vec::new();
prepared.try_reserve_exact(layer_count).map_err(|_| {
DecoderKvCacheError::LayerAllocationFailed {
layers: layer_count,
}
})?;
let mut score_values = 0usize;
for (layer, (block, cache)) in self
.model
.blocks()
.iter()
.zip(&self.cache.layers)
.enumerate()
{
let attention_norm = block
.attention_norm()
.forward(¤t)
.map_err(|source| DecoderKvCacheError::AttentionNorm { layer, source })?;
let ticket = block
.attention()
.prepare_incremental_bound(&attention_norm, cache)
.map_err(|source| DecoderKvCacheError::IncrementalAttention {
layer,
source,
})?;
score_values = checked_add(
score_values,
ticket.attention_score_values(),
DecoderKvCacheCounter::AttentionScoreValues,
)?;
let after_attention = residual_add(¤t, ticket.output())
.map_err(|source| DecoderKvCacheError::AttentionResidual { layer, source })?;
let feed_forward_norm = block
.feed_forward_norm()
.forward(&after_attention)
.map_err(|source| DecoderKvCacheError::FeedForwardNorm { layer, source })?;
let feed_forward = block
.feed_forward()
.forward(&feed_forward_norm)
.map_err(|source| DecoderKvCacheError::FeedForward { layer, source })?;
current = residual_add(&after_attention, &feed_forward)
.map_err(|source| DecoderKvCacheError::FeedForwardResidual { layer, source })?;
prepared.push(ticket);
}
let final_norm = self
.model
.final_norm()
.forward(¤t)
.map_err(DecoderKvCacheError::FinalNorm)?;
let tied_weight = self
.model
.tied_embedding()
.tensor()
.transpose(0, 1)
.map_err(|source| DecoderKvCacheError::Autodiff {
stage: CachedDecoderStage::TiedWeightTranspose,
source,
})?;
let logits = final_norm.matmul(&tied_weight).map_err(|source| {
DecoderKvCacheError::Autodiff {
stage: CachedDecoderStage::TiedVocabularyProjection,
source,
}
})?;
Ok::<_, DecoderKvCacheError>((logits, prepared, score_values))
})?;
for (layer, (ticket, cache)) in prepared.iter().zip(&self.cache.layers).enumerate() {
if ticket.cache_len() != position + 1 || !ticket.matches_cache(cache) {
return Err(DecoderKvCacheError::PreparedCacheChanged { layer });
}
}
let next_len = checked_add(position, 1, DecoderKvCacheCounter::TokenForwards)?;
let next_work = self.cache.next_work(phase, score_values)?;
for (ticket, cache) in prepared.into_iter().zip(&mut self.cache.layers) {
let _ = ticket.commit(cache);
}
self.cache.len = next_len;
self.cache.work = next_work;
Ok(CachedDecoderOutput {
logits,
position,
cache_len: next_len,
attention_score_values: score_values,
})
}
/// Clears sequence state while retaining this session's exact model binding.
pub fn reset(&mut self) {
self.cache.reset();
}
/// Exposes the bound cache for deterministic state and work inspection.
pub fn cache(&self) -> &DecoderKvCache {
self.cache
}
}
fn valid_token(token_id: u32, vocabulary_size: usize) -> bool {
usize::try_from(token_id)
.ok()
.is_some_and(|token| token < vocabulary_size)
}
fn same_config(left: DecoderModelConfig, right: DecoderModelConfig) -> bool {
left.vocabulary_size() == right.vocabulary_size()
&& left.model_width() == right.model_width()
&& left.heads() == right.heads()
&& left.feed_forward_width() == right.feed_forward_width()
&& left.layers() == right.layers()
&& left.max_positions() == right.max_positions()
&& left.rope_base().to_bits() == right.rope_base().to_bits()
&& left.rms_epsilon().to_bits() == right.rms_epsilon().to_bits()
}
fn checked_add(
left: usize,
right: usize,
counter: DecoderKvCacheCounter,
) -> Result<usize, DecoderKvCacheError> {
left.checked_add(right)
.ok_or(DecoderKvCacheError::WorkOverflow { counter })
}
fn checked_mul(
left: usize,
right: usize,
counter: DecoderKvCacheCounter,
) -> Result<usize, DecoderKvCacheError> {
left.checked_mul(right)
.ok_or(DecoderKvCacheError::WorkOverflow { counter })
}generate_cached reuses the Chapter 36 selection rule and checks stop conditions
in the order EOS, token limit, then context limit. Only if none applies does it
send the selected token through decode to obtain later logits. The second
in the loaded fixture is therefore returned at the context boundary while the
full cache remains at length .
rust/crates/llm-from-scratch/src/generation/kv_cache.rs#cached-generation /// One selected token and the categorical evidence used by cached generation.
#[derive(Clone, Debug, PartialEq)]
pub struct CachedGenerationStep {
prefix_length: usize,
token_id: u32,
unit_draw: Option<f64>,
interval_start: f64,
interval_end: f64,
}
impl CachedGenerationStep {
pub const fn prefix_length(&self) -> usize {
self.prefix_length
}
pub const fn token_id(&self) -> u32 {
self.token_id
}
pub const fn unit_draw(&self) -> Option<f64> {
self.unit_draw
}
pub const fn interval_start(&self) -> f64 {
self.interval_start
}
pub const fn interval_end(&self) -> f64 {
self.interval_end
}
}
/// Exact cached work and its dense complete-prefix attention-score baseline.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct CachedGenerationWork {
prefill_tokens: usize,
decode_tokens: usize,
layer_cache_count: usize,
cache_appends: usize,
qkv_projection_rows: usize,
cached_attention_score_values: usize,
complete_prefix_attention_score_values: usize,
}
impl CachedGenerationWork {
pub const fn prefill_tokens(self) -> usize {
self.prefill_tokens
}
pub const fn decode_tokens(self) -> usize {
self.decode_tokens
}
pub const fn layer_cache_count(self) -> usize {
self.layer_cache_count
}
pub const fn cache_appends(self) -> usize {
self.cache_appends
}
pub const fn qkv_projection_rows(self) -> usize {
self.qkv_projection_rows
}
pub const fn cached_attention_score_values(self) -> usize {
self.cached_attention_score_values
}
pub const fn complete_prefix_attention_score_values(self) -> usize {
self.complete_prefix_attention_score_values
}
}
/// Cached tokens, stops, final logical state, and exact work evidence.
#[derive(Clone, Debug, PartialEq)]
pub struct CachedGenerationResult {
prompt: Vec<u32>,
generated: Vec<u32>,
steps: Vec<CachedGenerationStep>,
stop: GenerationStop,
final_cache_len: usize,
work: CachedGenerationWork,
}
impl CachedGenerationResult {
pub fn prompt(&self) -> &[u32] {
&self.prompt
}
pub fn generated(&self) -> &[u32] {
&self.generated
}
pub fn steps(&self) -> &[CachedGenerationStep] {
&self.steps
}
pub const fn stop(&self) -> GenerationStop {
self.stop
}
pub const fn final_cache_len(&self) -> usize {
self.final_cache_len
}
pub const fn work(&self) -> CachedGenerationWork {
self.work
}
}
/// A cached-generation request, model step, sampling step, or work count failed.
#[derive(Debug, PartialEq)]
pub enum CachedGenerationError {
Cache(DecoderKvCacheError),
Sampling(SamplingError),
EmptyPrompt,
PromptTooLong {
tokens: usize,
max_positions: usize,
},
PromptTokenOutOfBounds {
position: usize,
token_id: u32,
vocabulary_size: usize,
},
EosTokenOutOfBounds {
token_id: u32,
vocabulary_size: usize,
},
LogitCountMismatch {
expected: usize,
actual: usize,
},
AllocationFailed {
values: usize,
},
WorkOverflow {
counter: DecoderKvCacheCounter,
},
}
impl fmt::Display for CachedGenerationError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Cache(source) => source.fmt(formatter),
Self::Sampling(source) => source.fmt(formatter),
Self::EmptyPrompt => formatter.write_str("cached generation needs a nonempty prompt"),
Self::PromptTooLong {
tokens,
max_positions,
} => write!(
formatter,
"cached generation prompt has {tokens} tokens, exceeding context capacity {max_positions}"
),
Self::PromptTokenOutOfBounds {
position,
token_id,
vocabulary_size,
} => write!(
formatter,
"cached generation prompt token {token_id} at position {position} is out of bounds for vocabulary {vocabulary_size}"
),
Self::EosTokenOutOfBounds {
token_id,
vocabulary_size,
} => write!(
formatter,
"cached generation EOS token {token_id} is out of bounds for vocabulary {vocabulary_size}"
),
Self::LogitCountMismatch { expected, actual } => write!(
formatter,
"cached last-position logits need {expected} values, received {actual}"
),
Self::AllocationFailed { values } => write!(
formatter,
"cannot allocate cached generation evidence for {values} values"
),
Self::WorkOverflow { counter } => {
write!(formatter, "cached generation {counter} counter overflows")
}
}
}
}
impl Error for CachedGenerationError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::Cache(source) => Some(source),
Self::Sampling(source) => Some(source),
_ => None,
}
}
}
impl From<DecoderKvCacheError> for CachedGenerationError {
fn from(source: DecoderKvCacheError) -> Self {
Self::Cache(source)
}
}
impl From<SamplingError> for CachedGenerationError {
fn from(source: SamplingError) -> Self {
Self::Sampling(source)
}
}
/// Generates with one prompt prefill followed by only the needed one-token decodes.
pub fn generate_cached(
model: &DecoderModel,
prompt: &[u32],
config: GenerationConfig,
rng: &mut SplitMix64,
) -> Result<CachedGenerationResult, CachedGenerationError> {
let model_config = model.config();
let vocabulary_size = model_config.vocabulary_size();
let max_positions = model_config.max_positions();
validate_generation_request(vocabulary_size, max_positions, prompt, config)?;
let planned_steps = config.max_new_tokens().min(
max_positions
.checked_sub(prompt.len())
.and_then(|remaining| remaining.checked_add(1))
.ok_or(CachedGenerationError::AllocationFailed { values: usize::MAX })?,
);
let mut prompt_copy = Vec::new();
prompt_copy.try_reserve_exact(prompt.len()).map_err(|_| {
CachedGenerationError::AllocationFailed {
values: prompt.len(),
}
})?;
prompt_copy.extend_from_slice(prompt);
let mut generated = Vec::new();
generated.try_reserve_exact(planned_steps).map_err(|_| {
CachedGenerationError::AllocationFailed {
values: planned_steps,
}
})?;
let mut steps = Vec::new();
steps.try_reserve_exact(planned_steps).map_err(|_| {
CachedGenerationError::AllocationFailed {
values: planned_steps,
}
})?;
if config.max_new_tokens() == 0 {
return Ok(CachedGenerationResult {
prompt: prompt_copy,
generated,
steps,
stop: GenerationStop::TokenLimit,
final_cache_len: 0,
work: CachedGenerationWork {
layer_cache_count: model_config.layers(),
..CachedGenerationWork::default()
},
});
}
let mut cache = DecoderKvCache::new(model)?;
let mut session = cache.bind(model)?;
let mut current = session.prefill(prompt)?;
let mut prefix_length = prompt.len();
let mut complete_prefix_attention_score_values = 0usize;
let stop = loop {
complete_prefix_attention_score_values = checked_complete_prefix_scores(
complete_prefix_attention_score_values,
model_config.layers(),
model_config.heads(),
prefix_length,
)?;
let decision = {
let logits = current.logits().value();
if logits.len() != vocabulary_size {
return Err(CachedGenerationError::LogitCountMismatch {
expected: vocabulary_size,
actual: logits.len(),
});
}
sample_next_token(logits.as_slice(), config.mode(), rng)?
};
let token_id = decision.token_id();
generated.push(token_id);
steps.push(CachedGenerationStep {
prefix_length,
token_id,
unit_draw: decision.unit_draw(),
interval_start: decision.interval_start(),
interval_end: decision.interval_end(),
});
prefix_length =
prefix_length
.checked_add(1)
.ok_or(CachedGenerationError::WorkOverflow {
counter: DecoderKvCacheCounter::TokenForwards,
})?;
if config.eos_token() == Some(token_id) {
break GenerationStop::Eos;
}
if generated.len() == config.max_new_tokens() {
break GenerationStop::TokenLimit;
}
if prefix_length > max_positions {
break GenerationStop::ContextLimit;
}
current = session.decode(token_id)?;
};
let cache_work = session.cache().work();
Ok(CachedGenerationResult {
prompt: prompt_copy,
generated,
steps,
stop,
final_cache_len: session.cache().len(),
work: CachedGenerationWork {
prefill_tokens: cache_work.prefill_tokens(),
decode_tokens: cache_work.decode_tokens(),
layer_cache_count: session.cache().layer_count(),
cache_appends: cache_work.cache_appends(),
qkv_projection_rows: cache_work.qkv_projection_rows(),
cached_attention_score_values: cache_work.attention_score_values(),
complete_prefix_attention_score_values,
},
})
}
fn validate_generation_request(
vocabulary_size: usize,
max_positions: usize,
prompt: &[u32],
config: GenerationConfig,
) -> Result<(), CachedGenerationError> {
if prompt.is_empty() {
return Err(CachedGenerationError::EmptyPrompt);
}
if prompt.len() > max_positions {
return Err(CachedGenerationError::PromptTooLong {
tokens: prompt.len(),
max_positions,
});
}
for (position, &token_id) in prompt.iter().enumerate() {
if !valid_token(token_id, vocabulary_size) {
return Err(CachedGenerationError::PromptTokenOutOfBounds {
position,
token_id,
vocabulary_size,
});
}
}
if let Some(token_id) = config.eos_token()
&& !valid_token(token_id, vocabulary_size)
{
return Err(CachedGenerationError::EosTokenOutOfBounds {
token_id,
vocabulary_size,
});
}
if let SamplingMode::TemperatureTopK { temperature, top_k } = config.mode() {
if !temperature.is_finite() || temperature <= 0.0 {
return Err(SamplingError::InvalidTemperature { value: temperature }.into());
}
if top_k == 0 || top_k > vocabulary_size {
return Err(SamplingError::InvalidTopK {
top_k,
vocabulary_size,
}
.into());
}
}
Ok(())
}
fn checked_complete_prefix_scores(
current: usize,
layers: usize,
heads: usize,
prefix_length: usize,
) -> Result<usize, CachedGenerationError> {
let counter = DecoderKvCacheCounter::CompletePrefixAttentionScoreValues;
let square = prefix_length
.checked_mul(prefix_length)
.ok_or(CachedGenerationError::WorkOverflow { counter })?;
let values = layers
.checked_mul(heads)
.and_then(|factor| factor.checked_mul(square))
.ok_or(CachedGenerationError::WorkOverflow { counter })?;
current
.checked_add(values)
.ok_or(CachedGenerationError::WorkOverflow { counter })
}DecoderKvSession::reset clears logical length, phase, and work while preserving
the backing allocation, stored K/V values outside the empty logical prefix, and
the current session’s model/cache relationship. Decode before prefill,
prefill into nonempty state, and overflow return typed operation errors without
changing the committed state. A rebuilt model or changed decoder configuration
is rejected earlier, when code tries to bind the cache.
While a session exists, an AdamW step reaches its existing fallible write
boundary and returns ParameterValueBorrowed. Parameter values, optimizer state,
and cache state remain unchanged. The failed update does not invalidate the
session, so any otherwise valid later decode can continue.
Dropping the session releases its read-only borrows. AdamW may then update the
model and advance its value revisions. A later attempt to bind the old cache
returns ModelParameterRevisionMismatch; reset cannot make that stale cache
compatible. Construct a new DecoderKvCache from the updated model before
creating the next session.
The executable checks prefill and decode newest-position logits within tolerance, records that the two layer buffers are distinct, measures cached and reference attention tensors, compares the restored fixture’s generation decisions and EOS behavior, resets and replays, and checks rejected binds and operations against snapshots captured before each call.
rust/demos/ch38-cached-generation/src/lib.rs#learner-evidence /// Checks model-wide cache coherence, loaded generation parity, reset, and errors.
pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
let model = fixture_model()?;
let config = model.config();
let mut cache = DecoderKvCache::new(&model)?;
let mut session = cache.bind(&model)?;
let prefill_output = session.prefill(&PROMPT)?;
let prefill = phase_evidence(&model, session.cache(), &prefill_output, &PROMPT, 0, 0)?;
let cached_scores_after_prefill = session.cache().work().attention_score_values();
let decode_output = session.decode(DECODE_TOKEN)?;
let decode_prefix = [PROMPT[0], PROMPT[1], DECODE_TOKEN];
let decode = phase_evidence(
&model,
session.cache(),
&decode_output,
&decode_prefix,
PROMPT.len(),
cached_scores_after_prefill,
)?;
let layer_storage_distinct = session
.cache()
.layer_cache(0)
.zip(session.cache().layer_cache(1))
.is_some_and(|(left, right)| {
left.key_storage().as_ptr() != right.key_storage().as_ptr()
&& left.value_storage().as_ptr() != right.value_storage().as_ptr()
});
let work = session.cache().work();
let complete_prefix_attention_score_values = prefill
.complete_prefix_attention_score_values
.checked_add(decode.complete_prefix_attention_score_values)
.ok_or(FixtureError::Invariant(
"measured complete-prefix score total overflowed",
))?;
let loaded = loaded_generation_evidence()?;
let reset = reset_evidence(&mut session, &decode)?;
let errors = error_evidence(&model)?;
let cached_retained_lengths = [1, prefill.cache_after, decode.cache_after];
let complete_prefix_lengths = [prefill.cache_after, decode.cache_after];
let history = historical_cache_contrast(
config,
&cached_retained_lengths,
&complete_prefix_lengths,
work.attention_score_values(),
complete_prefix_attention_score_values,
)?;
require(
prefill.max_abs_difference <= TOLERANCE,
"prefill logits differ",
)?;
require(
decode.max_abs_difference <= TOLERANCE,
"decode logits differ",
)?;
require(layer_storage_distinct, "layer caches share storage")?;
require(
work.prefill_tokens() == 2
&& work.decode_tokens() == 1
&& work.cache_appends() == 6
&& work.qkv_projection_rows() == 18
&& work.attention_score_values() == 24
&& complete_prefix_attention_score_values == 52,
"two-layer work counters changed",
)?;
require(
reset.after == 0
&& reset.allocation_reused
&& reset.storage_unchanged
&& reset.work_zeroed
&& reset.replay_identical,
"reset evidence changed",
)?;
require(errors.unchanged, "a rejected cache operation changed state")?;
require(
loaded.cached.work().cached_attention_score_values() == 6
&& loaded
.cached
.work()
.complete_prefix_attention_score_values()
== 10,
"loaded score counts changed",
)?;
Ok(LearnerEvidence {
config,
prefill,
decode,
layer_storage_distinct,
work,
complete_prefix_attention_score_values,
loaded,
reset,
errors,
history,
})
}Run cargo run --quiet --locked -p ch38-cached-generation. It prints the exact
learner report:
chapter=38-cached-generation
config=layers:2 heads:2 model_width:4 context:4 tolerance:0.000000000002
prefill=prompt:[0,1] cache:0->2 layer_lengths:[2,2] shape:[1,2,2,2] cached_scores:12 complete_prefix_scores:16 max_abs_diff:0.000000000000 logits:[1.768374438,0.208825256,1.056205728,-0.451857108,0.388467944]
decode=token:2 position:2 cache:2->3 layer_lengths:[3,3] shape:[1,2,3,2] cached_scores:12 complete_prefix_scores:36 max_abs_diff:0.000000000000 logits:[0.032908910,-0.679583624,1.408381841,0.525525421,-0.588014095]
work=prefill_tokens:2 decode_tokens:1 layer_caches:2 cache_appends:6 qkv_rows:18 cached_scores:24 complete_prefix_scores:52 layer_storage_distinct:true
loaded=checkpoint_bytes:6330 context_capacity:2 rng_state:0x9e3779b97f4a7c38 prompt:[0] generated:[4,4] text:44 prefixes:[1,2] stop:context-limit final_cache:2 prefill_tokens:1 decode_tokens:1 cached_scores:6 complete_prefix_scores:10 tokens_match:true rng_match:true
eos=token:4 generated:[4] stop:eos final_cache:1 decode_tokens:0 tokens_match:true rng_match:true
reset=before:3 after:0 allocation_reused:true storage_unchanged:true work_zeroed:true replay_identical:true
errors=decode_before_prefill:true prefill_nonempty:true overflow:true rebuilt_model:true changed_config:true unchanged:true
history=lanes:4 cached_lengths:[1,2,3] cached_scores:24 complete_prefix_lengths:[2,3] complete_prefix_scores:52 avoided_scores:28
next=assemble the complete end-to-end LLM pipelinerust/demos/ch38-cached-generation/src/main.rs fn main() -> Result<(), Box<dyn std::error::Error>> {
print!("{}", ch38_cached_generation::learner_report()?);
Ok(())
}Follow prefill into one-token decode
The figure begins with prompt positions and in two distinct block caches, then follows selected token into absolute position . Solid-border prefill layer summaries and double-border decode layer summaries distinguish the phases even without color. Each cached logit coordinate sits beside its complete-prefix counterpart, making the within-tolerance comparison direct.
The work cards compare cached score values with reference values. The remaining evidence keeps two different stop boundaries separate. In the loaded run, prefill reaches cache length , the first selected is decoded to length , and the second selected is returned before context-limit stops without another decoder forward. The EOS run returns its first selected token and also stops before another decoder forward.
Prefill every layer for the prompt; decode every layer together
The exact Rust trace follows a two-token prompt and one later token through two distinct decoder-block caches, checks newest-position logits within tolerance, and compares measured attention-score work plus stopping and reset behavior.
- prompt prefill — solid border
- new one-token decode — double border
- newest logits match within tolerance
Move from prompt prefill to one-token decode
The prompt initializes both block caches at positions zero and one. Token two then enters at position two, and both block lengths advance coherently.
prompt prefill — solid border
Prompt token IDs: [0,1]
Absolute position:
Logical cache length:
Layer-cache shape:
- Decoder block Logical cache length owns separate K/V storage
- Decoder block Logical cache length owns separate K/V storage
Cached path:
Complete-prefix reference:
newest logits match within tolerance Maximum absolute difference:
new one-token decode — double border
Selected token:
Absolute position:
Logical cache length:
Layer-cache shape:
- Decoder block Logical cache length owns separate K/V storage
- Decoder block Logical cache length owns separate K/V storage
Cached path:
Complete-prefix reference:
newest logits match within tolerance Maximum absolute difference:
Count attention-score values only
The fixture counts score values formed inside attention. It does not treat those counts as total runtime or a hardware measurement.
Cached path
Four layer-head lanes form one score row at each retained length.
cache_appends=6 qkv_rows=18 Complete-prefix reference
The two reference calls rebuild square score grids at prompt length two and decoded length three.
layer_caches=2 Compare the restored fixture's sampling and stops
For the restored fixture, cached and complete-prefix paths make the same token and sampling-draw decisions, finish with the same RNG state, and agree on context and EOS stops.
Restored checkpoint: matching decisions and exact context boundary
Prompt token IDs: [0]
Generated token IDs and text: [4,4] -> 44
Stopping reason: context-limit
Logical cache length:
Attention-score values:
tokens_match=true rng_match=true EOS: select the token, then stop before an unnecessary decode
Selected token:
Generated token IDs and text: [4]
Stopping reason: eos
Logical cache length:
tokens_match=true rng_match=true Reset or reject without corrupting model-wide state
Reset retains allocations while clearing logical state and work. Invalid phase or capacity operations and incompatible bind attempts leave committed state unchanged.
logical state returns to zero
allocation_reused=true storage_unchanged=true work_zeroed=true replay_identical=true rejected calls commit no state change
decode_before_prefill=true prefill_nonempty=true overflow=true rebuilt_model=true changed_config=true unchanged=true Predict before checking the results
- How many layer caches does a two-block decoder own?
- After prefill with two tokens, what is each layer cache’s logical length?
- Which absolute position does the first decode token use?
- At which boundary does a cache from rebuilt equal-valued weights fail, and why?
- Why does a successful session retain read-only borrows of every parameter value after checking the captured value revisions?
- How many cache appends occur across two prompt rows and one decode row?
- Why does the fixture form cached attention-score values?
- Does caching make the newest attention query constant-time?
- If the first selected token is EOS, must that token be decoded?
- Which compatibility checks happen once when a session is bound, which checks still happen for each operation, and what relationship does reset retain?
Misconception: caches with equal shapes can be shared by different blocks. In
fact, each retained row depends on the hidden state entering its block, so block
cannot reuse block ‘s K/V rows. DecoderKvCache stores compatibility
evidence for the complete model-wide state, and DecoderKvSession checks and
retains the exact decoder/cache relationship. Each nested layer cache also
records its own block’s attention parameters, geometry, and RoPE configuration.
Check the ten predictions
- The decoder owns layer caches, one for each block.
- Both caches reach logical length and shape .
- The first later token uses absolute position , equal to the old cache length.
bindperforms the identity check. Even when values, shapes, and decoder configuration agree, rebuilding creates different parameter nodes, so the cache cannot open a session with that rebuilt model.- The read-only borrows prevent AdamW from changing a value after
bindchecks its revision but before a later decoder operation reads it. They preserve the checked relationship for the full session without copying the weights. - blocks times token rows gives coherent cache appends.
- There are batch-layer-head lanes, and each forms score values, giving .
- No. The newest query still reads retained keys at prefix length .
- No. EOS is returned as the selected token, and no later logits are needed.
bindchecks stable configuration, parameter identity and revisions, layer lengths and geometry, attention bindings, and RoPE once for that session. Each operation still checks its prompt, phase, token, remaining capacity, counters, and prepared-ticket storage and length. Reset clears logical length, phase, and counters while retaining the same model/cache relationship, backing allocation, and old K/V bits outside the empty logical prefix.
Connect inference to the whole pipeline
The complete decoder can now bind compatible graph-free K/V state across all blocks for one session, prefill a prompt, and decode selected tokens without receiving the model again. The session keeps using the exact decoder it bound and retains read-only access to its parameter values. After the session ends, a weight update makes the old cache stale, so the updated model needs a newly constructed cache. The exact fixture still matches the complete-prefix generation decisions. Chapter 39 will connect this inference path to the full pipeline; inside that execution test cannot affect the selected state, while Chapter 39’s checked-in decoder-lower-than-bigram loss ordering is retained only as fixed-fixture regression evidence.
That final chapter will exercise the whole road in one program: partition data, learn and apply BPE, train and select the decoder, run a locally isolated test evaluation, save and reload the checkpoint, prefill the prompt, generate through the model-wide cache, and decode the resulting token IDs back to text. That within-run boundary does not turn Chapter 39’s checked-in decoder-versus-bigram result into a new independent generalization estimate when later executions repeat the comparison.