inferencelayer 0.2.13

//! `kld-eval` — the quantization ruler's driver: teacher-forced KL divergence between a REFERENCE
//! arm and a candidate ARM over a calibration corpus.
//!
//! ```text
//! kld-eval --model <dir> --calib bench/calib_v1.txt [--ref gpu|cpu] [--arm gpu]
//!          [--arm-model <dir>] [--max-docs N] [--max-tokens N] [--label NAME] [--out FILE]
//! ```
//!
//! ## Why teacher forcing (and not "generate and compare")
//!
//! Two arms that differ at ONE position diverge into different continuations, after which every
//! later comparison is measuring two different texts rather than two different models. Feeding
//! both arms the same ground-truth token at every step keeps them on the identical prefix, so the
//! per-position distributions are comparable and the statistic is the model difference alone.
//! This is what `llama-perplexity --kl-divergence` does and why its numbers are stable.
//!
//! ## Why the arms run interleaved, position by position
//!
//! Storing one arm's full logit rows for later comparison is not affordable — 20k positions ×
//! 150k vocab × f32 is ~12 GB per arm — and an f16 dump loses ~3 decimal digits, the same order
//! as the KLD values being compared. So both arms are resident and stepped in lockstep, and only
//! the f64 summary statistics survive each position. The cost is holding two models at once,
//! which is why `--max-tokens` exists for the cases that do not fit.
//!
//! ## Arms
//!
//! * `gpu` — [`Lfm2Gpu::forward`], the engine's own M=1 decode path (Q4 today).
//! * `cpu` — the f32 CPU oracle the kernels are verified against ([`Qwen35Ref`] for the Qwen3.5
//!   hybrid, [`RefModel`] for LFM2), i.e. a TRUE full-precision reference rather than a
//!   higher-precision quantization. Slow by construction: it is the spec, not a fast path.
//!
//! `--ref gpu --arm gpu` with one model is the SELF-GATE: two independently constructed engines
//! over the same weights must produce exactly-zero divergence.

use anyhow::{Context, Result};
use inferencelayer::kld::{KldAccum, KldSummary, compare_step};

/// Which implementation produces an arm's logits.
enum Arm {
    Gpu(Box<inferencelayer::Lfm2Gpu>),
    CpuQwen35(Box<inferencelayer::reference_qwen35::Qwen35Ref>),
    CpuLfm2(Box<inferencelayer::reference::RefModel>),
}

impl Arm {
    fn load(kind: &str, ctx: &inferencelayer::GpuCtx, dir: &std::path::Path) -> Result<Self> {
        match kind {
            // `gpu` honours OSFKB_DECODE_PRECISION; `gpu-q4`/`gpu-f16` pin the storage dtype
            // explicitly, which is what a two-arm comparison needs — both arms live in ONE
            // process, so an environment variable cannot distinguish them.
            "gpu" | "gpu-q4" | "gpu-f16" | "gpu-q4-head-f16" => {
                use inferencelayer::weights::{Precision, WDtype};
                let w = match kind {
                    "gpu-q4" => inferencelayer::Weights::load_with_precision(
                        ctx,
                        dir,
                        Precision::uniform(WDtype::Q4),
                    )?,
                    "gpu-f16" => inferencelayer::Weights::load_with_precision(
                        ctx,
                        dir,
                        Precision::uniform(WDtype::F16),
                    )?,
                    // Unsloth's canonical recovery lever: Q4 body, f16 head.
                    "gpu-q4-head-f16" => inferencelayer::Weights::load_with_precision(
                        ctx,
                        dir,
                        Precision::parse("q4,head=f16").expect("static spec parses"),
                    )?,
                    _ => inferencelayer::Weights::load(ctx, dir)?,
                };
                Ok(Self::Gpu(Box::new(inferencelayer::Lfm2Gpu::new(ctx, w))))
            }
            // Arbitrary precision policy: `gpu-prec:<spec>` with the full OSFKB_DECODE_PRECISION
            // grammar, e.g. `gpu-prec:q4,blk:0-1=f16,head=f16` — the arm the per-layer
            // sensitivity sweeps drive.
            k if k.starts_with("gpu-prec:") => {
                use inferencelayer::weights::Precision;
                let spec = &k["gpu-prec:".len()..];
                let w = inferencelayer::Weights::load_with_precision(
                    ctx,
                    dir,
                    Precision::parse(spec)?,
                )?;
                Ok(Self::Gpu(Box::new(inferencelayer::Lfm2Gpu::new(ctx, w))))
            }
            // A llama-arch GGUF (incl. an Unsloth Dynamic K-quant mix) repacked to f16/Q4. `dir`
            // is the .gguf FILE. `gpu-gguf-f16` is the near-lossless repack: comparing it to a
            // safetensors `gpu-f16` arm isolates the GGUF's OWN quantization damage.
            // NATIVE: the artifact's own block formats served as-is (padded, values untouched).
            "gpu-gguf-native" => {
                let w = inferencelayer::gguf::load_gguf_native(ctx, dir)?;
                Ok(Self::Gpu(Box::new(inferencelayer::Lfm2Gpu::new(ctx, w))))
            }
            "gpu-gguf-f16" | "gpu-gguf-q4" => {
                use inferencelayer::weights::WDtype;
                let dt = if kind.ends_with("f16") {
                    WDtype::F16
                } else {
                    WDtype::Q4
                };
                let w = inferencelayer::gguf::load_gguf(ctx, dir, dt)?;
                Ok(Self::Gpu(Box::new(inferencelayer::Lfm2Gpu::new(ctx, w))))
            }
            // The f32 oracle: try the Qwen3.5 hybrid reference first, fall back to LFM2's. Both
            // expose the same `forward(token, pos) -> Vec<f32>` shape.
            "cpu" => match inferencelayer::reference_qwen35::Qwen35Ref::load(dir) {
                Ok(m) => Ok(Self::CpuQwen35(Box::new(m))),
                Err(e_q) => match inferencelayer::reference::RefModel::load(dir) {
                    Ok(m) => Ok(Self::CpuLfm2(Box::new(m))),
                    Err(e_l) => anyhow::bail!(
                        "no CPU f32 reference for this checkpoint (qwen35: {e_q}; lfm2: {e_l})"
                    ),
                },
            },
            other => anyhow::bail!(
                "unknown arm `{other}` (expected `gpu`, `gpu-q4`, `gpu-f16` or `cpu`)"
            ),
        }
    }

    fn label(&self) -> &'static str {
        use inferencelayer::weights::WDtype;
        match self {
            Self::Gpu(g) => {
                if !g.w.cfg.layer_kinds.is_empty() {
                    return "gpu-gguf-native";
                }
                if !g.w.cfg.layer_wdtype.is_empty() {
                    return "gpu-blk-mixed"; // per-layer override in force; --label carries detail
                }
                let (b, h) = (g.w.cfg.wdtype, g.w.cfg.head_wdtype);
                // Name the head separately when it differs from the body — the pin must record
                // that this is a body/head mix, not a uniform arm.
                match (b, h) {
                    (WDtype::Q4, WDtype::F16) => "gpu-q4-head-f16",
                    (WDtype::F16, WDtype::Q4) => "gpu-f16-head-q4",
                    _ if b == h => match b {
                        WDtype::F16 => "gpu-f16",
                        WDtype::Q8 => "gpu-q8",
                        WDtype::Q4 => "gpu-q4",
                    },
                    _ => "gpu-mixed",
                }
            }
            Self::CpuQwen35(_) | Self::CpuLfm2(_) => "cpu-f32",
        }
    }

    fn reset(&mut self, ctx: &inferencelayer::GpuCtx) {
        match self {
            Self::Gpu(g) => g.reset(ctx),
            Self::CpuQwen35(m) => m.reset(),
            Self::CpuLfm2(m) => m.reset(),
        }
    }

    fn forward(
        &mut self,
        ctx: &inferencelayer::GpuCtx,
        token: u32,
        pos: usize,
    ) -> Result<Vec<f32>> {
        match self {
            Self::Gpu(g) => g.forward(ctx, token, pos),
            Self::CpuQwen35(m) => Ok(m.forward(token, pos)),
            Self::CpuLfm2(m) => Ok(m.forward(token, pos)),
        }
    }
}

struct Args {
    model: String,
    arm_model: Option<String>,
    calib: String,
    ref_kind: String,
    arm_kind: String,
    max_docs: usize,
    max_tokens: usize,
    label: Option<String>,
    out: Option<String>,
}

fn parse_args() -> Result<Args> {
    let mut a = Args {
        model: String::new(),
        arm_model: None,
        calib: "bench/calib_v1.txt".into(),
        ref_kind: "cpu".into(),
        arm_kind: "gpu".into(),
        max_docs: usize::MAX,
        max_tokens: 512,
        label: None,
        out: None,
    };
    let mut it = std::env::args().skip(1);
    while let Some(k) = it.next() {
        match k.as_str() {
            "--model" => a.model = it.next().context("--model value")?,
            "--arm-model" => a.arm_model = it.next(),
            "--calib" => a.calib = it.next().context("--calib value")?,
            "--ref" => a.ref_kind = it.next().context("--ref value")?,
            "--arm" => a.arm_kind = it.next().context("--arm value")?,
            "--max-docs" => a.max_docs = it.next().context("--max-docs value")?.parse()?,
            "--max-tokens" => a.max_tokens = it.next().context("--max-tokens value")?.parse()?,
            "--label" => a.label = it.next(),
            "--out" => a.out = it.next(),
            other => anyhow::bail!("unknown argument {other}"),
        }
    }
    anyhow::ensure!(!a.model.is_empty(), "--model <dir> required");
    Ok(a)
}

fn main() -> Result<()> {
    let args = parse_args()?;
    let dir = std::path::Path::new(&args.model);
    let arm_dir = args
        .arm_model
        .as_ref()
        .map_or(dir, |s| std::path::Path::new(s));

    let ctx = inferencelayer::GpuCtx::new()?;
    eprintln!("backend: {}", ctx.backend);

    // ── Calibration corpus ──────────────────────────────────────────────────────────────────
    let text = std::fs::read_to_string(&args.calib)
        .with_context(|| format!("calibration corpus {}", args.calib))?;
    let docs: Vec<&str> = text
        .lines()
        .map(str::trim)
        .filter(|l| !l.is_empty() && !l.starts_with('#'))
        .take(args.max_docs)
        .collect();
    anyhow::ensure!(!docs.is_empty(), "calibration corpus has no documents");
    let tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
        .map_err(|e| anyhow::anyhow!("tokenizer: {e}"))?;

    // ── Both arms resident (see the module docs: no affordable dump format) ──────────────────
    let t0 = std::time::Instant::now();
    let mut refa = Arm::load(&args.ref_kind, &ctx, dir)?;
    let mut arma = Arm::load(&args.arm_kind, &ctx, arm_dir)?;
    eprintln!(
        "arms loaded in {:.1}s: ref={} arm={}",
        t0.elapsed().as_secs_f64(),
        refa.label(),
        arma.label()
    );

    let mut acc = KldAccum::new();
    let (mut ref_secs, mut arm_secs) = (0.0f64, 0.0f64);
    let run0 = std::time::Instant::now();
    for (di, doc) in docs.iter().enumerate() {
        let ids = tok
            .encode(*doc, true)
            .map_err(|e| anyhow::anyhow!("encode: {e}"))?
            .get_ids()
            .to_vec();
        let ids: Vec<u32> = ids.into_iter().take(args.max_tokens).collect();
        if ids.len() < 2 {
            continue; // no scorable position
        }
        refa.reset(&ctx);
        arma.reset(&ctx);
        // Teacher forcing: both arms see ids[pos], and the logits are scored against ids[pos+1].
        // The final token is fed to neither — it has no ground-truth successor.
        for pos in 0..ids.len() - 1 {
            let t = std::time::Instant::now();
            let rl = refa.forward(&ctx, ids[pos], pos)?;
            ref_secs += t.elapsed().as_secs_f64();
            let t = std::time::Instant::now();
            let al = arma.forward(&ctx, ids[pos], pos)?;
            arm_secs += t.elapsed().as_secs_f64();
            acc.push(compare_step(&rl, &al, ids[pos + 1]));
        }
        eprintln!(
            "doc {}/{}: {} positions scored ({} total)",
            di + 1,
            docs.len(),
            ids.len() - 1,
            acc.len()
        );
    }
    let wall = run0.elapsed().as_secs_f64();
    let s = acc.summary().context("no positions scored")?;
    let n = s.n as f64;
    report(
        &args,
        &s,
        refa.label(),
        arma.label(),
        wall,
        ref_secs,
        arm_secs,
    );
    eprintln!(
        "\nperf: {:.1}s wall | ref {:.1} tok/s | arm {:.1} tok/s",
        wall,
        n / ref_secs.max(1e-9),
        n / arm_secs.max(1e-9)
    );
    Ok(())
}

fn report(
    args: &Args,
    s: &KldSummary,
    ref_label: &str,
    arm_label: &str,
    wall: f64,
    ref_secs: f64,
    arm_secs: f64,
) {
    println!("\n── KLD ruler ────────────────────────────────────────────");
    println!("model      {}", args.model);
    println!("calib      {} ({} positions)", args.calib, s.n);
    println!("ref / arm  {ref_label} / {arm_label}");
    println!("KLD mean   {:.6} nats", s.kld_mean);
    println!("    median {:.6}", s.kld_median);
    println!("    p99    {:.6}", s.kld_p99);
    println!("    max    {:.6}", s.kld_max);
    println!("top-1      {:.4} agreement", s.top1_agreement);
    println!("top-10     {:.4} overlap", s.topk_overlap);
    println!("ppl ref    {:.4}", s.ppl_ref);
    println!("ppl arm    {:.4}", s.ppl_arm);

    if let Some(out) = &args.out {
        let row = serde_json::json!({
            "label": args.label.clone().unwrap_or_else(|| format!("{ref_label}-vs-{arm_label}")),
            "model": args.model,
            "arm_model": args.arm_model,
            "calib": args.calib,
            "ref": ref_label,
            "arm": arm_label,
            "n": s.n,
            "kld_mean": s.kld_mean,
            "kld_median": s.kld_median,
            "kld_p99": s.kld_p99,
            "kld_max": s.kld_max,
            "top1_agreement": s.top1_agreement,
            "top10_overlap": s.topk_overlap,
            "ppl_ref": s.ppl_ref,
            "ppl_arm": s.ppl_arm,
            "wall_s": wall,
            "ref_tok_s": s.n as f64 / ref_secs.max(1e-9),
            "arm_tok_s": s.n as f64 / arm_secs.max(1e-9),
        });
        // Append-mode pin file: one JSON row per line, so parallel runs never clobber a shared
        // array and the history of every arm stays readable (the scoreboard-pins convention).
        match std::fs::OpenOptions::new()
            .create(true)
            .append(true)
            .open(out)
        {
            Ok(mut f) => {
                use std::io::Write;
                let _ = writeln!(f, "{row}");
                eprintln!("pinned → {out}");
            }
            Err(e) => eprintln!("WARN: could not write {out}: {e}"),
        }
    }
}