oxibonsai-model 0.1.1

//! Quantized KV cache: INT8 per-row quantization for keys and values.
//!
//! Memory reduction: 4× vs FP32, 2× vs FP16.
//! Accuracy: ~0.1% error vs FP32 for typical activation ranges.
//!
//! # Layout
//! For each layer, each head, each token position:
//!   - keys_i8: [seq_len, num_kv_heads, head_dim] as i8
//!   - key_scales: [seq_len, num_kv_heads] as f32  (per-row scale)
//!   - values_i8: [seq_len, num_kv_heads, head_dim] as i8
//!   - value_scales: [seq_len, num_kv_heads] as f32

/// Error types for quantized KV cache operations.
#[derive(Debug, thiserror::Error)]
pub enum QuantKvError {
    #[error("capacity exceeded: capacity {capacity}, tried to push token {pos}")]
    CapacityExceeded { capacity: usize, pos: usize },

    #[error("token position {0} out of range")]
    PositionOutOfRange(usize),

    #[error("head index {head} out of range (num_kv_heads = {num_heads})")]
    HeadOutOfRange { head: usize, num_heads: usize },

    #[error("layer {layer} out of range (num_layers = {num_layers})")]
    LayerOutOfRange { layer: usize, num_layers: usize },

    #[error("key/value shape mismatch: expected {expected}, got {actual}")]
    ShapeMismatch { expected: usize, actual: usize },
}

// ─── Primitive quantization helpers ──────────────────────────────────────────

/// Quantize a slice to INT8 with a single per-row scale.
///
/// Returns `(quantized: Vec<i8>, scale: f32)`.
///
/// `scale = max(|x|) / 127.0`, clamped to at least [`f32::EPSILON`] to avoid
/// division-by-zero. All values are symmetrically clamped to `[-127, 127]` so
/// that rounding can never produce the asymmetric `i8::MIN` (-128).
pub fn quantize_row_i8(row: &[f32]) -> (Vec<i8>, f32) {
    if row.is_empty() {
        return (Vec::new(), f32::EPSILON);
    }

    let max_abs = row.iter().map(|x| x.abs()).fold(0.0_f32, f32::max);

    // Clamp scale to at least EPSILON to avoid division by zero for all-zero rows.
    let scale = (max_abs / 127.0_f32).max(f32::EPSILON);

    let quantized = row
        .iter()
        .map(|&x| (x / scale).round().clamp(-127.0, 127.0) as i8)
        .collect();

    (quantized, scale)
}

/// Dequantize INT8 back to f32 using the row scale.
///
/// Each element is simply multiplied by `scale`. If `scale` is zero or
/// near-zero the output will be all zeros, which is the correct representation
/// for an all-zero input row.
pub fn dequantize_row_i8(quantized: &[i8], scale: f32) -> Vec<f32> {
    quantized.iter().map(|&q| q as f32 * scale).collect()
}

/// Mean absolute error (MAE) between the original f32 slice and the
/// dequantized version of the quantized INT8 representation.
///
/// Returns `0.0` for an empty slice.
pub fn quant_error_mae(original: &[f32], quantized: &[i8], scale: f32) -> f32 {
    let n = original.len().min(quantized.len());
    if n == 0 {
        return 0.0;
    }
    let sum: f32 = original
        .iter()
        .zip(quantized.iter())
        .map(|(&o, &q)| (o - q as f32 * scale).abs())
        .sum();
    sum / n as f32
}

// ─── Per-layer quantized KV storage ──────────────────────────────────────────

/// A single layer's INT8-quantized KV cache.
///
/// Memory layout for the INT8 data arrays uses the token-major order
/// `[token_pos * num_kv_heads * head_dim]`, so sequential decode steps
/// append contiguous blocks. Scale arrays use `[token_pos * num_kv_heads]`.
#[derive(Debug)]
pub struct QuantizedKvLayer {
    /// Quantized key data: `[capacity * num_kv_heads * head_dim]` as i8.
    keys_i8: Vec<i8>,
    /// Per-row key scales: `[capacity * num_kv_heads]` as f32.
    key_scales: Vec<f32>,
    /// Quantized value data: `[capacity * num_kv_heads * head_dim]` as i8.
    values_i8: Vec<i8>,
    /// Per-row value scales: `[capacity * num_kv_heads]` as f32.
    value_scales: Vec<f32>,
    /// Number of KV attention heads.
    pub num_kv_heads: usize,
    /// Dimension of each attention head.
    pub head_dim: usize,
    /// Maximum number of token positions pre-allocated.
    pub capacity: usize,
    /// Number of token positions actually stored so far.
    pub len: usize,
}

impl QuantizedKvLayer {
    /// Allocate an empty quantized KV layer with the given dimensions.
    ///
    /// Pre-allocates all storage so that subsequent [`push`](Self::push) calls
    /// do not allocate.
    pub fn new(capacity: usize, num_kv_heads: usize, head_dim: usize) -> Self {
        let data_len = capacity * num_kv_heads * head_dim;
        let scale_len = capacity * num_kv_heads;

        Self {
            keys_i8: vec![0i8; data_len],
            key_scales: vec![0.0_f32; scale_len],
            values_i8: vec![0i8; data_len],
            value_scales: vec![0.0_f32; scale_len],
            num_kv_heads,
            head_dim,
            capacity,
            len: 0,
        }
    }

    /// Append keys and values for the next token position.
    ///
    /// `keys` must be a flat slice of shape `[num_kv_heads * head_dim]` (heads
    /// first, then dims). `values` must have the same shape.
    ///
    /// Each head's row is quantized independently with its own scale.
    ///
    /// # Errors
    /// - [`QuantKvError::CapacityExceeded`] if `self.len == self.capacity`.
    /// - [`QuantKvError::ShapeMismatch`] if `keys` or `values` length is wrong.
    pub fn push(&mut self, keys: &[f32], values: &[f32]) -> Result<(), QuantKvError> {
        let expected = self.num_kv_heads * self.head_dim;

        if keys.len() != expected {
            return Err(QuantKvError::ShapeMismatch {
                expected,
                actual: keys.len(),
            });
        }
        if values.len() != expected {
            return Err(QuantKvError::ShapeMismatch {
                expected,
                actual: values.len(),
            });
        }
        if self.len >= self.capacity {
            return Err(QuantKvError::CapacityExceeded {
                capacity: self.capacity,
                pos: self.len,
            });
        }

        let token_pos = self.len;

        for head in 0..self.num_kv_heads {
            let row_start = head * self.head_dim;
            let row_end = row_start + self.head_dim;

            // Compute offsets before any mutable borrows to satisfy the borrow checker.
            let data_off = self.data_offset(token_pos, head);
            let scale_off = self.scale_offset(token_pos, head);

            // Keys
            let key_row = &keys[row_start..row_end];
            let (kq, ks) = quantize_row_i8(key_row);
            self.keys_i8[data_off..data_off + self.head_dim].copy_from_slice(&kq);
            self.key_scales[scale_off] = ks;

            // Values
            let val_row = &values[row_start..row_end];
            let (vq, vs) = quantize_row_i8(val_row);
            self.values_i8[data_off..data_off + self.head_dim].copy_from_slice(&vq);
            self.value_scales[scale_off] = vs;
        }

        self.len += 1;
        Ok(())
    }

    /// Get dequantized keys for a specific token position and head.
    ///
    /// Returns a `Vec<f32>` of length `head_dim`.
    ///
    /// # Errors
    /// - [`QuantKvError::PositionOutOfRange`] if `token_pos >= self.len`.
    /// - [`QuantKvError::HeadOutOfRange`] if `head >= self.num_kv_heads`.
    pub fn get_key(&self, token_pos: usize, head: usize) -> Result<Vec<f32>, QuantKvError> {
        self.validate_pos_head(token_pos, head)?;
        let data_off = self.data_offset(token_pos, head);
        let scale = self.key_scales[self.scale_offset(token_pos, head)];
        Ok(dequantize_row_i8(
            &self.keys_i8[data_off..data_off + self.head_dim],
            scale,
        ))
    }

    /// Get dequantized values for a specific token position and head.
    ///
    /// Returns a `Vec<f32>` of length `head_dim`.
    ///
    /// # Errors
    /// - [`QuantKvError::PositionOutOfRange`] if `token_pos >= self.len`.
    /// - [`QuantKvError::HeadOutOfRange`] if `head >= self.num_kv_heads`.
    pub fn get_value(&self, token_pos: usize, head: usize) -> Result<Vec<f32>, QuantKvError> {
        self.validate_pos_head(token_pos, head)?;
        let data_off = self.data_offset(token_pos, head);
        let scale = self.value_scales[self.scale_offset(token_pos, head)];
        Ok(dequantize_row_i8(
            &self.values_i8[data_off..data_off + self.head_dim],
            scale,
        ))
    }

    /// Get all dequantized keys for a token position (all heads, interleaved).
    ///
    /// Returns a flat `Vec<f32>` of length `num_kv_heads * head_dim`.
    ///
    /// # Errors
    /// - [`QuantKvError::PositionOutOfRange`] if `token_pos >= self.len`.
    pub fn get_keys_at(&self, token_pos: usize) -> Result<Vec<f32>, QuantKvError> {
        if token_pos >= self.len {
            return Err(QuantKvError::PositionOutOfRange(token_pos));
        }
        let mut out = Vec::with_capacity(self.num_kv_heads * self.head_dim);
        for head in 0..self.num_kv_heads {
            let data_off = self.data_offset(token_pos, head);
            let scale = self.key_scales[self.scale_offset(token_pos, head)];
            out.extend(dequantize_row_i8(
                &self.keys_i8[data_off..data_off + self.head_dim],
                scale,
            ));
        }
        Ok(out)
    }

    /// Get all dequantized values for a token position (all heads, interleaved).
    ///
    /// Returns a flat `Vec<f32>` of length `num_kv_heads * head_dim`.
    ///
    /// # Errors
    /// - [`QuantKvError::PositionOutOfRange`] if `token_pos >= self.len`.
    pub fn get_values_at(&self, token_pos: usize) -> Result<Vec<f32>, QuantKvError> {
        if token_pos >= self.len {
            return Err(QuantKvError::PositionOutOfRange(token_pos));
        }
        let mut out = Vec::with_capacity(self.num_kv_heads * self.head_dim);
        for head in 0..self.num_kv_heads {
            let data_off = self.data_offset(token_pos, head);
            let scale = self.value_scales[self.scale_offset(token_pos, head)];
            out.extend(dequantize_row_i8(
                &self.values_i8[data_off..data_off + self.head_dim],
                scale,
            ));
        }
        Ok(out)
    }

    /// Memory used by this layer in bytes (INT8 data + f32 scales).
    ///
    /// Only accounts for the pre-allocated storage slabs, not struct overhead.
    pub fn memory_bytes(&self) -> usize {
        // INT8 data: 1 byte per element
        let data_bytes = self.keys_i8.len() + self.values_i8.len();
        // f32 scales: 4 bytes each
        let scale_bytes = (self.key_scales.len() + self.value_scales.len()) * 4;
        data_bytes + scale_bytes
    }

    /// Equivalent memory if the same data were stored as FP32 (no scales).
    ///
    /// `2 * capacity * num_kv_heads * head_dim * 4 bytes`
    pub fn fp32_memory_bytes(&self) -> usize {
        // Keys + values, each element 4 bytes
        2 * self.capacity * self.num_kv_heads * self.head_dim * 4
    }

    /// Compression ratio versus FP32 storage.
    ///
    /// Values approaching 4.0 indicate near-ideal INT8 compression. The ratio
    /// is slightly below 4.0 because per-row f32 scales add overhead.
    pub fn compression_ratio(&self) -> f32 {
        let quant = self.memory_bytes();
        if quant == 0 {
            return 1.0;
        }
        self.fp32_memory_bytes() as f32 / quant as f32
    }

    // ── Internal helpers ──────────────────────────────────────────────────────

    /// Flat index into the INT8 data arrays for `(token_pos, head, 0)`.
    ///
    /// Layout: `[token_pos][head][dim]` → `(token_pos * num_kv_heads + head) * head_dim`
    #[inline]
    fn data_offset(&self, token_pos: usize, head: usize) -> usize {
        (token_pos * self.num_kv_heads + head) * self.head_dim
    }

    /// Flat index into the scale arrays for `(token_pos, head)`.
    ///
    /// Layout: `[token_pos][head]` → `token_pos * num_kv_heads + head`
    #[inline]
    fn scale_offset(&self, token_pos: usize, head: usize) -> usize {
        token_pos * self.num_kv_heads + head
    }

    /// Validate that `token_pos < self.len` and `head < self.num_kv_heads`.
    fn validate_pos_head(&self, token_pos: usize, head: usize) -> Result<(), QuantKvError> {
        if token_pos >= self.len {
            return Err(QuantKvError::PositionOutOfRange(token_pos));
        }
        if head >= self.num_kv_heads {
            return Err(QuantKvError::HeadOutOfRange {
                head,
                num_heads: self.num_kv_heads,
            });
        }
        Ok(())
    }
}

// ─── Multi-layer quantized KV cache ──────────────────────────────────────────

/// Full multi-layer INT8-quantized KV cache for autoregressive decoding.
///
/// Wraps one [`QuantizedKvLayer`] per transformer layer and exposes a
/// unified decode-step interface through [`push_step`](Self::push_step).
#[derive(Debug)]
pub struct QuantizedKvCache {
    layers: Vec<QuantizedKvLayer>,
    /// Number of transformer layers.
    pub num_layers: usize,
    /// Number of KV attention heads per layer.
    pub num_kv_heads: usize,
    /// Dimension of each attention head.
    pub head_dim: usize,
}

impl QuantizedKvCache {
    /// Allocate a new quantized KV cache for `num_layers` transformer layers.
    ///
    /// Each layer is pre-allocated for `capacity` token positions.
    pub fn new(num_layers: usize, capacity: usize, num_kv_heads: usize, head_dim: usize) -> Self {
        let layers = (0..num_layers)
            .map(|_| QuantizedKvLayer::new(capacity, num_kv_heads, head_dim))
            .collect();

        Self {
            layers,
            num_layers,
            num_kv_heads,
            head_dim,
        }
    }

    /// Append KV tensors for all layers at the current decode step.
    ///
    /// `all_keys[layer]` must be a flat slice of shape `[num_kv_heads * head_dim]`.
    /// `all_values[layer]` must have the same shape.
    ///
    /// # Errors
    /// - [`QuantKvError::LayerOutOfRange`] if `all_keys.len() != self.num_layers`.
    /// - Propagates [`QuantKvError`] from each layer's [`push`](QuantizedKvLayer::push).
    pub fn push_step(
        &mut self,
        all_keys: &[Vec<f32>],
        all_values: &[Vec<f32>],
    ) -> Result<(), QuantKvError> {
        if all_keys.len() != self.num_layers {
            return Err(QuantKvError::LayerOutOfRange {
                layer: all_keys.len(),
                num_layers: self.num_layers,
            });
        }
        if all_values.len() != self.num_layers {
            return Err(QuantKvError::LayerOutOfRange {
                layer: all_values.len(),
                num_layers: self.num_layers,
            });
        }

        for (layer_idx, (layer, (keys, values))) in self
            .layers
            .iter_mut()
            .zip(all_keys.iter().zip(all_values.iter()))
            .enumerate()
        {
            layer.push(keys, values).map_err(|e| match e {
                // Re-attach layer context to capacity errors
                QuantKvError::CapacityExceeded { capacity, pos } => {
                    QuantKvError::CapacityExceeded { capacity, pos }
                }
                QuantKvError::ShapeMismatch { expected, actual } => {
                    QuantKvError::ShapeMismatch { expected, actual }
                }
                // Pass through other errors; we could enrich them with layer_idx
                // but the error types don't carry that field — keep as is.
                other => {
                    let _ = layer_idx;
                    other
                }
            })?;
        }
        Ok(())
    }

    /// Get dequantized keys for a specific layer, token position, and head.
    ///
    /// # Errors
    /// - [`QuantKvError::LayerOutOfRange`] if `layer >= self.num_layers`.
    /// - Propagates position/head errors from the underlying layer.
    pub fn get_key(
        &self,
        layer: usize,
        token_pos: usize,
        head: usize,
    ) -> Result<Vec<f32>, QuantKvError> {
        self.validate_layer(layer)?;
        self.layers[layer].get_key(token_pos, head)
    }

    /// Get dequantized values for a specific layer, token position, and head.
    ///
    /// # Errors
    /// - [`QuantKvError::LayerOutOfRange`] if `layer >= self.num_layers`.
    /// - Propagates position/head errors from the underlying layer.
    pub fn get_value(
        &self,
        layer: usize,
        token_pos: usize,
        head: usize,
    ) -> Result<Vec<f32>, QuantKvError> {
        self.validate_layer(layer)?;
        self.layers[layer].get_value(token_pos, head)
    }

    /// Total memory used across all layers in bytes.
    pub fn total_memory_bytes(&self) -> usize {
        self.layers.iter().map(|l| l.memory_bytes()).sum()
    }

    /// FP32-equivalent memory across all layers.
    pub fn total_fp32_memory_bytes(&self) -> usize {
        self.layers.iter().map(|l| l.fp32_memory_bytes()).sum()
    }

    /// Overall compression ratio vs FP32.
    pub fn compression_ratio(&self) -> f32 {
        let quant = self.total_memory_bytes();
        if quant == 0 {
            return 1.0;
        }
        self.total_fp32_memory_bytes() as f32 / quant as f32
    }

    /// Number of token positions currently stored (taken from layer 0).
    ///
    /// Returns `0` if there are no layers.
    pub fn seq_len(&self) -> usize {
        self.layers.first().map(|l| l.len).unwrap_or(0)
    }

    // ── Internal helpers ──────────────────────────────────────────────────────

    fn validate_layer(&self, layer: usize) -> Result<(), QuantKvError> {
        if layer >= self.num_layers {
            return Err(QuantKvError::LayerOutOfRange {
                layer,
                num_layers: self.num_layers,
            });
        }
        Ok(())
    }
}