Один раз заполните кэши, затем декодируйте по одному токену

/// 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,
    })
}

/// 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)
}

/// 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(&current)
                    .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(&current, 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(&current)
                .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 })
}