geographdb_core/algorithms/

kv_cache_mps.rs

//! MPS-inspired KV Cache via Online Low-Rank Basis Factorisation.
//!
//! Represents the KV history `[(K_0,V_0), …, (K_{T-1},V_{T-1})]` as a
//! rank-`χ` factorisation:
//!
//!   K ≈ B_k · C_k^T     (d_k × T via d_k×χ basis · T×χ coefficients)
//!   V ≈ B_v · C_v^T
//!
//! This is the canonical 1-D MPS with bond dimension χ viewed in its SVD
//! decomposed form: `B_k` contains the left singular vectors (physical-index
//! basis), and `C_k` contains the right-singular-vector projections scaled by
//! singular values (the "bond" coordinates of each token).
//!
//! **Memory:** O(χ · d_kv + T · χ)  vs O(T · d_kv) flat.
//! For χ << d_kv and T large the storage ratio ≈ d_kv / χ.
//!
//! **Append:** Gram-Schmidt extends each basis independently if the current
//! rank < chi_max; otherwise the new token is projected (lossy).
//!
//! **Attend:** For each token, approximate K[t] ≈ B_k · c_k[t] and
//! V[t] ≈ B_v · c_v[t]. Compute softmax(Q·K^T) · V entirely in basis space
//! without materialising individual KV vectors.

// ── Public API ────────────────────────────────────────────────────────────────

/// MPS-compressed KV cache backed by independent online low-rank bases for
/// keys and values.
pub struct KvCacheMps {
    /// Key basis, column-major: `basis_k[col * d_k + row]`.  Shape d_k × chi_k.
    basis_k: Vec<f32>,
    /// Value basis, column-major: `basis_v[col * d_v + row]`. Shape d_v × chi_v.
    basis_v: Vec<f32>,
    /// Per-token key coefficients: `coeff_k[t]` has length `chi_k` at time t.
    coeff_k: Vec<Vec<f32>>,
    /// Per-token value coefficients: `coeff_v[t]` has length `chi_v` at time t.
    coeff_v: Vec<Vec<f32>>,
    /// Current number of key basis vectors (≤ chi_max).
    chi_k: usize,
    /// Current number of value basis vectors (≤ chi_max).
    chi_v: usize,
    chi_max: usize,
    d_k: usize,
    d_v: usize,
}

impl KvCacheMps {
    /// Create an empty cache.
    ///
    /// `d_k` / `d_v`: key / value head dimensions.
    /// `chi_max`: maximum bond dimension (compression rank).
    pub fn new(d_k: usize, d_v: usize, chi_max: usize) -> Self {
        assert!(chi_max >= 1, "chi_max must be ≥ 1");
        Self {
            basis_k: Vec::new(),
            basis_v: Vec::new(),
            coeff_k: Vec::new(),
            coeff_v: Vec::new(),
            chi_k: 0,
            chi_v: 0,
            chi_max,
            d_k,
            d_v,
        }
    }

    /// Append a new `(k, v)` token pair.
    ///
    /// If the respective basis is not yet full (`chi < chi_max`), a new
    /// orthonormal direction is added via Gram-Schmidt; otherwise the token is
    /// projected onto the existing basis (lossy compression).
    pub fn append(&mut self, k: &[f32], v: &[f32]) {
        assert_eq!(k.len(), self.d_k);
        assert_eq!(v.len(), self.d_v);

        let c_k = project_and_extend(
            &mut self.basis_k,
            &mut self.chi_k,
            self.chi_max,
            self.d_k,
            k,
        );
        self.coeff_k.push(c_k);

        let c_v = project_and_extend(
            &mut self.basis_v,
            &mut self.chi_v,
            self.chi_max,
            self.d_v,
            v,
        );
        self.coeff_v.push(c_v);
    }

    /// Number of tokens stored.
    pub fn token_count(&self) -> usize {
        self.coeff_k.len()
    }

    /// Current bond dimension (max of key and value ranks; ≤ chi_max).
    pub fn max_bond_dim(&self) -> usize {
        self.chi_k.max(self.chi_v)
    }

    /// Compute softmax attention over the compressed KV cache.
    ///
    /// `query` has length `d_k`. Returns a vector of length `d_v`.
    ///
    /// Scores are computed as `(B_k^T · query) · c_k[t]`, which equals the
    /// dot product `query · K_approx[t]` exactly (the basis is orthonormal).
    pub fn attend(&self, query: &[f32], scale: f32) -> Vec<f32> {
        let n = self.token_count();
        if n == 0 {
            return vec![0.0; self.d_v];
        }

        // Compress the query into the key basis: q_comp[k] = basis_k[:,k] · query
        let q_comp: Vec<f32> = (0..self.chi_k)
            .map(|k| {
                let col = &self.basis_k[k * self.d_k..(k + 1) * self.d_k];
                col.iter().zip(query).map(|(b, q)| b * q).sum()
            })
            .collect();

        // Scores: score[t] = (q_comp · c_k[t]) * scale.
        // Older tokens may have shorter c_k (basis grew after they were added);
        // zip stops at the shorter length, which correctly skips basis vectors
        // that are orthogonal to the token (their true coefficient is 0).
        let mut scores: Vec<f32> = self
            .coeff_k
            .iter()
            .map(|c| {
                q_comp
                    .iter()
                    .zip(c.iter())
                    .map(|(q, ci)| q * ci)
                    .sum::<f32>()
                    * scale
            })
            .collect();

        // Softmax
        let max_s = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let mut weights: Vec<f32> = scores.iter_mut().map(|s| (*s - max_s).exp()).collect();
        let sum_w: f32 = weights.iter().sum();
        for w in &mut weights {
            *w /= sum_w;
        }

        // Weighted V sum: out = Σ_t w[t] * (B_v · c_v[t])
        let mut out = vec![0.0f32; self.d_v];
        for (w, c_v) in weights.iter().zip(self.coeff_v.iter()) {
            for (k, &ck) in c_v.iter().enumerate() {
                let wck = w * ck;
                let col = &self.basis_v[k * self.d_v..(k + 1) * self.d_v];
                for (o, b) in out.iter_mut().zip(col) {
                    *o += wck * b;
                }
            }
        }
        out
    }

    /// Compression ratio: flat storage bytes / compressed storage bytes.
    ///
    /// Values > 1 mean the cache uses less memory than storing KV pairs flat.
    /// The ratio grows with T; for large T it approaches `d_kv / chi_max`.
    pub fn compression_ratio(&self) -> f64 {
        let n = self.token_count();
        if n == 0 {
            return 1.0;
        }
        let flat = n * (self.d_k + self.d_v);
        // Basis storage + per-token coefficient storage
        let compressed =
            self.chi_k * self.d_k + self.chi_v * self.d_v + n * self.chi_k + n * self.chi_v;
        flat as f64 / compressed as f64
    }
}

// ── Gram-Schmidt helpers ──────────────────────────────────────────────────────

/// Project `vec` onto the current columns of `basis` (d × chi, column-major),
/// and — if `chi < chi_max` and the residual is non-trivial — extend the basis
/// with the normalised residual direction.
///
/// Returns the coefficient vector (length = chi after any extension).
fn project_and_extend(
    basis: &mut Vec<f32>,
    chi: &mut usize,
    chi_max: usize,
    d: usize,
    vec: &[f32],
) -> Vec<f32> {
    // Project onto existing columns
    let mut coeff: Vec<f32> = (0..*chi)
        .map(|k| {
            let col = &basis[k * d..(k + 1) * d];
            col.iter().zip(vec).map(|(b, v)| b * v).sum()
        })
        .collect();

    if *chi >= chi_max {
        return coeff;
    }

    // Residual: r = vec − Σ_k coeff[k] * basis[:,k]
    let mut residual = vec.to_vec();
    for k in 0..*chi {
        let c = coeff[k];
        let col = &basis[k * d..(k + 1) * d];
        for (r, b) in residual.iter_mut().zip(col) {
            *r -= c * b;
        }
    }

    let norm: f32 = residual.iter().map(|x| x * x).sum::<f32>().sqrt();
    // Relative threshold: skip extension if the residual is tiny compared to the
    // input. f32 Gram-Schmidt accumulates ~O(chi * macheps) relative error, so
    // the residual of a truly in-span vector can reach ~1e-5 × ‖vec‖ after several
    // steps. Using 1e-4 provides a generous safety margin without confusing
    // genuinely small new directions with numerical noise.
    let vec_norm: f32 = vec.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm < 1e-4 * vec_norm.max(1e-12) {
        // vec is already in the span of the basis
        return coeff;
    }

    // Add normalised residual as a new basis column
    for r in residual.iter_mut() {
        *r /= norm;
    }
    basis.extend_from_slice(&residual);
    *chi += 1;

    // The new coefficient for this token in the new dimension equals the
    // pre-normalisation residual norm.
    coeff.push(norm);
    coeff
}

// ── Truncated SVD (retained for recompression experiments) ───────────────────

/// Rank-`rank` truncated SVD of `m` (rows × cols, row-major).
/// Returns `(U [rows×rank], sigma [rank], Vt [rank×cols], rank_kept)`.
pub fn svd_truncated(
    m: &[f32],
    rows: usize,
    cols: usize,
    rank: usize,
) -> (Vec<f32>, Vec<f32>, Vec<f32>, usize) {
    let rank = rank.min(rows).min(cols);
    if rows <= cols {
        svd_via_gram_left(m, rows, cols, rank)
    } else {
        svd_via_gram_right(m, rows, cols, rank)
    }
}

fn svd_via_gram_left(
    m: &[f32],
    rows: usize,
    cols: usize,
    rank: usize,
) -> (Vec<f32>, Vec<f32>, Vec<f32>, usize) {
    let mut g = vec![0.0f32; rows * rows];
    for i in 0..rows {
        for k in 0..cols {
            let mik = m[i * cols + k];
            for j in 0..rows {
                g[i * rows + j] += mik * m[j * cols + k];
            }
        }
    }
    let (eigvecs, eigvals) = power_iteration_symmetric(&g, rows, rank);

    const SVD_EPS: f32 = 1e-7;
    let sigma_max = eigvals[0].max(0.0).sqrt();
    let mut rank_kept = rank;
    while rank_kept > 1 && eigvals[rank_kept - 1].max(0.0).sqrt() < SVD_EPS * sigma_max {
        rank_kept -= 1;
    }
    rank_kept = rank_kept.max(1);

    let sigma: Vec<f32> = (0..rank_kept).map(|k| eigvals[k].max(0.0).sqrt()).collect();

    let mut u = vec![0.0f32; rows * rank_kept];
    for i in 0..rows {
        for k in 0..rank_kept {
            u[i * rank_kept + k] = eigvecs[i * rank + k];
        }
    }

    let mut vt = vec![0.0f32; rank_kept * cols];
    for k in 0..rank_kept {
        if sigma[k] < 1e-15 {
            continue;
        }
        for j in 0..cols {
            let mut val = 0.0f32;
            for i in 0..rows {
                val += m[i * cols + j] * u[i * rank_kept + k];
            }
            vt[k * cols + j] = val / sigma[k];
        }
    }
    (u, sigma, vt, rank_kept)
}

fn svd_via_gram_right(
    m: &[f32],
    rows: usize,
    cols: usize,
    rank: usize,
) -> (Vec<f32>, Vec<f32>, Vec<f32>, usize) {
    let mut g = vec![0.0f32; cols * cols];
    for k in 0..rows {
        for i in 0..cols {
            let mki = m[k * cols + i];
            for j in 0..cols {
                g[i * cols + j] += mki * m[k * cols + j];
            }
        }
    }
    let (eigvecs, eigvals) = power_iteration_symmetric(&g, cols, rank);

    const SVD_EPS: f32 = 1e-7;
    let sigma_max = eigvals[0].max(0.0).sqrt();
    let mut rank_kept = rank;
    while rank_kept > 1 && eigvals[rank_kept - 1].max(0.0).sqrt() < SVD_EPS * sigma_max {
        rank_kept -= 1;
    }
    rank_kept = rank_kept.max(1);

    let sigma: Vec<f32> = (0..rank_kept).map(|k| eigvals[k].max(0.0).sqrt()).collect();

    let mut v = vec![0.0f32; cols * rank_kept];
    for i in 0..cols {
        for k in 0..rank_kept {
            v[i * rank_kept + k] = eigvecs[i * rank + k];
        }
    }

    let mut vt = vec![0.0f32; rank_kept * cols];
    for k in 0..rank_kept {
        for j in 0..cols {
            vt[k * cols + j] = v[j * rank_kept + k];
        }
    }

    let mut u = vec![0.0f32; rows * rank_kept];
    for k in 0..rank_kept {
        if sigma[k] < 1e-15 {
            continue;
        }
        for i in 0..rows {
            let mut val = 0.0f32;
            for j in 0..cols {
                val += m[i * cols + j] * v[j * rank_kept + k];
            }
            u[i * rank_kept + k] = val / sigma[k];
        }
    }
    (u, sigma, vt, rank_kept)
}

fn power_iteration_symmetric(g: &[f32], n: usize, rank: usize) -> (Vec<f32>, Vec<f32>) {
    const MAX_ITER: usize = 64;
    const TOL: f32 = 1e-6;

    let mut eigvecs = vec![0.0f32; n * rank];
    let mut eigvals = vec![0.0f32; rank];
    let mut deflated = g.to_vec();

    for k in 0..rank {
        let mut v: Vec<f32> = (0..n).map(|i| if i == k % n { 1.0 } else { 0.0 }).collect();
        normalise(&mut v);
        let mut lambda = 0.0f32;
        for _ in 0..MAX_ITER {
            let w = matvec_sq(&deflated, &v, n);
            let lambda_new: f32 = v.iter().zip(w.iter()).map(|(vi, wi)| vi * wi).sum();
            let mut w2 = w;
            normalise(&mut w2);
            let diff: f32 = v.iter().zip(w2.iter()).map(|(a, b)| (a - b).abs()).sum();
            v = w2;
            lambda = lambda_new;
            if diff < TOL {
                break;
            }
        }
        eigvals[k] = lambda;
        for i in 0..n {
            eigvecs[i * rank + k] = v[i];
        }
        for i in 0..n {
            for j in 0..n {
                deflated[i * n + j] -= lambda * v[i] * v[j];
            }
        }
    }
    (eigvecs, eigvals)
}

fn matvec_sq(m: &[f32], x: &[f32], n: usize) -> Vec<f32> {
    let mut y = vec![0.0f32; n];
    for i in 0..n {
        for j in 0..n {
            y[i] += m[i * n + j] * x[j];
        }
    }
    y
}

fn normalise(v: &mut [f32]) {
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 1e-12 {
        for x in v.iter_mut() {
            *x /= norm;
        }
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;

    fn vec_norm(v: &[f32]) -> f32 {
        v.iter().map(|x| x * x).sum::<f32>().sqrt()
    }
    fn vec_err(a: &[f32], b: &[f32]) -> f32 {
        a.iter()
            .zip(b)
            .map(|(x, y)| (x - y).powi(2))
            .sum::<f32>()
            .sqrt()
    }
    fn approx_eq(a: f32, b: f32, tol: f32) -> bool {
        (a - b).abs() < tol
    }

    // ── svd_truncated ─────────────────────────────────────────────────────────

    #[test]
    fn test_svd_rank1_reconstruction() {
        let u_vec = [1.0f32, 0.0, 0.0];
        let s = 3.0f32;
        let vt_vec = [0.0f32, 1.0, 0.0, 0.0];
        let mut m = vec![0.0f32; 3 * 4];
        for i in 0..3 {
            for j in 0..4 {
                m[i * 4 + j] = u_vec[i] * s * vt_vec[j];
            }
        }
        let (u_out, sigma_out, vt_out, rank) = svd_truncated(&m, 3, 4, 2);
        assert!(rank >= 1);
        assert!(approx_eq(sigma_out[0], s, 0.1), "sigma={}", sigma_out[0]);
        let mut recon = vec![0.0f32; 3 * 4];
        for k in 0..rank {
            for i in 0..3 {
                for j in 0..4 {
                    recon[i * 4 + j] += u_out[i * rank + k] * sigma_out[k] * vt_out[k * 4 + j];
                }
            }
        }
        let err = vec_err(&recon, &m) / (vec_norm(&m) + 1e-8);
        assert!(err < 0.05, "reconstruction error = {err:.4}");
    }

    #[test]
    fn test_svd_identity_singular_values() {
        let mut eye = vec![0.0f32; 16];
        for i in 0..4 {
            eye[i * 4 + i] = 1.0;
        }
        let (_, sigma, _, rank) = svd_truncated(&eye, 4, 4, 4);
        assert_eq!(rank, 4);
        for s in &sigma {
            assert!(approx_eq(*s, 1.0, 0.1), "sigma={s}");
        }
    }

    // ── KvCacheMps ────────────────────────────────────────────────────────────

    #[test]
    fn test_empty_cache_attend_zero() {
        let cache = KvCacheMps::new(4, 4, 8);
        let out = cache.attend(&[1.0, 0.0, 0.0, 0.0], 1.0);
        assert_eq!(out, vec![0.0; 4]);
    }

    #[test]
    fn test_single_token_attend_returns_value() {
        let mut cache = KvCacheMps::new(4, 4, 8);
        let k = vec![1.0f32, 0.0, 0.0, 0.0];
        let v = vec![0.0f32, 0.0, 1.0, 0.0];
        cache.append(&k, &v);
        let out = cache.attend(&[1.0, 0.0, 0.0, 0.0], 1.0);
        // Single token: softmax weight = 1.0 → output exactly equals v
        assert!(vec_err(&out, &v) < 1e-4, "expected {v:?}, got {out:?}");
    }

    #[test]
    fn test_token_count_increments() {
        let mut cache = KvCacheMps::new(4, 4, 8);
        for i in 0..5 {
            cache.append(&[i as f32, 0.0, 0.0, 0.0], &[0.0, i as f32, 0.0, 0.0]);
            assert_eq!(cache.token_count(), i + 1);
        }
    }

    #[test]
    fn test_chi_max_bounds_bond_dimension() {
        let chi_max = 4;
        let mut cache = KvCacheMps::new(8, 8, chi_max);
        for i in 0..32 {
            let k: Vec<f32> = (0..8).map(|j| ((i + j) as f32) * 0.1).collect();
            let v: Vec<f32> = (0..8).map(|j| ((i * 2 + j) as f32) * 0.1).collect();
            cache.append(&k, &v);
        }
        assert!(
            cache.max_bond_dim() <= chi_max,
            "bond dim {} > chi_max {chi_max}",
            cache.max_bond_dim()
        );
    }

    #[test]
    fn test_compression_ratio_exceeds_one() {
        // chi_max=1, T=20, d_k=d_v=8:
        //   flat    = 20 * 16 = 320
        //   compressed = chi_k*d_k + chi_v*d_v + T*chi_k + T*chi_v
        //              = 1*8 + 1*8 + 20*1 + 20*1 = 56
        //   ratio   = 320/56 ≈ 5.7
        let mut cache = KvCacheMps::new(8, 8, 1);
        for i in 0..20 {
            let k: Vec<f32> = (0..8).map(|j| (i + j) as f32).collect();
            let v: Vec<f32> = (0..8).map(|j| (i * 2 + j) as f32).collect();
            cache.append(&k, &v);
        }
        let ratio = cache.compression_ratio();
        assert!(ratio > 1.0, "expected ratio > 1, got {ratio:.3}");
    }

    #[test]
    fn test_higher_chi_lower_attend_error() {
        // Diverse KV sequence with many orthogonal-ish directions
        let d = 8;
        let tokens: Vec<(Vec<f32>, Vec<f32>)> = (0..16)
            .map(|i| {
                let k: Vec<f32> = (0..d)
                    .map(|j| (((i * 3 + j * 7) % 11) as f32 - 5.0) * 0.3)
                    .collect();
                let v: Vec<f32> = (0..d)
                    .map(|j| (((i * 5 + j * 3) % 7) as f32 - 3.0) * 0.2)
                    .collect();
                (k, v)
            })
            .collect();

        let query: Vec<f32> = (0..d).map(|i| i as f32 * 0.1).collect();
        let scale = 1.0 / (d as f32).sqrt();

        // Reference: exact full-rank attention
        let ref_out = direct_attend(&tokens, &query, scale, d);

        let error_for_chi = |chi: usize| {
            let mut cache = KvCacheMps::new(d, d, chi);
            for (k, v) in &tokens {
                cache.append(k, v);
            }
            let out = cache.attend(&query, scale);
            vec_err(&out, &ref_out) / (vec_norm(&ref_out) + 1e-8)
        };

        let err1 = error_for_chi(1);
        let err8 = error_for_chi(8);

        // chi=8 must not be much worse than chi=1 (more rank ≥ less rank)
        assert!(
            err8 <= err1 + 0.1,
            "chi=8 error {err8:.4} should be ≤ chi=1 error {err1:.4} + 0.1"
        );
        // Full-rank capture (chi=d) should give near-exact attention
        let err_full = error_for_chi(d);
        assert!(
            err_full < 0.02,
            "full-rank attend error {err_full:.4} should be < 2%"
        );
    }

    /// Reference: exact direct attention without compression (test helper only).
    fn direct_attend(
        tokens: &[(Vec<f32>, Vec<f32>)],
        query: &[f32],
        scale: f32,
        d: usize,
    ) -> Vec<f32> {
        let scores_raw: Vec<f32> = tokens
            .iter()
            .map(|(k, _)| k.iter().zip(query).map(|(ki, qi)| ki * qi).sum::<f32>() * scale)
            .collect();
        let max_s = scores_raw.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let exp: Vec<f32> = scores_raw.iter().map(|s| (s - max_s).exp()).collect();
        let sum_exp: f32 = exp.iter().sum();
        let weights: Vec<f32> = exp.iter().map(|e| e / sum_exp).collect();
        let mut out = vec![0.0f32; d];
        for (w, (_, v)) in weights.iter().zip(tokens.iter()) {
            for (o, vi) in out.iter_mut().zip(v.iter()) {
                *o += w * vi;
            }
        }
        out
    }
}