ferrox-models 0.16.0

//! The two CPU schedulers must produce the same tokens.
//!
//! Issue #27 replaces rayon's per-operation fork-join with a persistent
//! worker pool (`ferrox_core::cpu_pool`, selected by `FERROX_CPU_POOL`).
//! The reason that change exists is speed, and **an agent may not
//! benchmark**: measurement needs a quiet host and a loaded one reads
//! 25-45% low. So this suite asserts the other half, the half that can
//! be settled without a stopwatch -- that the new path computes exactly
//! what the old one did, on a real checkpoint, greedily, token for
//! token.
//!
//! It is not a smoke test of the pool (`ferrox_core::cpu_pool`'s unit
//! tests cover deadlock, dropped work, panics and lifetime). It is the
//! end-to-end statement: every rewritten matvec, every K-quant GEMV
//! tail, prefill and decode, agree bit-for-bit across the switch.
//!
//! # Why subprocesses
//!
//! `FERROX_CPU_POOL` is read once and cached for the process, which is
//! deliberate -- a scheduler that could change mid-run would make a
//! before/after measurement meaningless. So one process cannot exercise
//! both arms, and the comparison runs each in a child of this test
//! binary. Both children are spawned the same way, so neither arm gets
//! the advantage of being the one that ran in-process.

use std::path::{Path, PathBuf};
use std::process::Command;

use ferrox_core::cache::KvCache;
use ferrox_gguf::ShardedGguf;
use ferrox_models::config::ModelConfig;
use ferrox_models::decoder::Decoder;
use ferrox_models::tokenizer::GgufBpeTokenizer;

const PROMPT: &str = "The capital of France is";
const MAX_NEW_TOKENS: usize = 12;

/// Set on a child to make it actually run the generation.
const CHILD: &str = "FERROX_TEST_CPU_POOL_CHILD";
/// The checkpoint a child should generate from.
const CHILD_GGUF: &str = "FERROX_TEST_CPU_POOL_GGUF";
/// The line a child prints its token ids on.
const MARKER: &str = "FERROX_TOKENS";

/// `FERROX_TEST_MODELS_DIR`, else the workspace's `models/`. The
/// override exists because a git worktree does not carry the (ignored)
/// checkpoint directory, and a suite that silently skips there is a
/// suite that never runs while the work is being done.
fn models_dir() -> PathBuf {
    if let Some(dir) = std::env::var_os("FERROX_TEST_MODELS_DIR") {
        return PathBuf::from(dir);
    }
    Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models")
}

/// The checkpoints this runs over, when they are present.
///
/// Two quantizations rather than one on purpose: Q8_0 and Q4_K_M take
/// *different* rewritten kernels (`gemv_q8_0x4_group` and the
/// interleaved `Q4_KX8` group plus its scalar tail), and a chunking bug
/// in one would not show up in the other.
fn candidates() -> Vec<PathBuf> {
    let dir = models_dir();
    [
        "hf_test/SmolLM2-135M-Instruct-Q8_0.gguf",
        "hf_test/SmolLM2-135M-Instruct-Q4_K_M.gguf",
    ]
    .iter()
    .map(|name| dir.join(name))
    .filter(|path| path.exists())
    .collect()
}

fn argmax(logits: &[f32]) -> usize {
    logits
        .iter()
        .enumerate()
        .max_by(|(_, a), (_, b)| a.partial_cmp(b).expect("logits are finite"))
        .map(|(i, _)| i)
        .expect("non-empty logits")
}

/// Greedy continuation of [`PROMPT`]: prefill through `forward_batch`,
/// then `forward_token` per step, exactly the shape decode has.
fn greedy_tokens(path: &Path) -> Vec<usize> {
    let file = ShardedGguf::open(path).expect("open GGUF");
    let config = ModelConfig::from_gguf(&file).expect("model config");
    let tok = GgufBpeTokenizer::from_gguf(&file).expect("tokenizer");
    let decoder = Decoder::from_gguf(path, config.clone()).expect("load decoder");

    let prompt: Vec<usize> = tok.encode(PROMPT).into_iter().map(|t| t as usize).collect();
    let mut caches: Vec<KvCache> = (0..config.n_layers)
        .map(|_| KvCache::new(config.n_kv_heads, config.head_dim))
        .collect();

    let logits = decoder.forward_batch(&prompt, 0, &mut caches);
    let mut last = logits.last().expect("non-empty prompt").clone();

    let mut generated = Vec::with_capacity(MAX_NEW_TOKENS);
    for step in 0..MAX_NEW_TOKENS {
        let next = argmax(&last);
        generated.push(next);
        last = decoder.forward_token(next, prompt.len() + step, &mut caches);
    }
    generated
}

/// Run this test binary again with `backend` selected, and read back the
/// tokens it generated.
fn tokens_from_child(path: &Path, backend: &str) -> Vec<usize> {
    let exe = std::env::current_exe().expect("test binary path");
    let out = Command::new(exe)
        .args(["--exact", "generates_tokens_for_the_parent", "--nocapture"])
        .env(CHILD, "1")
        .env(CHILD_GGUF, path)
        .env("FERROX_CPU_POOL", backend)
        // Both arms run the integer `vec_dot` kernels, which is what the
        // product ships and where all the rewritten GEMVs live. Without
        // this the library default (reference-exact f32 dequant-dot)
        // would leave the interleaved paths untested.
        .env("FERROX_CPU_INT_DOT", "1")
        .env("FERROX_TEST_MODELS_DIR", models_dir())
        .output()
        .expect("spawn child test process");
    let stdout = String::from_utf8_lossy(&out.stdout);
    assert!(
        out.status.success(),
        "child ({backend}) failed: {}\n{}",
        out.status,
        String::from_utf8_lossy(&out.stderr)
    );
    let line = stdout
        .lines()
        .find_map(|l| l.strip_prefix(MARKER))
        .unwrap_or_else(|| panic!("child ({backend}) printed no {MARKER} line:\n{stdout}"));
    line.split_whitespace()
        .map(|t| t.parse::<usize>().expect("token id"))
        .collect()
}

/// The child half. Inert unless [`CHILD`] is set, so an ordinary
/// `cargo test` run does not load a model twice for nothing.
#[test]
fn generates_tokens_for_the_parent() {
    if std::env::var_os(CHILD).is_none() {
        return;
    }
    let path = PathBuf::from(std::env::var_os(CHILD_GGUF).expect("child needs a GGUF path"));
    let tokens = greedy_tokens(&path);
    let ids: Vec<String> = tokens.iter().map(|t| t.to_string()).collect();
    println!("{MARKER} {}", ids.join(" "));
}

/// **The assertion this PR rests on.**
///
/// Sabotage: make `ferrox_core::par`'s spin arm drop its last task (e.g.
/// `hi = ((t + 1) * per).min(n) - 1`) and this goes red on the first
/// checkpoint, because a matvec that skips rows changes the argmax.
#[test]
fn the_two_cpu_schedulers_produce_token_identical_output() {
    let models = candidates();
    if models.is_empty() {
        eprintln!(
            "skip: no checkpoint under {} -- this suite needs a real GGUF",
            models_dir().display()
        );
        return;
    }
    for path in models {
        let rayon = tokens_from_child(&path, "rayon");
        let spin = tokens_from_child(&path, "spin");
        assert_eq!(
            rayon.len(),
            MAX_NEW_TOKENS,
            "{}: the rayon arm generated nothing to compare",
            path.display()
        );
        assert_eq!(
            rayon,
            spin,
            "{}: the persistent pool and rayon disagree about the greedy continuation",
            path.display()
        );
    }
}