inferencelayer 0.2.3

//! TurboQuant KV-cache quantization (Zandieh–Daliri–Hadian–Mirrokni, arXiv 2504.19874):
//! calibration-free, data-oblivious vector quantization built from (1) a fixed random rotation
//! `Π` (orthogonal — inner products are preserved in the rotated domain, so attention never
//! un-rotates keys: it rotates the QUERY once and dots against quantized coordinates), and (2) a
//! per-coordinate Lloyd–Max quantizer whose centroids are precomputed ONCE on the analytic
//! high-dimensional coordinate density (≈ N(0, 1/d) for unit vectors) — no calibration data.
//!
//! Stage 1 (`quantize_mse`) is the MSE-optimal path used for VALUES (and, at 4 bits, a strong
//! baseline for keys). Stage 2 (`quantize_unbiased`) adds a 1-bit QJL sign sketch of the residual,
//! making inner-product estimates UNBIASED — the paper's key path. The GPU kernels mirror stage 1
//! bit-exactly (same centroid table, same nearest-centroid rule); stage 2's CPU reference and its
//! statistical gates live here so the GPU port has a pinned spec.
//!
//! Distortion bounds verified by the unit gates (paper Thm. 1): per-vector relative MSE
//! `≤ 2.7 · 4^-b` (within ~2.7× of the optimal `4^-b`).

/// Lloyd–Max centroids for a STANDARD NORMAL coordinate, `2^bits` levels, computed by the classic
/// fixed-point iteration (centroid = conditional mean of its Voronoi cell) on the analytic density.
/// Deterministic (no data, no RNG); callers scale by the coordinate std `1/√d · ‖x‖`.
pub fn lloyd_max_centroids(bits: u32) -> Vec<f32> {
    let k = 1usize << bits;
    // Fine grid over ±6σ with Gaussian weights (analytic, not sampled).
    const N: usize = 24_001;
    let (lo, hi) = (-6.0f64, 6.0f64);
    let step = (hi - lo) / (N - 1) as f64;
    let xs: Vec<f64> = (0..N).map(|i| lo + i as f64 * step).collect();
    let ws: Vec<f64> = xs.iter().map(|x| (-0.5 * x * x).exp()).collect();
    // Init: uniform quantiles of the grid mass.
    let total: f64 = ws.iter().sum();
    let mut cents: Vec<f64> = {
        let mut acc = 0.0;
        let mut c = Vec::with_capacity(k);
        let mut next = total / (2.0 * k as f64);
        for (x, w) in xs.iter().zip(&ws) {
            acc += w;
            while c.len() < k && acc >= next {
                c.push(*x);
                next += total / k as f64;
            }
        }
        while c.len() < k {
            c.push(hi);
        }
        c
    };
    for _ in 0..200 {
        // Voronoi boundaries = midpoints; new centroid = cell's conditional mean.
        let mut sums = vec![0.0f64; k];
        let mut mass = vec![0.0f64; k];
        let mut cell = 0usize;
        for (x, w) in xs.iter().zip(&ws) {
            while cell + 1 < k && *x > 0.5 * (cents[cell] + cents[cell + 1]) {
                cell += 1;
            }
            sums[cell] += w * x;
            mass[cell] += w;
        }
        let mut moved = 0.0f64;
        for i in 0..k {
            if mass[i] > 0.0 {
                let nc = sums[i] / mass[i];
                moved = moved.max((nc - cents[i]).abs());
                cents[i] = nc;
            }
        }
        // reset scan cell order (cents stay sorted by construction of the iteration)
        if moved < 1e-12 {
            break;
        }
    }
    cents.iter().map(|c| *c as f32).collect()
}

/// Deterministic orthogonal rotation `Π ∈ R^{d×d}` (row-major): QR (modified Gram–Schmidt, f64) of
/// a seeded Gaussian matrix. Fixed per (d, seed) — precomputed once, shared by quantizer and
/// attention (which rotates queries with the SAME matrix).
pub fn orthogonal_rotation(d: usize, seed: u64) -> Vec<f32> {
    // xorshift64* → Box–Muller normals, f64 for QR stability.
    let mut state = seed | 1;
    let mut next_u = move || {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        (state.wrapping_mul(0x2545F4914F6CDD1D) >> 11) as f64 / (1u64 << 53) as f64
    };
    let mut g = vec![0.0f64; d * d];
    let mut i = 0;
    while i < d * d {
        let (u1, u2) = (next_u().max(1e-300), next_u());
        let r = (-2.0 * u1.ln()).sqrt();
        let (s, c) = (2.0 * std::f64::consts::PI * u2).sin_cos();
        g[i] = r * c;
        if i + 1 < d * d {
            g[i + 1] = r * s;
        }
        i += 2;
    }
    // Modified Gram–Schmidt on rows.
    for r in 0..d {
        for p in 0..r {
            let dot: f64 = (0..d).map(|j| g[r * d + j] * g[p * d + j]).sum();
            for j in 0..d {
                g[r * d + j] -= dot * g[p * d + j];
            }
        }
        let nrm: f64 = (0..d)
            .map(|j| g[r * d + j] * g[r * d + j])
            .sum::<f64>()
            .sqrt();
        for j in 0..d {
            g[r * d + j] /= nrm;
        }
    }
    g.iter().map(|v| *v as f32).collect()
}

/// Stage-1 (MSE) quantization of one vector: rotate, normalize, nearest Lloyd–Max centroid per
/// coordinate. Returns (4-bit indices packed 8/u32 in coordinate order, per-vector scale
/// `‖x‖/√d`). Dequantized coordinate `j` (ROTATED domain) = `centroid[idx_j] · scale`.
pub fn quantize_mse(x: &[f32], rot: &[f32], cents: &[f32]) -> (Vec<u32>, f32) {
    let d = x.len();
    let norm = x.iter().map(|v| (*v as f64).powi(2)).sum::<f64>().sqrt() as f32;
    let scale = norm / (d as f32).sqrt();
    let inv = if scale > 0.0 { 1.0 / scale } else { 0.0 };
    let mut packed = vec![0u32; d.div_ceil(8)];
    for j in 0..d {
        // y_j = (Π·x)_j, normalized to the standard-normal domain of the centroid table.
        let y: f32 = (0..d).map(|t| rot[j * d + t] * x[t]).sum();
        let z = y * inv;
        let mut best = 0u32;
        let mut bd = f32::INFINITY;
        for (ci, c) in cents.iter().enumerate() {
            let dd = (z - c).abs();
            if dd < bd {
                bd = dd;
                best = ci as u32;
            }
        }
        packed[j / 8] |= (best & 0xF) << ((j % 8) * 4);
    }
    (packed, scale)
}

/// Dequantize stage-1 coordinates back to the ROTATED domain (attention dots against these).
pub fn dequant_rotated(packed: &[u32], scale: f32, cents: &[f32], d: usize) -> Vec<f32> {
    (0..d)
        .map(|j| cents[((packed[j / 8] >> ((j % 8) * 4)) & 0xF) as usize] * scale)
        .collect()
}

/// Stage-2 (unbiased inner product) CPU reference: (b−1)-bit stage-1 + 1-bit QJL sign sketch of the
/// residual. Returns (mse indices, sign bits packed 32/u32, residual norm γ, scale). Estimator:
/// `x̃ = x̃_mse + (√(π/2)/d)·γ·Sᵀ·q` — `E[⟨y, x̃⟩] = ⟨y, x⟩` (verified statistically by the gate).
pub fn quantize_unbiased(
    x: &[f32],
    rot: &[f32],
    sketch: &[f32],
    cents_lo: &[f32],
) -> (Vec<u32>, Vec<u32>, f32, f32) {
    let d = x.len();
    let (packed, scale) = quantize_mse_with(x, rot, cents_lo);
    let deq_rot = dequant_rotated_any(&packed, scale, cents_lo, d);
    // Residual in the ORIGINAL domain: r = x − Πᵀ·ŷ.
    let mut r = vec![0.0f32; d];
    for t in 0..d {
        let back: f32 = (0..d).map(|j| rot[j * d + t] * deq_rot[j]).sum();
        r[t] = x[t] - back;
    }
    let gamma = r.iter().map(|v| (*v as f64).powi(2)).sum::<f64>().sqrt() as f32;
    let mut signs = vec![0u32; d.div_ceil(32)];
    for j in 0..d {
        let s: f32 = (0..d).map(|t| sketch[j * d + t] * r[t]).sum();
        if s >= 0.0 {
            signs[j / 32] |= 1 << (j % 32);
        }
    }
    (packed, signs, gamma, scale)
}

/// Stage-2 dequantization to the ORIGINAL domain.
#[allow(clippy::too_many_arguments)]
pub fn dequant_unbiased(
    packed: &[u32],
    signs: &[u32],
    gamma: f32,
    scale: f32,
    rot: &[f32],
    sketch: &[f32],
    cents_lo: &[f32],
    d: usize,
) -> Vec<f32> {
    let deq_rot = dequant_rotated_any(packed, scale, cents_lo, d);
    let coef = (std::f64::consts::PI / 2.0).sqrt() as f32 / d as f32 * gamma;
    (0..d)
        .map(|t| {
            let base: f32 = (0..d).map(|j| rot[j * d + t] * deq_rot[j]).sum();
            let corr: f32 = (0..d)
                .map(|j| {
                    let sgn = if (signs[j / 32] >> (j % 32)) & 1 == 1 {
                        1.0
                    } else {
                        -1.0
                    };
                    sketch[j * d + t] * sgn
                })
                .sum();
            base + coef * corr
        })
        .collect()
}

/// i.i.d. N(0,1) sketch matrix `S` for QJL (deterministic per (d, seed)).
pub fn gaussian_sketch(d: usize, seed: u64) -> Vec<f32> {
    let mut state = seed | 1;
    let mut next_u = move || {
        state ^= state << 13;
        state ^= state >> 7;
        state ^= state << 17;
        (state.wrapping_mul(0x2545F4914F6CDD1D) >> 11) as f64 / (1u64 << 53) as f64
    };
    (0..d * d)
        .map(|_| {
            let (u1, u2) = (next_u().max(1e-300), next_u());
            ((-2.0 * u1.ln()).sqrt() * (2.0 * std::f64::consts::PI * u2).cos()) as f32
        })
        .collect()
}

// Generic-width helpers (stage 2 uses (b−1)-bit tables; stage 1 uses 4-bit — same nearest rule).
fn quantize_mse_with(x: &[f32], rot: &[f32], cents: &[f32]) -> (Vec<u32>, f32) {
    let d = x.len();
    let norm = x.iter().map(|v| (*v as f64).powi(2)).sum::<f64>().sqrt() as f32;
    let scale = norm / (d as f32).sqrt();
    let inv = if scale > 0.0 { 1.0 / scale } else { 0.0 };
    let mut idx = vec![0u32; d.div_ceil(8)];
    for j in 0..d {
        let y: f32 = (0..d).map(|t| rot[j * d + t] * x[t]).sum();
        let z = y * inv;
        let mut best = 0u32;
        let mut bd = f32::INFINITY;
        for (ci, c) in cents.iter().enumerate() {
            let dd = (z - c).abs();
            if dd < bd {
                bd = dd;
                best = ci as u32;
            }
        }
        idx[j / 8] |= (best & 0xF) << ((j % 8) * 4);
    }
    (idx, scale)
}

fn dequant_rotated_any(packed: &[u32], scale: f32, cents: &[f32], d: usize) -> Vec<f32> {
    (0..d)
        .map(|j| cents[((packed[j / 8] >> ((j % 8) * 4)) & 0xF) as usize] * scale)
        .collect()
}

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

    fn rngv(seed: u64, n: usize) -> Vec<f32> {
        let mut s = seed | 1;
        (0..n)
            .map(|_| {
                s ^= s << 13;
                s ^= s >> 7;
                s ^= s << 17;
                ((s.wrapping_mul(0x2545F4914F6CDD1D) >> 40) as u32 as f32 / (1u32 << 24) as f32)
                    - 0.5
            })
            .collect()
    }

    #[test]
    fn rotation_is_orthogonal_and_preserves_dots() {
        for d in [64usize, 128] {
            let r = orthogonal_rotation(d, 0xDEAD);
            // Rows orthonormal.
            for a in 0..8 {
                for b in 0..8 {
                    let dot: f32 = (0..d).map(|j| r[a * d + j] * r[b * d + j]).sum();
                    let expect = if a == b { 1.0 } else { 0.0 };
                    assert!((dot - expect).abs() < 1e-4, "d={d} rows {a},{b}: {dot}");
                }
            }
            // ⟨Πx, Πy⟩ == ⟨x, y⟩ — the property the attention path relies on.
            let x = rngv(1, d);
            let y = rngv(2, d);
            let rx: Vec<f32> = (0..d)
                .map(|j| (0..d).map(|t| r[j * d + t] * x[t]).sum())
                .collect();
            let ry: Vec<f32> = (0..d)
                .map(|j| (0..d).map(|t| r[j * d + t] * y[t]).sum())
                .collect();
            let d0: f32 = x.iter().zip(&y).map(|(a, b)| a * b).sum();
            let d1: f32 = rx.iter().zip(&ry).map(|(a, b)| a * b).sum();
            assert!((d0 - d1).abs() < 1e-3 * d0.abs().max(1.0));
        }
    }

    #[test]
    fn centroids_are_sorted_symmetric_and_match_known_optimum() {
        let c4 = lloyd_max_centroids(4);
        assert_eq!(c4.len(), 16);
        for w in c4.windows(2) {
            assert!(w[0] < w[1], "sorted");
        }
        for i in 0..8 {
            assert!(
                (c4[i] + c4[15 - i]).abs() < 1e-3,
                "symmetric: {} vs {}",
                c4[i],
                c4[15 - i]
            );
        }
        // 1-bit Lloyd–Max optimum for N(0,1) is ±√(2/π) ≈ ±0.7979 — the classic closed form.
        let c1 = lloyd_max_centroids(1);
        assert!((c1[1] - 0.7978846).abs() < 1e-3, "got {}", c1[1]);
    }

    #[test]
    fn stage1_relative_mse_within_paper_bound() {
        // Paper: D_mse ≤ 2.7·4^-b (relative, per unit norm). Check b=4 over random vectors.
        let d = 128;
        let rot = orthogonal_rotation(d, 7);
        let cents = lloyd_max_centroids(4);
        // The theorem bounds EXPECTED relative distortion; per-vector draws fluctuate. Gate the
        // MEAN against the paper bound and the worst single vector against a 2× sanity cap.
        let bound = 2.7 * 4.0f64.powi(-4);
        let (mut sum, mut worst) = (0.0f64, 0.0f64);
        let n = 64;
        for s in 0..n {
            let x = rngv(1000 + s, d);
            let (packed, scale) = quantize_mse(&x, &rot, &cents);
            let deq_rot = dequant_rotated(&packed, scale, &cents, d);
            // Compare in the rotated domain (orthogonality ⇒ same MSE as original domain).
            let rx: Vec<f32> = (0..d)
                .map(|j| (0..d).map(|t| rot[j * d + t] * x[t]).sum())
                .collect();
            let err: f64 = rx
                .iter()
                .zip(&deq_rot)
                .map(|(a, b)| ((a - b) as f64).powi(2))
                .sum();
            let nrm: f64 = rx.iter().map(|v| (*v as f64).powi(2)).sum();
            sum += err / nrm;
            worst = worst.max(err / nrm);
        }
        let mean = sum / n as f64;
        assert!(
            mean <= bound,
            "mean relative MSE {mean:.5} exceeds paper bound {bound:.5} (worst {worst:.5})"
        );
        assert!(
            worst <= 2.0 * bound,
            "worst-vector relative MSE {worst:.5} exceeds 2x sanity cap"
        );
    }

    #[test]
    fn stage2_inner_products_are_unbiased_and_tighter() {
        // E[⟨y, x̃⟩] = ⟨y, x⟩: the mean signed error over many (x, y) pairs must be ~0, and far
        // below the mean |error| (bias ≪ noise). Uses 3-bit MSE + 1-bit QJL = 4 bits total.
        let d = 128;
        let rot = orthogonal_rotation(d, 11);
        let sketch = gaussian_sketch(d, 13);
        let cents3 = lloyd_max_centroids(3);
        let (mut signed, mut absolute, mut n) = (0.0f64, 0.0f64, 0);
        for s in 0..64 {
            let x = rngv(5000 + s, d);
            let y = rngv(9000 + s, d);
            let (p, sg, gamma, scale) = quantize_unbiased(&x, &rot, &sketch, &cents3);
            let xt = dequant_unbiased(&p, &sg, gamma, scale, &rot, &sketch, &cents3, d);
            let exact: f64 = x
                .iter()
                .zip(&y)
                .map(|(a, b)| (*a as f64) * (*b as f64))
                .sum();
            let est: f64 = xt
                .iter()
                .zip(&y)
                .map(|(a, b)| (*a as f64) * (*b as f64))
                .sum();
            signed += est - exact;
            absolute += (est - exact).abs();
            n += 1;
        }
        let (bias, noise) = (signed / n as f64, absolute / n as f64);
        assert!(
            bias.abs() < 0.35 * noise,
            "estimator looks biased: mean signed {bias:.5} vs mean |err| {noise:.5}"
        );
    }
}