keyhog-scanner 0.5.44

//! Mixture-of-experts model weights for the ML secret scorer.
//!
//! Architecture: gate Linear(55,6) → Softmax plus 6 experts of
//! Linear(55,32) → ReLU → Linear(32,16) → ReLU → Linear(16,1)
//!
//! Generated by: `python3 ml/train_classifier.py --features 55 --write`
//! Serialized as little-endian `f32` values so parsing is deterministic across
//! host architectures. INPUT_DIM (55) is owned by `model_arch.rs`; the feature
//! layout (including the decode-structure feature) is defined in `ml_features.rs`,
//! which lets the model discount base64/hex-wrapped binary.

include!(concat!(env!("OUT_DIR"), "/model_version.rs"));

// `weights.bin` stays at `src/` (build.rs reads `src/weights.bin`); this file
// now lives in `src/ml_scorer/`, so the embed path goes up one level.
pub(crate) const WEIGHTS: &[u8] = include_bytes!("../weights.bin");

// Model architecture dimensions come from the single owner `super::model_arch`
// (the flat `weights.bin` layout below AND the forward pass's const-generic
// layer sizes both read the SAME consts, so a retrain that changes the
// architecture cannot silently disagree between the buffer strides here and the
// scorer). The parse-time buffer-size check (`parse_weights`) fails closed on any
// stride mismatch.
use super::model_arch::{
    EXPERTS_OFF, EXPERT_COUNT, EXPERT_FC1_B_COUNT, EXPERT_FC1_OUT, EXPERT_FC1_W_COUNT,
    EXPERT_FC2_B_COUNT, EXPERT_FC2_OUT, EXPERT_FC2_W_COUNT, EXPERT_FC3_B_COUNT, EXPERT_FC3_W_COUNT,
    EXPERT_PARAM_COUNT, GATE_B_COUNT, GATE_B_OFF, GATE_W_COUNT, GATE_W_OFF, INPUT_DIM,
    TOTAL_F32_COUNT,
};

/// Parse the embedded little-endian `f32` weight buffer, failing closed on a
/// size mismatch OR any non-finite (NaN/inf) value.
///
/// Model integrity: the size check guarantees the weights file matches the
/// expected architecture. We intentionally do NOT use a SHA-256 checksum because
/// it breaks the retrain→deploy workflow, every retrain produces new weights
/// with a different hash. But size alone accepts ANY bit pattern, so a corrupt
/// `weights.bin` carrying NaN/inf would flow untouched into the forward pass and
/// silently yield NaN confidences (softmax over a NaN sum, NaN-propagating
/// sigmoid). The one-time finite scan closes that value-domain gap: it is a
/// single pass at first model access, negligible against the forward passes it
/// guards, and catches corruption the size check cannot.
pub(crate) fn parse_weights(raw: &[u8]) -> Result<Box<[f32]>, String> {
    if raw.len() != TOTAL_F32_COUNT * 4 {
        return Err(format!(
            "weights.bin size {} bytes does not match expected {} bytes ({} f32 values)",
            raw.len(),
            TOTAL_F32_COUNT * 4,
            TOTAL_F32_COUNT
        ));
    }
    let parsed: Box<[f32]> = raw
        .chunks_exact(4)
        .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
        .collect();
    if let Some(index) = parsed.iter().position(|weight| !weight.is_finite()) {
        return Err(format!(
            "weights.bin f32 at index {index} is non-finite (NaN/inf); embedded model is corrupt"
        ));
    }
    Ok(parsed)
}

fn all_weights() -> &'static [f32] {
    static PARSED: std::sync::OnceLock<Box<[f32]>> = std::sync::OnceLock::new();

    PARSED.get_or_init(|| match parse_weights(WEIGHTS) {
        Ok(weights) => weights,
        // Fail closed: `weights.bin` is `include_bytes!`-embedded Tier-B data, so
        // an invalid buffer is a build/corruption bug, not a runtime input.
        Err(error) => panic!("{error}"),
    })
}

fn load_f32_slice(offset: usize, count: usize) -> &'static [f32] {
    &all_weights()[offset..offset + count]
}

/// One expert's three dense layers as `&'static` slices into the parsed weight
/// buffer. Sliced ONCE at first model access (see [`model`]) instead of being
/// re-bounds-checked and re-sliced out of the flat buffer on every candidate.
///
/// `fc1_weight_t`/`fc2_weight_t` are the dense-layer weights in COLUMN-major
/// (transposed) layout, input `k`'s fan-out to every output packed contiguously
/// at `[k*OUT .. (k+1)*OUT]`. The forward pass (`ml_scorer.rs`) runs these layers
/// "output-stationary": for each input `k`, scale that contiguous output row by
/// `input[k]` and accumulate into the `OUT` running sums. The inner loop over
/// outputs is dependency-free (each `acc[o]` is independent), so the compiler
/// auto-vectorizes it across the output lanes. WITHOUT reassociating any single
/// output's reduction. Each `acc[o]` still sums its inputs in `k = 0,1,..,IN-1`
/// order with separate round(mul) then round(add) steps (no FMA fusion), so the
/// result is BIT-IDENTICAL to the row-major scalar dot product, not merely close.
/// (A previous AVX2+FMA attempt fused the multiply-add and reassociated lanes,
/// which was sub-ULP DIFFERENT and regressed ~30 ML-gated contracts; this layout
/// gets the SIMD width with none of that divergence.) `fc3_weight` stays ROW-major
/// (a single output (nothing to vectorize across), as does the 6-output gate).
pub(crate) struct ExpertWeights {
    pub fc1_weight_t: &'static [f32],
    pub fc1_bias: &'static [f32],
    pub fc2_weight_t: &'static [f32],
    pub fc2_bias: &'static [f32],
    pub fc3_weight: &'static [f32],
    pub fc3_bias: f32,
}

/// The full mixture-of-experts model with every weight/bias slice resolved
/// once. The CPU forward pass borrows `&'static MoeModel` a single time per
/// call and indexes these fields directly, so the 37 per-candidate accessor +
/// `OnceLock`-acquire + re-slice calls collapse to one acquire and zero
/// re-slicing. The slices are the SAME bytes the per-layer accessors return,
/// so scores are bit-identical to the pre-hoist path.
pub(crate) struct MoeModel {
    pub gate_weight: &'static [f32],
    pub gate_bias: &'static [f32],
    pub experts: [ExpertWeights; EXPERT_COUNT],
}

/// Transpose a row-major `rows x cols` weight matrix into a `'static`
/// column-major `cols x rows` buffer: element `(o, k)` at row-major
/// `[o*cols + k]` moves to column-major `[k*rows + o]`. Computed exactly once
/// per layer during `model()` init and leaked to `'static` - the model lives
/// for the whole process, so this is a one-time allocation, not a per-scan one.
/// The values are copied verbatim (no arithmetic), so the transposed buffer is
/// bit-identical to the original weights, just reordered for the output-stationary
/// forward kernel in `ml_scorer.rs`.
fn transpose_static(src: &[f32], rows: usize, cols: usize) -> &'static [f32] {
    assert_eq!(src.len(), rows * cols, "transpose dimensions must match");
    let mut out = vec![0.0f32; rows * cols];
    for o in 0..rows {
        for k in 0..cols {
            out[k * rows + o] = src[o * cols + k];
        }
    }
    Box::leak(out.into_boxed_slice())
}

/// Return the process-wide hoisted MoE model. The slices are computed exactly
/// once (behind the same `OnceLock` as the parsed buffer) and shared for the
/// life of the process; every subsequent call is a single atomic-acquire load
/// of an already-initialized pointer.
pub(crate) fn model() -> &'static MoeModel {
    static MODEL: std::sync::OnceLock<MoeModel> = std::sync::OnceLock::new();
    MODEL.get_or_init(|| {
        // Touch the parsed buffer once so the per-layer accessors below resolve
        // against an initialized `&'static [f32]`.
        let _ = all_weights(); // LAW10: forces lazy-static/regex eager init (warm-up); not a fallback
        let experts = std::array::from_fn(|expert_idx| ExpertWeights {
            fc1_weight_t: transpose_static(
                expert_fc1_weight(expert_idx),
                EXPERT_FC1_OUT,
                INPUT_DIM,
            ),
            fc1_bias: expert_fc1_bias(expert_idx),
            fc2_weight_t: transpose_static(
                expert_fc2_weight(expert_idx),
                EXPERT_FC2_OUT,
                EXPERT_FC1_OUT,
            ),
            fc2_bias: expert_fc2_bias(expert_idx),
            fc3_weight: expert_fc3_weight(expert_idx),
            fc3_bias: expert_fc3_bias(expert_idx)[0],
        });
        MoeModel {
            gate_weight: gate_weight(),
            gate_bias: gate_bias(),
            experts,
        }
    })
}

/// Return the full flattened weight buffer (used by GPU batch inference).
#[cfg(feature = "gpu")]
pub(crate) fn all_weights_slice() -> &'static [f32] {
    all_weights()
}

/// Return the gate-layer weight matrix.
pub(crate) fn gate_weight() -> &'static [f32] {
    load_f32_slice(GATE_W_OFF, GATE_W_COUNT)
}

/// Return the gate-layer bias vector.
pub(crate) fn gate_bias() -> &'static [f32] {
    load_f32_slice(GATE_B_OFF, GATE_B_COUNT)
}

fn expert_base_offset(expert_idx: usize) -> usize {
    assert!(expert_idx < EXPERT_COUNT, "expert index out of range");
    EXPERTS_OFF + expert_idx * EXPERT_PARAM_COUNT
}

/// Return the first dense-layer weights for `expert_idx`.
pub(crate) fn expert_fc1_weight(expert_idx: usize) -> &'static [f32] {
    let base = expert_base_offset(expert_idx);
    load_f32_slice(base, EXPERT_FC1_W_COUNT)
}

/// Return the first dense-layer bias vector for `expert_idx`.
pub(crate) fn expert_fc1_bias(expert_idx: usize) -> &'static [f32] {
    let base = expert_base_offset(expert_idx) + EXPERT_FC1_W_COUNT;
    load_f32_slice(base, EXPERT_FC1_B_COUNT)
}

/// Return the second dense-layer weights for `expert_idx`.
pub(crate) fn expert_fc2_weight(expert_idx: usize) -> &'static [f32] {
    let base = expert_base_offset(expert_idx) + EXPERT_FC1_W_COUNT + EXPERT_FC1_B_COUNT;
    load_f32_slice(base, EXPERT_FC2_W_COUNT)
}

/// Return the second dense-layer bias vector for `expert_idx`.
pub(crate) fn expert_fc2_bias(expert_idx: usize) -> &'static [f32] {
    let base = expert_base_offset(expert_idx)
        + EXPERT_FC1_W_COUNT
        + EXPERT_FC1_B_COUNT
        + EXPERT_FC2_W_COUNT;
    load_f32_slice(base, EXPERT_FC2_B_COUNT)
}

/// Return the output-layer weights for `expert_idx`.
pub(crate) fn expert_fc3_weight(expert_idx: usize) -> &'static [f32] {
    let base = expert_base_offset(expert_idx)
        + EXPERT_FC1_W_COUNT
        + EXPERT_FC1_B_COUNT
        + EXPERT_FC2_W_COUNT
        + EXPERT_FC2_B_COUNT;
    load_f32_slice(base, EXPERT_FC3_W_COUNT)
}

/// Return the output-layer bias for `expert_idx`.
pub(crate) fn expert_fc3_bias(expert_idx: usize) -> &'static [f32] {
    let base = expert_base_offset(expert_idx)
        + EXPERT_FC1_W_COUNT
        + EXPERT_FC1_B_COUNT
        + EXPERT_FC2_W_COUNT
        + EXPERT_FC2_B_COUNT
        + EXPERT_FC3_W_COUNT;
    load_f32_slice(base, EXPERT_FC3_B_COUNT)
}