mod common;
use common::{
assert_all_three_paths_match_within, assert_decoder_matches_on_all_three_paths, graph_caches,
kl_vs_golden, load_graph_fixture, worst_vs, GRAPH_PROMPT,
};
use frink_models::capability::{resolve_architecture, ArchPath, BIASED_LAYER_NORM};
use frink_models::config::RopeLayout;
use frink_models::norm::NormOp;
use frink_models::parallel_residual::ParallelNorm;
use frink_models::proj_bias::{Presence, OUTPUT_BIAS_CREATORS};
use frink_models::{Decoder, FfnActivation};
const PHI2: &str = "phi2";
const PHI2_FUSED: &str = "phi2_fused";
const GELU_TABLE_TOL_BIASED: f32 = 1e-2;
const PHI2_GOLDEN: [f32; 48] = [
-2.703486,
-6.481769,
-1.7040737,
1.4405402,
1.0704855,
1.539808,
-4.351997,
-1.717865,
3.3333216,
1.2004324,
2.103045,
-1.4339964,
1.5129533,
-3.1447134,
-1.5147669,
3.5478418,
1.0659884,
2.2116117,
-0.15988258,
-6.2265286,
0.7879978,
1.7139347,
1.8020049,
1.3367853,
2.9507613,
-1.4044921,
-2.086196,
0.06277639,
5.179845,
3.6416256,
-1.7835414,
-1.7161425,
-1.1067437,
-0.42966536,
-0.56274307,
-0.045191407,
1.3812802,
-3.0431554,
2.634523,
3.1698017,
1.0008919,
1.6742024,
-1.3365207,
0.79682136,
-2.5168622,
-4.505355,
1.7743274,
0.75843775,
];
fn decode(decoder: &Decoder) -> Vec<f32> {
let mut kv = graph_caches(decoder);
let mut out = Vec::new();
for (pos, &tok) in GRAPH_PROMPT.iter().enumerate() {
out = decoder.forward_token(tok, pos, &mut kv);
}
out
}
#[test]
fn phi2_matches_llama_cpp_on_all_three_paths() {
assert_all_three_paths_match_within(PHI2, &PHI2_GOLDEN, GELU_TABLE_TOL_BIASED);
}
#[test]
fn the_fused_qkv_spelling_matches_the_same_golden() {
assert_all_three_paths_match_within(PHI2_FUSED, &PHI2_GOLDEN, GELU_TABLE_TOL_BIASED);
}
#[test]
fn report_kl_against_llama_cpp() {
let out = decode(&load_graph_fixture(PHI2));
println!(
"{PHI2}: KL(llama.cpp || frink) = {:.3e}, max |delta| = {:.3e}",
kl_vs_golden(&out, &PHI2_GOLDEN),
worst_vs(&out, &PHI2_GOLDEN)
);
}
#[test]
fn the_loaded_decoder_is_the_graph() {
assert!(matches!(
resolve_architecture(PHI2),
Some(ArchPath::GenericGqa {
rope: RopeLayout::Neox
})
));
assert!(BIASED_LAYER_NORM.contains(&PHI2));
assert!(OUTPUT_BIAS_CREATORS.contains(&(PHI2, Presence::Required)));
for name in [PHI2, PHI2_FUSED] {
let d = load_graph_fixture(name);
let bias = d.output_bias.as_ref().expect("phi2.cpp:22: REQUIRED");
assert_eq!(bias.len(), d.output_head.rows());
assert!(bias.iter().any(|b| b.abs() > 0.5), "the file's, not zeros");
assert!(d.config.parallel_residual);
assert_eq!(d.config.rope_dim, Some(4), "partial_rotary_factor 0.5 of 8");
assert_eq!(d.config.ffn_activation, FfnActivation::GeluUngated);
for layer in &d.layers {
assert_eq!(layer.moe.parallel, Some(ParallelNorm::SharedNorm));
assert!(matches!(
layer.attn.norm_weight,
NormOp::LayerNormBias { .. }
));
assert!(matches!(layer.moe.norm_weight, NormOp::None));
assert!(layer.attn.q_bias.is_some() && layer.attn.k_bias.is_some());
assert!(layer.attn.v_bias.is_some() && layer.attn.o_bias.is_some());
let b = layer.moe.dense_bias.as_ref().expect("REQUIRED");
assert!(b.up.is_some() && b.down.is_some() && b.gate.is_none());
}
assert!(matches!(d.final_norm, NormOp::LayerNormBias { .. }));
}
}
#[test]
fn the_output_bias_and_the_residual_are_visible_in_the_logits() {
let mut d = load_graph_fixture(PHI2);
assert_decoder_matches_on_all_three_paths(&d, &PHI2_GOLDEN, GELU_TABLE_TOL_BIASED, "baseline");
let saved = d.output_bias.take();
let worst = worst_vs(&decode(&d), &PHI2_GOLDEN);
assert!(worst > 1.0, "output.bias not seen: {worst}");
d.output_bias = saved;
for l in d.layers.iter_mut() {
l.moe.parallel = None;
}
let worst = worst_vs(&decode(&d), &PHI2_GOLDEN);
assert!(worst > 1e-1, "the parallel residual not seen: {worst}");
for l in d.layers.iter_mut() {
l.moe.parallel = Some(ParallelNorm::SharedNorm);
}
assert_decoder_matches_on_all_three_paths(&d, &PHI2_GOLDEN, GELU_TABLE_TOL_BIASED, "restored");
}