inferencelayer 0.2.8

//! The quantization RULER: KL divergence between a reference model's next-token distribution and
//! a candidate arm's, accumulated over a calibration corpus.
//!
//! Why this module exists: every precision decision in this engine is a speed/fidelity trade, and
//! until now the fidelity side was judged by cosine on a handful of prompts (encoder rows) or by
//! top-1 agreement on one small model (the `WDtype` ladder in `weights.rs`). Neither ranks
//! *quantization schemes* well — cosine is insensitive to the tail that sampling actually draws
//! from, and top-1 agreement saturates long before a scheme stops hurting. The practice that does
//! rank them (llama.cpp's quant comparisons, Unsloth's Dynamic-2.0 methodology) is **KL divergence
//! against a reference model on a fixed calibration set**, which is what this computes.
//!
//! Direction matters and is fixed here: `KLD(P_ref ‖ Q_arm) = Σ p·(log p − log q)`. This is the
//! "excess surprise of serving the arm to a consumer who expects the reference" — it weights by
//! the REFERENCE's mass, so an arm that shifts probability onto tokens the reference considers
//! dead is penalized exactly as much as that mass deserves, and no more. The reverse direction is
//! mode-seeking and would flatter aggressive quantization; `kld_reverse` is provided for
//! diagnosis, never for ranking.
//!
//! Everything here is pure f64 math over logit rows, GPU-free and wasm-safe, so it is unit-tested
//! without a device. The teacher-forced loop that produces the rows lives in `src/bin/kld_eval.rs`.

/// The default top-k overlap width. 10 is the llama.cpp convention and is wide enough that a
/// scheme has to disturb real sampling mass to move it, unlike top-1.
pub const TOPK: usize = 10;

/// Numerically stable natural log-softmax in f64 (max-subtracted). Input is the f32 logit row the
/// GPU head produces; the f64 widening happens here, ONCE, so every downstream statistic is
/// computed in the wider type — at 250k-wide vocabularies an f32 log-sum-exp loses ~3 decimal
/// digits, which is the same order as the KLD values being compared.
pub fn log_softmax_f64(logits: &[f32]) -> Vec<f64> {
    let mut max = f64::NEG_INFINITY;
    for &l in logits {
        let l = f64::from(l);
        if l > max {
            max = l;
        }
    }
    let mut sum = 0.0f64;
    let mut out = Vec::with_capacity(logits.len());
    for &l in logits {
        let d = f64::from(l) - max;
        sum += d.exp();
        out.push(d);
    }
    let lse = sum.ln();
    for v in &mut out {
        *v -= lse;
    }
    out
}

/// Indices of the `k` largest entries, descending by value. Ties break by index so the result is
/// deterministic (a pinned metric must not depend on sort stability).
pub fn topk_indices(logits: &[f32], k: usize) -> Vec<u32> {
    let mut idx: Vec<u32> = (0..logits.len() as u32).collect();
    let k = k.min(idx.len());
    idx.sort_unstable_by(|&a, &b| {
        logits[b as usize]
            .partial_cmp(&logits[a as usize])
            .unwrap_or(std::cmp::Ordering::Equal)
            .then(a.cmp(&b))
    });
    idx.truncate(k);
    idx
}

/// One teacher-forced position's comparison between the reference and the arm.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct StepCompare {
    /// `KLD(P_ref ‖ Q_arm)` in nats. Zero iff the two logit rows induce the same distribution.
    pub kld: f64,
    /// Did the two arms pick the same argmax token? (What greedy decoding would see.)
    pub top1_agree: bool,
    /// `|top-k(ref) ∩ top-k(arm)| / k` — the sampling-relevant overlap.
    pub topk_overlap: f64,
    /// Reference log-probability of the GROUND-TRUTH next token (for the reference's perplexity).
    pub ref_logprob_target: f64,
    /// Arm log-probability of the ground-truth next token (for the arm's perplexity).
    pub arm_logprob_target: f64,
}

/// Compare one position. `target` is the ground-truth next token (teacher forcing), used only for
/// the perplexity terms — the KLD itself is over the full distributions and does not depend on it.
///
/// # Panics
/// If the two logit rows differ in length (comparing different vocabularies is a caller bug that
/// must never be silently averaged into a pin).
pub fn compare_step(ref_logits: &[f32], arm_logits: &[f32], target: u32) -> StepCompare {
    assert_eq!(
        ref_logits.len(),
        arm_logits.len(),
        "vocab mismatch between reference and arm"
    );
    let t = target as usize;
    assert!(t < ref_logits.len(), "target {target} outside vocab");
    let lp = log_softmax_f64(ref_logits);
    let lq = log_softmax_f64(arm_logits);
    // Σ p·(log p − log q), accumulated in f64. Terms where p underflows to 0 contribute nothing
    // (0·log 0 ≡ 0 by convention and by the limit), and skipping them also protects the sum from
    // a −inf·0 NaN when the arm assigns a token −inf.
    let mut kld = 0.0f64;
    for (a, b) in lp.iter().zip(&lq) {
        let p = a.exp();
        if p > 0.0 {
            kld += p * (a - b);
        }
    }
    let ref_top = topk_indices(ref_logits, TOPK);
    let arm_top = topk_indices(arm_logits, TOPK);
    let hits = ref_top.iter().filter(|i| arm_top.contains(i)).count();
    StepCompare {
        kld,
        top1_agree: ref_top[0] == arm_top[0],
        topk_overlap: hits as f64 / ref_top.len() as f64,
        ref_logprob_target: lp[t],
        arm_logprob_target: lq[t],
    }
}

/// `KLD(Q_arm ‖ P_ref)` — the reverse (mode-seeking) direction, for diagnosis only. A large
/// forward/reverse asymmetry says the arm moved mass ONTO tokens the reference considers dead,
/// which is the failure mode that produces confident nonsense.
pub fn kld_reverse(ref_logits: &[f32], arm_logits: &[f32]) -> f64 {
    compare_step(arm_logits, ref_logits, 0).kld
}

/// Streaming accumulator over a calibration run. Keeps the per-step KLDs because the DISTRIBUTION
/// of divergence is the interesting part — a scheme with a low mean and a fat p99 breaks rarely
/// and badly, which a mean alone hides.
#[derive(Clone, Debug, Default)]
pub struct KldAccum {
    klds: Vec<f64>,
    top1_hits: usize,
    topk_overlap_sum: f64,
    ref_nll: f64,
    arm_nll: f64,
}

impl KldAccum {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn push(&mut self, c: StepCompare) {
        self.klds.push(c.kld);
        self.top1_hits += usize::from(c.top1_agree);
        self.topk_overlap_sum += c.topk_overlap;
        self.ref_nll -= c.ref_logprob_target;
        self.arm_nll -= c.arm_logprob_target;
    }

    pub fn len(&self) -> usize {
        self.klds.len()
    }

    pub fn is_empty(&self) -> bool {
        self.klds.is_empty()
    }

    /// Reduce to the pinned summary. Returns `None` for an empty run rather than emitting NaNs
    /// into a pin file.
    pub fn summary(&self) -> Option<KldSummary> {
        if self.klds.is_empty() {
            return None;
        }
        let n = self.klds.len();
        let mut sorted = self.klds.clone();
        sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
        Some(KldSummary {
            n,
            kld_mean: self.klds.iter().sum::<f64>() / n as f64,
            kld_median: quantile(&sorted, 0.5),
            kld_p99: quantile(&sorted, 0.99),
            kld_max: sorted[n - 1],
            top1_agreement: self.top1_hits as f64 / n as f64,
            topk_overlap: self.topk_overlap_sum / n as f64,
            ppl_ref: (self.ref_nll / n as f64).exp(),
            ppl_arm: (self.arm_nll / n as f64).exp(),
        })
    }
}

/// Nearest-rank quantile of an ASCENDING slice (no interpolation — the pins are compared across
/// runs, and nearest-rank is the definition that stays stable as `n` changes).
fn quantile(sorted: &[f64], q: f64) -> f64 {
    let n = sorted.len();
    let rank = (q * n as f64).ceil() as usize;
    sorted[rank.clamp(1, n) - 1]
}

/// The row written to `bench/kld_pins.json`.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct KldSummary {
    /// Scored positions (calibration tokens minus one per document — the last token has no
    /// ground-truth successor).
    pub n: usize,
    pub kld_mean: f64,
    pub kld_median: f64,
    pub kld_p99: f64,
    pub kld_max: f64,
    pub top1_agreement: f64,
    pub topk_overlap: f64,
    /// Reference perplexity on the calibration set — the sanity anchor. If this is not in the
    /// plausible range for the model/corpus, the harness is mis-wired and no KLD below it means
    /// anything.
    pub ppl_ref: f64,
    pub ppl_arm: f64,
}

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

    /// Logits for a 2-symbol distribution with the given probabilities.
    fn logits2(p0: f64, p1: f64) -> Vec<f32> {
        vec![p0.ln() as f32, p1.ln() as f32]
    }

    #[test]
    fn self_vs_self_is_exactly_zero() {
        // The primary gate: an arm compared against itself must produce EXACTLY 0.0, not
        // "small". Bitwise-identical rows give (a − b) == 0.0 for every term, so any nonzero
        // result means the two sides were not computed the same way.
        let l: Vec<f32> = (0..1024).map(|i| ((i * 37 % 101) as f32) * 0.1 - 5.0).collect();
        let c = compare_step(&l, &l, 7);
        assert_eq!(c.kld, 0.0, "self-KLD must be exactly zero");
        assert!(c.top1_agree);
        assert_eq!(c.topk_overlap, 1.0);
        assert_eq!(c.ref_logprob_target, c.arm_logprob_target);
    }

    #[test]
    fn matches_hand_computed_two_point_kld() {
        // p = (0.5, 0.5), q = (0.75, 0.25)
        // KLD = 0.5·ln(0.5/0.75) + 0.5·ln(0.5/0.25) = 0.5·(−0.4054651) + 0.5·(0.6931472)
        //     = 0.1438410
        let c = compare_step(&logits2(0.5, 0.5), &logits2(0.75, 0.25), 0);
        assert!(
            (c.kld - 0.143_841_0).abs() < 1e-6,
            "hand-computed KLD mismatch: {}",
            c.kld
        );
    }

    #[test]
    fn kld_is_asymmetric_and_reverse_agrees() {
        // KLD(q‖p) = 0.75·ln(0.75/0.5) + 0.25·ln(0.25/0.5) = 0.3040 − 0.1733 = 0.1308
        let (p, q) = (logits2(0.5, 0.5), logits2(0.75, 0.25));
        let fwd = compare_step(&p, &q, 0).kld;
        let rev = kld_reverse(&p, &q);
        assert!((rev - 0.130_812_0).abs() < 1e-6, "reverse KLD: {rev}");
        assert!(
            (fwd - rev).abs() > 1e-3,
            "the two directions must not be conflated"
        );
    }

    #[test]
    fn kld_is_shift_invariant() {
        // Softmax is invariant to a constant logit offset; a naive exp() without max-subtraction
        // overflows and this catches it. The shift must exceed f64's exp range (~709) to have
        // teeth — an earlier +80 was sized for f32 and passed against a max-subtraction-free
        // implementation, i.e. tested nothing (found by mutation).
        let a: Vec<f32> = vec![1.0, 2.0, 3.0, 0.5];
        let b: Vec<f32> = vec![1.5, 2.0, 2.5, 0.0];
        let base = compare_step(&a, &b, 2).kld;
        assert!(base > 1e-3, "base case must be a real divergence: {base}");
        let sa: Vec<f32> = a.iter().map(|v| v + 800.0).collect();
        let sb: Vec<f32> = b.iter().map(|v| v + 800.0).collect();
        let shifted = compare_step(&sa, &sb, 2).kld;
        assert!(base.is_finite() && shifted.is_finite());
        assert!(
            (base - shifted).abs() < 1e-12,
            "shift changed KLD: {base} vs {shifted}"
        );
    }

    #[test]
    fn perturbation_raises_kld_and_drops_agreement() {
        // The plan's synthetic gate: a perturbed arm must register as WORSE on every axis.
        let refl: Vec<f32> = (0..512).map(|i| ((i % 17) as f32) * 0.3).collect();
        let mut arm = refl.clone();
        arm[3] += 9.0; // force a different argmax
        for (i, v) in arm.iter_mut().enumerate() {
            *v += ((i % 5) as f32) * 0.05; // and a broad low-level distortion
        }
        let c = compare_step(&refl, &arm, 11);
        assert!(c.kld > 0.0, "perturbation must show positive KLD");
        assert!(!c.top1_agree, "argmax should have moved");
        assert!(c.topk_overlap < 1.0, "top-k set should have moved");
    }

    #[test]
    fn kld_is_non_negative_on_random_pairs() {
        // Gibbs' inequality — the invariant that catches a sign or direction error. Deterministic
        // pseudo-random pairs (no rng dependency, reproducible in CI).
        let mut seed = 0x2545_F491_4F6C_DD1Du64;
        let mut next = || {
            seed ^= seed << 13;
            seed ^= seed >> 7;
            seed ^= seed << 17;
            (seed >> 40) as f32 / 1024.0 - 12.0
        };
        for _ in 0..64 {
            let a: Vec<f32> = (0..256).map(|_| next()).collect();
            let b: Vec<f32> = (0..256).map(|_| next()).collect();
            let k = compare_step(&a, &b, 0).kld;
            assert!(k >= -1e-9, "KLD went negative: {k}");
        }
    }

    #[test]
    fn topk_overlap_counts_set_intersection_not_order() {
        // Same top-3 SET, different internal order → overlap 1.0 at k=3 (order is top1_agree's
        // job, not the overlap's).
        let a: Vec<f32> = vec![5.0, 4.0, 3.0, 0.0, 0.0];
        let b: Vec<f32> = vec![3.0, 4.0, 5.0, 0.0, 0.0];
        assert_eq!(topk_indices(&a, 3), vec![0, 1, 2]);
        assert_eq!(topk_indices(&b, 3), vec![2, 1, 0]);
        let hits = topk_indices(&a, 3)
            .iter()
            .filter(|i| topk_indices(&b, 3).contains(i))
            .count();
        assert_eq!(hits, 3);
    }

    #[test]
    fn topk_ties_break_by_index_deterministically() {
        let flat: Vec<f32> = vec![1.0; 8];
        assert_eq!(topk_indices(&flat, 3), vec![0, 1, 2]);
    }

    #[test]
    fn uniform_distribution_has_perplexity_equal_to_vocab() {
        // A uniform 64-way distribution must score ppl = 64 exactly (mean NLL = ln 64).
        let flat: Vec<f32> = vec![0.0; 64];
        let mut acc = KldAccum::new();
        for t in 0..10u32 {
            acc.push(compare_step(&flat, &flat, t));
        }
        let s = acc.summary().expect("non-empty");
        assert!((s.ppl_ref - 64.0).abs() < 1e-9, "ppl_ref {}", s.ppl_ref);
        assert!((s.ppl_arm - 64.0).abs() < 1e-9, "ppl_arm {}", s.ppl_arm);
        assert_eq!(s.kld_mean, 0.0);
        assert_eq!(s.top1_agreement, 1.0);
        assert_eq!(s.n, 10);
    }

    #[test]
    fn summary_order_statistics_are_nearest_rank() {
        let mut acc = KldAccum::new();
        // KLDs 0.0 … 0.99 via 100 synthetic steps built from real comparisons would be slow;
        // drive the accumulator directly (the fields it reduces are what is under test).
        for i in 0..100 {
            acc.push(StepCompare {
                kld: i as f64 / 100.0,
                top1_agree: i < 90,
                topk_overlap: 1.0,
                ref_logprob_target: -1.0,
                arm_logprob_target: -1.0,
            });
        }
        let s = acc.summary().expect("non-empty");
        assert_eq!(s.n, 100);
        assert!((s.kld_median - 0.49).abs() < 1e-12, "median {}", s.kld_median);
        assert!((s.kld_p99 - 0.98).abs() < 1e-12, "p99 {}", s.kld_p99);
        assert!((s.kld_max - 0.99).abs() < 1e-12);
        assert!((s.top1_agreement - 0.90).abs() < 1e-12);
        assert!((s.ppl_ref - std::f64::consts::E).abs() < 1e-12);
    }

    #[test]
    fn quantile_uses_nearest_rank_not_floor() {
        // The order-statistics test above uses n=100, where 0.5·n and 0.99·n are exact integers
        // and ceil == floor — so it could not tell nearest-rank from a floor implementation
        // (found by mutation). Seven values make every rank fractional.
        let v: Vec<f64> = (0..7).map(|i| i as f64).collect();
        assert_eq!(quantile(&v, 0.5), 3.0, "rank ceil(3.5) = 4th smallest");
        assert_eq!(quantile(&v, 0.99), 6.0, "rank ceil(6.93) = 7th smallest");
        assert_eq!(quantile(&v, 1.0), 6.0);
        assert_eq!(quantile(&v, 0.0), 0.0, "rank clamps up to 1");
    }

    #[test]
    fn empty_accumulator_yields_no_summary() {
        assert!(KldAccum::new().summary().is_none());
    }
}