37 · Content revision 5
Keep the prefix, project only the new row
Learn how one layer-bound KV cache appends rotated keys and unrotated values while reproducing full-prefix attention at the newest position.
Predict the third result before running the example
The fixture has one batch, model width , heads, head width , and cache capacity . Before the first call the cache is empty. Each call accepts one row shaped , so the logical key/value shapes advance through , , and .
The first two cached outputs are and . Before revealing the third output, predict three facts:
- The old length is , so the new and use absolute RoPE position .
- Each head scores its newest query against keys, then forms its output from the corresponding values.
- The cache appends one rotated key row and one unrotated value row; it does not retain the query.
The result is . An independent full-prefix pass produces the same displayed values, and the unrounded maximum absolute difference is at most .
Append along the position axis
For attention layer , the cache transition is
The semicolon means concatenate along the sequence-position axis. The new key has already received RoPE for its absolute position. The new value has been projected and split into heads but is not rotated. Queries are needed only for the current calculation, so this implementation does not store them in the cache.
The implementation uses a zero-based offset equal to the old cache length. The formula uses conventional one-based prefix notation, so mathematical position is one greater than that offset.
Keep layer, logical length, and capacity separate
- is layer ‘s rotated key prefix after the append.
- is the same layer’s value prefix after the append.
- identifies the decoder-block attention layer that owns the state.
- identifies the newest position in one-based mathematical prefix notation.
- includes every retained position through the newest one.
- and are the unchanged earlier rows.
- and are the one-row candidate pair.
- means append after along the position axis.
The physical buffers have shape , where is batch size, is head count, is fixed capacity, and is head width. A logical snapshot has shape . Reset changes to without changing or allocating replacement buffers. It also leaves the parameter-node identities and value revisions captured at cache construction unchanged.
From causal attention to managed LLM inference state
Causal Transformer attention lets each newest decoder position read the known prefix. Without retained projections, however, a generation loop can recompute unchanged earlier key and value rows at every step.
Attention Is All You Need describes the full-prefix causal computation used as the reference here. Vaswani and colleagues define scaled dot-product attention, describe an autoregressive decoder that emits one element at a time, and mask decoder self-attention so a position can use only the known prefix; the paper does not specify KV caching or RoPE.
Fast Transformer Decoding: One Write-Head is All You Need makes the reuse boundary explicit. Shazeer’s incremental multi-head self-attention takes previous key and value tensors, appends the current projected pair, and returns the enlarged tensors; its analysis identifies repeatedly loading those tensors as a memory-bandwidth bottleneck before proposing multi-query attention.
Efficient Memory Management for Large Language Model Serving with PagedAttention continues the road into modern serving. Kwon and colleagues describe sequential LLM generation in which earlier key and value vectors are cached and only the newest pair is computed, then organize dynamically growing KV caches as logical blocks mapped to non-contiguous physical memory.
Incremental decoding made one-step reuse explicit by retaining per-layer key/value tensors; later LLM serving systems treated the growing KV cache as a central memory-management object.
Modern cached generation keeps one compatible cache per decoder block: each new query reads the retained prefix while only the newest key and value are projected.
The approach developed from causal masked attention, through explicit incremental KV reuse, to serving systems organized around dynamically growing caches. This implementation additionally requires fixed capacity, exact tensor layout, the old cache length as the RoPE offset, unchanged bound parameters, allocation-preserving reset, and typed errors.
rust/demos/ch37-incremental-attention/src/lib.rs#historical-kv-contrast /// Measures complete-prefix recomputation against one-row incremental projection.
pub fn historical_kv_contrast(
steps: &[StepEvidence],
) -> Result<HistoricalKvContrast, FixtureError> {
require(
!steps.is_empty(),
"history evidence needs at least one step",
)?;
let mut newest_query_key_rows = Vec::new();
newest_query_key_rows
.try_reserve_exact(steps.len())
.map_err(|_| FixtureError::Invariant("cannot allocate history evidence"))?;
for step in steps {
let rows = step
.heads
.first()
.ok_or(FixtureError::Invariant(
"history step has no attention head",
))?
.weights
.len();
require(
step.heads.iter().all(|head| head.weights.len() == rows),
"attention heads disagree about retained key rows",
)?;
require(
rows == step.cache_after && step.cache_shape.get(2) == Some(&step.cache_after),
"attention span disagrees with logical cache length",
)?;
newest_query_key_rows.push(rows);
}
let complete_prefix_rows_per_projection =
steps.iter().map(|step| step.full_rows_per_projection).sum();
let incremental_rows_per_projection = steps
.iter()
.map(|step| step.incremental_rows_per_projection)
.sum();
let reused_rows = steps.iter().map(|step| step.reused_key_value_rows).sum();
require(
complete_prefix_rows_per_projection == incremental_rows_per_projection + reused_rows,
"projection-row contrast is inconsistent",
)?;
Ok(HistoricalKvContrast {
newest_query_key_rows,
complete_prefix_rows_per_projection,
incremental_rows_per_projection,
reused_key_rows: reused_rows,
reused_value_rows: reused_rows,
})
}Bind the state, calculate completely, then commit
LayerKvCache::new takes the actual attention layer. It records the batch size,
model width, head count, cache capacity, and head width; owns fixed key/value
buffers; and, for each of the layer’s four parameters, captures both the node
identity and its current value revision. It also records the exact RoPE
configuration.
A rebuilt layer contains different parameter nodes, so the cache returns
CacheLayerMismatch. A successful AdamW step updates the existing nodes in
place: their identities stay the same, but their value revisions advance. The
old cache is then stale and returns CacheLayerRevisionMismatch. This second
check matters because retained K/V rows were projected with the earlier weight
values; mixing those rows with projections from updated weights would describe
no single model. Construct a new cache after any weight update. Calling reset
on the old cache does not make it compatible: reset clears logical length but
does not refresh the captured revisions or bind the cache to another layer.
rust/crates/llm-from-scratch/src/attention/incremental.rs#layer-kv-cache /// A rejected cache configuration, append, or logical snapshot.
#[derive(Clone, Debug, PartialEq)]
pub enum LayerKvCacheError {
ZeroBatchSize,
ZeroCapacity,
CapacityExceedsPositions {
capacity: usize,
max_positions: usize,
},
ElementCountOverflow {
batch_size: usize,
heads: usize,
capacity: usize,
head_width: usize,
},
AllocationFailed {
elements: usize,
},
Full {
capacity: usize,
},
KeyShapeMismatch {
expected: Vec<usize>,
actual: Vec<usize>,
},
ValueShapeMismatch {
expected: Vec<usize>,
actual: Vec<usize>,
},
NonFiniteKey {
index: usize,
value: f64,
},
NonFiniteValue {
index: usize,
value: f64,
},
Tensor(TensorError),
}
impl fmt::Display for LayerKvCacheError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ZeroBatchSize => formatter.write_str("KV cache batch size must be nonzero"),
Self::ZeroCapacity => formatter.write_str("KV cache capacity must be nonzero"),
Self::CapacityExceedsPositions {
capacity,
max_positions,
} => write!(
formatter,
"KV cache capacity {capacity} exceeds RoPE position capacity {max_positions}"
),
Self::ElementCountOverflow {
batch_size,
heads,
capacity,
head_width,
} => write!(
formatter,
"KV cache element count overflows for batch {batch_size}, {heads} heads, capacity {capacity}, and head width {head_width}"
),
Self::AllocationFailed { elements } => write!(
formatter,
"cannot allocate KV cache storage for {elements} f64 values"
),
Self::Full { capacity } => {
write!(formatter, "KV cache is full at capacity {capacity}")
}
Self::KeyShapeMismatch { expected, actual } => write!(
formatter,
"appended key must have shape {expected:?}, got {actual:?}"
),
Self::ValueShapeMismatch { expected, actual } => write!(
formatter,
"appended value must have shape {expected:?}, got {actual:?}"
),
Self::NonFiniteKey { index, value } => write!(
formatter,
"appended key value at flat index {index} must be finite, got {value:?}"
),
Self::NonFiniteValue { index, value } => write!(
formatter,
"appended value at flat index {index} must be finite, got {value:?}"
),
Self::Tensor(source) => source.fmt(formatter),
}
}
}
impl Error for LayerKvCacheError {
fn source(&self) -> Option<&(dyn Error + 'static)> {
match self {
Self::Tensor(source) => Some(source),
_ => None,
}
}
}
impl From<TensorError> for LayerKvCacheError {
fn from(source: TensorError) -> Self {
Self::Tensor(source)
}
}
/// Fixed-capacity rotated-key and value storage for one attention layer.
///
/// Physical storage has layout `[batch, heads, capacity, head_width]`. `len`
/// selects the logical prefix; reset keeps the allocation and moves that prefix
/// back to zero.
#[derive(Clone, Debug)]
pub struct LayerKvCache {
batch_size: usize,
model_width: usize,
heads: usize,
head_width: usize,
capacity: usize,
len: usize,
keys: Vec<f64>,
values: Vec<f64>,
parameter_bindings: [TensorValueBinding; 4],
rope_feature_width: usize,
rope_max_positions: usize,
rope_base_bits: u64,
}
impl PartialEq for LayerKvCache {
fn eq(&self, other: &Self) -> bool {
self.batch_size == other.batch_size
&& self.model_width == other.model_width
&& self.heads == other.heads
&& self.head_width == other.head_width
&& self.capacity == other.capacity
&& self.len == other.len
&& self.keys == other.keys
&& self.values == other.values
&& self.rope_feature_width == other.rope_feature_width
&& self.rope_max_positions == other.rope_max_positions
&& self.rope_base_bits == other.rope_base_bits
&& self
.parameter_bindings
.iter()
.zip(&other.parameter_bindings)
.all(|(left, right)| left.same_binding(right))
}
}
impl LayerKvCache {
pub fn new(
layer: &MultiHeadAttention,
batch_size: usize,
capacity: usize,
) -> Result<Self, LayerKvCacheError> {
if batch_size == 0 {
return Err(LayerKvCacheError::ZeroBatchSize);
}
if capacity == 0 {
return Err(LayerKvCacheError::ZeroCapacity);
}
if capacity > layer.rope().max_positions() {
return Err(LayerKvCacheError::CapacityExceedsPositions {
capacity,
max_positions: layer.rope().max_positions(),
});
}
let model_width = layer.model_width();
let heads = layer.heads();
let head_width = layer.head_width();
let elements = batch_size
.checked_mul(heads)
.and_then(|count| count.checked_mul(capacity))
.and_then(|count| count.checked_mul(head_width))
.ok_or(LayerKvCacheError::ElementCountOverflow {
batch_size,
heads,
capacity,
head_width,
})?;
let keys = zeroed_storage(elements)?;
let values = zeroed_storage(elements)?;
let parameters = layer.parameters();
debug_assert_eq!(parameters.len(), 4);
Ok(Self {
batch_size,
model_width,
heads,
head_width,
capacity,
len: 0,
keys,
values,
parameter_bindings: [
TensorValueBinding::capture(parameters[0].tensor()),
TensorValueBinding::capture(parameters[1].tensor()),
TensorValueBinding::capture(parameters[2].tensor()),
TensorValueBinding::capture(parameters[3].tensor()),
],
rope_feature_width: layer.rope().feature_width(),
rope_max_positions: layer.rope().max_positions(),
rope_base_bits: layer.rope().base().to_bits(),
})
}
fn validate_append(&self, key: &Tensor, value: &Tensor) -> Result<(), LayerKvCacheError> {
let expected = [self.batch_size, self.heads, 1, self.head_width];
if key.shape() != expected.as_slice() {
return Err(LayerKvCacheError::KeyShapeMismatch {
expected: expected.to_vec(),
actual: key.shape().to_vec(),
});
}
if value.shape() != expected.as_slice() {
return Err(LayerKvCacheError::ValueShapeMismatch {
expected: expected.to_vec(),
actual: value.shape().to_vec(),
});
}
if self.is_full() {
return Err(LayerKvCacheError::Full {
capacity: self.capacity,
});
}
if let Some((index, value)) = first_nonfinite(key.as_slice()) {
return Err(LayerKvCacheError::NonFiniteKey { index, value });
}
if let Some((index, value)) = first_nonfinite(value.as_slice()) {
return Err(LayerKvCacheError::NonFiniteValue { index, value });
}
Ok(())
}
fn append_prevalidated(&mut self, key: &Tensor, value: &Tensor) {
debug_assert!(self.validate_append(key, value).is_ok());
for batch in 0..self.batch_size {
for head in 0..self.heads {
let source_start = (batch * self.heads + head) * self.head_width;
let destination_start =
((batch * self.heads + head) * self.capacity + self.len) * self.head_width;
self.keys[destination_start..destination_start + self.head_width]
.copy_from_slice(&key.as_slice()[source_start..source_start + self.head_width]);
self.values[destination_start..destination_start + self.head_width]
.copy_from_slice(
&value.as_slice()[source_start..source_start + self.head_width],
);
}
}
self.len += 1;
}
/// Returns the logical rotated-key prefix as `[batch, heads, len, head_width]`.
pub fn keys(&self) -> Result<Tensor, LayerKvCacheError> {
self.logical_tensor(&self.keys)
}
/// Returns the logical value prefix as `[batch, heads, len, head_width]`.
pub fn values(&self) -> Result<Tensor, LayerKvCacheError> {
self.logical_tensor(&self.values)
}
fn logical_tensor(&self, storage: &[f64]) -> Result<Tensor, LayerKvCacheError> {
let elements = self
.batch_size
.checked_mul(self.heads)
.and_then(|count| count.checked_mul(self.len))
.and_then(|count| count.checked_mul(self.head_width))
.ok_or(LayerKvCacheError::ElementCountOverflow {
batch_size: self.batch_size,
heads: self.heads,
capacity: self.len,
head_width: self.head_width,
})?;
let mut logical = Vec::new();
logical
.try_reserve_exact(elements)
.map_err(|_| LayerKvCacheError::AllocationFailed { elements })?;
for batch in 0..self.batch_size {
for head in 0..self.heads {
let start = (batch * self.heads + head) * self.capacity * self.head_width;
let end = start + self.len * self.head_width;
logical.extend_from_slice(&storage[start..end]);
}
}
Tensor::from_vec(
vec![self.batch_size, self.heads, self.len, self.head_width],
logical,
)
.map_err(Into::into)
}
/// Empties the logical prefix without reallocating either backing buffer.
pub fn reset(&mut self) {
self.len = 0;
}
pub const fn batch_size(&self) -> usize {
self.batch_size
}
pub const fn model_width(&self) -> usize {
self.model_width
}
pub const fn heads(&self) -> usize {
self.heads
}
pub const fn head_width(&self) -> usize {
self.head_width
}
pub const fn capacity(&self) -> usize {
self.capacity
}
pub const fn len(&self) -> usize {
self.len
}
pub const fn is_empty(&self) -> bool {
self.len == 0
}
pub const fn is_full(&self) -> bool {
self.len == self.capacity
}
/// Exposes allocated key storage for deterministic state audits.
pub fn key_storage(&self) -> &[f64] {
&self.keys
}
/// Exposes allocated value storage for deterministic state audits.
pub fn value_storage(&self) -> &[f64] {
&self.values
}
}
fn zeroed_storage(elements: usize) -> Result<Vec<f64>, LayerKvCacheError> {
let mut values = Vec::new();
values
.try_reserve_exact(elements)
.map_err(|_| LayerKvCacheError::AllocationFailed { elements })?;
values.resize(elements, 0.0);
Ok(values)
}
fn first_nonfinite(values: &[f64]) -> Option<(usize, f64)> {
values
.iter()
.copied()
.enumerate()
.find(|(_, value)| !value.is_finite())
}forward_incremental is the checked standalone entry: its caller supplies one
attention layer, one input row, and one cache directly. Before calculating
attention, it checks input rank and one-token shape, input batch and width, cache
model width and head geometry, parameter-node identities and value revisions,
the exact RoPE configuration, and remaining capacity, in that order. An
arbitrary standalone caller could otherwise pair a valid input with an unrelated
layer or cache.
After those checks, the crate-private prepare_incremental_bound function runs
the one shared calculation. Inside no_grad, it projects the new row, splits
heads, rotates and at the old cache length, computes numerically stable
weights shaped , mixes the retained and candidate values, merges
heads, and applies the existing output projection. It returns the complete
output together with the candidate rotated key and unrotated value. Only a later
commit copies those rows into the next cache slot and increments logical length.
The crate-private entry is not an unchecked public shortcut and does not contain a second attention algorithm. Its caller must first establish every condition listed above and preserve that exact layer/cache pairing until the prepared row is committed or discarded. Chapter 38 will establish those persistent relationships at its model-wide session boundary.
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)
}Across prefix lengths , , and , the executable history contrast records that the newest queries read key rows. A full-prefix call visits input rows in each , , or projection. The incremental path visits rows per projection. It therefore reuses earlier key rows and earlier value rows. The newest query still reads the resulting prefix of positions: retained positions plus the current candidate position. This count therefore does not claim constant-time attention or a measured speedup.
Errors remain transactional on both sides of the shared calculation. Projection,
finite-value, allocation, head-layout, and downstream tensor failures all occur
before commit. The deterministic demo exercises two-token input, a
full cache, model or head mismatch, same-shaped rebuilt weights, a changed RoPE
base, the same base with a different position capacity, and finite input whose
projection becomes nonfinite. Every one of those calls returns before cached
values or logical length change. Separately, the revision comparison at the
start of forward_incremental rejects a layer updated in place before projection
or append. That mismatch remains an error after reset, because reset does not
rebind the cache. Raw key/value append is private, so external callers cannot
bypass the checked standalone boundary.
rust/demos/ch37-incremental-attention/src/lib.rs#cache-errors fn error_evidence(layer: &MultiHeadAttention) -> Result<ErrorEvidence, FixtureError> {
let single = constant(&[1, 1, MODEL_WIDTH], &INPUT_VALUES[..MODEL_WIDTH])?;
let mut two_token_cache = LayerKvCache::new(layer, 1, 2)?;
let two_tokens_rejected = unchanged_after_error(
layer,
&constant(&[1, 2, MODEL_WIDTH], &INPUT_VALUES[..2 * MODEL_WIDTH])?,
&mut two_token_cache,
);
let mut full_cache = LayerKvCache::new(layer, 1, 1)?;
layer.forward_incremental(&single, &mut full_cache)?;
let full_cache_rejected = unchanged_after_error(layer, &single, &mut full_cache);
let mut rng = SplitMix64::from_seed(37);
let wider = MultiHeadAttention::new("wider", 8, 2, MAX_POSITIONS, ROPE_BASE, &mut rng)?;
let mut model_cache = LayerKvCache::new(&wider, 1, CAPACITY)?;
let model_mismatch_rejected = unchanged_after_error(layer, &single, &mut model_cache);
let one_head = MultiHeadAttention::new(
"one_head",
MODEL_WIDTH,
1,
MAX_POSITIONS,
ROPE_BASE,
&mut rng,
)?;
let mut head_cache = LayerKvCache::new(&one_head, 1, CAPACITY)?;
let head_mismatch_rejected = unchanged_after_error(layer, &single, &mut head_cache);
let other_layer = fixture_layer()?;
let mut other_cache = LayerKvCache::new(&other_layer, 1, CAPACITY)?;
let layer_mismatch_rejected = unchanged_after_error(layer, &single, &mut other_cache);
let different_rope = MultiHeadAttention::from_parameters(
layer.parameters()[0].clone(),
layer.parameters()[1].clone(),
layer.parameters()[2].clone(),
layer.parameters()[3].clone(),
HEADS,
MAX_POSITIONS,
ROPE_BASE * 2.0,
)?;
let mut rope_cache = LayerKvCache::new(layer, 1, CAPACITY)?;
let rope_mismatch_rejected = unchanged_after_error(&different_rope, &single, &mut rope_cache);
let different_positions = MultiHeadAttention::from_parameters(
layer.parameters()[0].clone(),
layer.parameters()[1].clone(),
layer.parameters()[2].clone(),
layer.parameters()[3].clone(),
HEADS,
MAX_POSITIONS + 1,
ROPE_BASE,
)?;
let mut position_cache = LayerKvCache::new(layer, 1, CAPACITY)?;
let rope_positions_mismatch_rejected =
unchanged_after_error(&different_positions, &single, &mut position_cache);
let mut nonfinite_cache = LayerKvCache::new(layer, 1, CAPACITY)?;
let nonfinite_projection_rejected = unchanged_after_error(
layer,
&constant(&[1, 1, MODEL_WIDTH], &[f64::MAX; MODEL_WIDTH])?,
&mut nonfinite_cache,
);
let every_cache_unchanged = two_tokens_rejected
&& full_cache_rejected
&& model_mismatch_rejected
&& head_mismatch_rejected
&& layer_mismatch_rejected
&& rope_mismatch_rejected
&& rope_positions_mismatch_rejected
&& nonfinite_projection_rejected;
Ok(ErrorEvidence {
two_tokens_rejected,
full_cache_rejected,
model_mismatch_rejected,
head_mismatch_rejected,
layer_mismatch_rejected,
rope_mismatch_rejected,
rope_positions_mismatch_rejected,
nonfinite_projection_rejected,
every_cache_unchanged,
})
}The fixture performs all three cached calls and all three independent reference calls. It checks full logical key/value prefixes as well as outputs, resets to length without replacing storage, and replays the same evidence exactly with the unchanged layer. That replay proves allocation reuse; it does not turn reset into a compatibility refresh after training.
rust/demos/ch37-incremental-attention/src/lib.rs#cache-step /// Runs three single-row appends, full-prefix references, reset, replay, and errors.
pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
let layer = fixture_layer()?;
let mut cache = LayerKvCache::new(&layer, 1, CAPACITY)?;
let steps = collect_steps(&layer, &mut cache)?;
let history = historical_kv_contrast(&steps)?;
let work = WorkEvidence {
full_rows_per_projection: steps.iter().map(|step| step.full_rows_per_projection).sum(),
incremental_rows_per_projection: steps
.iter()
.map(|step| step.incremental_rows_per_projection)
.sum(),
reused_rows_per_key_value_projection: steps
.iter()
.map(|step| step.reused_key_value_rows)
.sum(),
avoided_rows_across_key_and_value: 2 * steps
.iter()
.map(|step| step.reused_key_value_rows)
.sum::<usize>(),
};
let before_reset = cache.len();
let key_pointer = cache.key_storage().as_ptr();
let value_pointer = cache.value_storage().as_ptr();
let key_storage = cache.key_storage().to_vec();
let value_storage = cache.value_storage().to_vec();
cache.reset();
let reset_after = cache.len();
let allocation_reused = cache.key_storage().as_ptr() == key_pointer
&& cache.value_storage().as_ptr() == value_pointer;
let storage_unchanged =
cache.key_storage() == key_storage && cache.value_storage() == value_storage;
let replay = collect_steps(&layer, &mut cache)?;
let replay_identical = replay == steps;
let reset = ResetEvidence {
before: before_reset,
after: reset_after,
allocation_reused,
storage_unchanged,
replay_identical,
};
require(reset.after == 0, "cache reset did not return to zero")?;
require(
reset.allocation_reused && reset.storage_unchanged && reset.replay_identical,
"cache reset or replay evidence failed",
)?;
let errors = error_evidence(&layer)?;
require(
errors.every_cache_unchanged,
"one rejected operation changed cache state",
)?;
Ok(LearnerEvidence {
steps,
work,
reset,
errors,
history,
})
}Run cargo run --quiet --locked -p ch37-incremental-attention. The central
evidence is emitted directly by the executable:
step=position:0 cache:0->1 shape:[1,2,1,2] max_abs_diff:0.000000000000 output:[1.000000000,0.000000000,1.000000000,0.000000000]
step=position:1 cache:1->2 shape:[1,2,2,2] max_abs_diff:0.000000000000 output:[0.213809009,0.786190991,0.770151153,-0.420735492]
step=position:2 cache:2->3 shape:[1,2,3,2] max_abs_diff:0.000000000000 output:[0.629044078,0.945303958,0.374718490,-0.583589471]
work=full_rows_per_projection:6 incremental_rows_per_projection:3 reused_rows_per_kv_projection:3 avoided_rows_across_kv:6
reset=before:3 after:0 allocation_reused:true storage_unchanged:true replay_identical:true
errors=two_tokens:true full_cache:true model_mismatch:true head_mismatch:true layer_mismatch:true rope_mismatch:true rope_positions_mismatch:true nonfinite_projection:true unchanged:truerust/demos/ch37-incremental-attention/src/main.rs fn main() -> Result<(), Box<dyn std::error::Error>> {
print!("{}", ch37_incremental_attention::learner_report()?);
Ok(())
}Follow retained rows into the newest query
The figure follows the exact values produced by the three Rust calls. Read from absolute position through position . Solid-border boxes are retained rows; the double box is the one new key/value pair. Each row includes its newest-query weight, and each step ends with cached and full-prefix output values beside the recorded maximum difference.
The work cards compare projection-row counts, not elapsed time. The final cards show that reset reuses allocation and that rejected calls preserve state. A solid border marks a retained row; a double border marks the newly appended pair.
Retain earlier key/value rows; append exactly one new pair
The exact Rust trace follows three absolute positions, shows both head caches and attention weights, matches each newest output to a full-prefix reference, and records that reset plus rejected calls preserve storage.
- retained earlier row - solid border
- newly appended row - double border
- newest outputs match within tolerance
Advance the layer cache one position at a time
The old logical length is the zero-based RoPE position. Every head retains its earlier rows, adds one key/value pair, and lets the newest query read the complete resulting prefix.
Absolute position
Cache length
Logical cache shape:
Head
- newly appended row - double border Rotated keys: Unrotated values: Newest-query weights:
Head
- newly appended row - double border Rotated keys: Unrotated values: Newest-query weights:
Cached output:
Full-prefix reference:
newest outputs match within tolerance Maximum absolute difference:
Absolute position
Cache length
Logical cache shape:
Head
- retained earlier row - solid border Rotated keys: Unrotated values: Newest-query weights:
- newly appended row - double border Rotated keys: Unrotated values: Newest-query weights:
Head
- retained earlier row - solid border Rotated keys: Unrotated values: Newest-query weights:
- newly appended row - double border Rotated keys: Unrotated values: Newest-query weights:
Cached output:
Full-prefix reference:
newest outputs match within tolerance Maximum absolute difference:
Absolute position
Cache length
Logical cache shape:
Head
- retained earlier row - solid border Rotated keys: Unrotated values: Newest-query weights:
- retained earlier row - solid border Rotated keys: Unrotated values: Newest-query weights:
- newly appended row - double border Rotated keys: Unrotated values: Newest-query weights:
Head
- retained earlier row - solid border Rotated keys: Unrotated values: Newest-query weights:
- retained earlier row - solid border Rotated keys: Unrotated values: Newest-query weights:
- newly appended row - double border Rotated keys: Unrotated values: Newest-query weights:
Cached output:
Full-prefix reference:
newest outputs match within tolerance Maximum absolute difference:
Count projection rows, not hardware time
For prefix lengths one through three, the reference projects six rows per branch. The cached path projects three new rows and reuses three earlier key rows plus three earlier value rows.
Full-prefix reference
complete-prefix rows per query, key, or value projection
Cached output
incremental rows per query, key, or value projection
retained earlier row - solid border
earlier rows avoided across key and value projections
Reset or reject without corrupting retained state
Logical reset reuses the same allocation. Full, incompatible, or invalid calls leave both length and storage unchanged.
Reset keeps the allocation and stored values
allocation_reused=true storage_unchanged=true The same rows replay identically
replay_identical=true newest outputs match within tolerance
Rejected calls commit no append
two_tokens=true full_cache=true model_mismatch=true head_mismatch=true layer_mismatch=true rope_mismatch=true rope_positions_mismatch=true nonfinite_projection=true unchanged=true Predict before checking the trace
- Which RoPE offset is used when logical cache length is ?
- Does the layer cache retain the newest query?
- At the third call, how many keys can each newest query read?
- How many rows does one full-prefix key projection visit across lengths , , and ?
- How many rows does the incremental key projection visit?
- What does reset change, and does it refresh the parameter binding?
- Can a same-shaped cache from rebuilt weights be reused?
- Can the old cache be reused after AdamW updates the same parameter nodes in place?
- What happens to cache state when the capacity is already full?
Check the nine predictions
- The offset is , exactly the old logical length.
- No. Only rotated and unrotated rows persist.
- Each query reads keys and mixes values.
- The reference visits key-projection rows.
- The incremental path visits key-projection rows.
- Reset changes logical length to and keeps the allocation and stored values. It does not refresh the captured parameter identities or value revisions, and it does not rebind the cache.
- No. A rebuilt layer fails the parameter-node identity check even when shapes and numeric values agree.
- No. An in-place AdamW update preserves the parameter nodes but advances their
value revisions. The old cache returns
CacheLayerRevisionMismatch; create a new cache from the updated layer. - The call returns a typed error before stored values or logical length change.
Give every decoder block its own state next
One attention layer can now preserve graph-free rotated keys and unrotated values across decode steps and reproduce its full-prefix newest-position output. It uses one attention calculation behind two trust boundaries: the standalone public call proves that its input, layer, and cache belong together, while the crate-private path requires an owning caller to have proved the same facts already. Reset clears logical state without reallocating or rebinding.
The current cumulative decoder still evaluates full prefixes. Chapter 38 will bind one cache from each block’s actual attention layer to a model-wide session, thread those caches through prefill and one-token decode, and compare complete cached generation with the uncached reference from Chapter 36.