#![allow(clippy::excessive_precision)]
mod common;
use common::{
assert_all_three_paths_match, assert_decoder_matches_on_all_three_paths, graph_caches,
graph_fixture_path, kl_vs_golden, load_graph_fixture, worst_vs, GRAPH_PROMPT, GRAPH_TOL,
};
use frink_gguf::GgufFile;
use frink_models::lora::{LoraAdapter, LoraError};
use frink_models::Decoder;
const BASE: &str = "lora_base";
fn base_file() -> GgufFile {
GgufFile::open(graph_fixture_path(BASE)).expect("base fixture opens")
}
fn adapter(name: &str) -> LoraAdapter {
LoraAdapter::open(graph_fixture_path(name)).expect("adapter fixture parses")
}
fn adapted(adapters: &[(&str, f32)]) -> Decoder {
let mut d = load_graph_fixture(BASE);
let base = base_file();
for (i, (name, scale)) in adapters.iter().enumerate() {
let id = d
.attach_lora(&base, adapter(name), *scale)
.unwrap_or_else(|e| panic!("{name} attaches: {e}"));
assert_eq!(id, i);
}
d
}
fn prefill(d: &Decoder) -> Vec<f32> {
let mut kv = graph_caches(d);
d.forward_batch_last(&GRAPH_PROMPT, 0, &mut kv)
}
const BASE_GOLDEN: [f32; 48] = [
-0.370904267,
-1.12565887,
-1.74006295,
-1.08758676,
0.476774454,
0.600389838,
-0.589864194,
1.80681932,
2.86660552,
-1.96880364,
1.63735247,
1.32705796,
-1.13267052,
0.198578566,
0.247464299,
0.335818946,
-0.311274469,
1.67584085,
-0.893444657,
0.864701152,
-0.48749727,
-0.480290055,
-0.626489639,
0.805029571,
-0.804177046,
-2.09429097,
1.3663398,
0.76220727,
0.343970269,
0.346324265,
-2.55922723,
0.282697946,
1.32494926,
0.266927719,
0.101378471,
0.038739264,
-1.10674012,
-0.653178751,
0.23567237,
-1.41555858,
-0.80852747,
1.30809903,
1.44235611,
-0.389747143,
0.473889053,
0.992623746,
0.52838558,
0.232206672,
];
const A_GOLDEN: [f32; 48] = [
-0.536873817,
-0.865892172,
-1.06427574,
-0.360642791,
-1.36866879,
1.12833381,
0.17011711,
0.503298402,
-0.845693052,
2.09086847,
-0.953029752,
-0.791327715,
-0.358710736,
0.315385282,
0.289216399,
0.27945292,
-2.62314343,
0.0533467531,
1.24105656,
0.890393138,
-1.35349703,
-0.287334263,
-0.469020128,
-0.771149993,
-2.01075006,
1.36664784,
-1.09304082,
0.679412365,
1.54256582,
2.4239862,
4.14771175,
0.736795902,
-2.54307985,
-1.21588445,
-0.00577360392,
-1.08792377,
1.53952384,
-0.979934931,
-0.481251597,
-0.00689542294,
0.65814662,
-0.798549235,
2.45526719,
-1.45747113,
1.3203969,
-0.118594617,
-0.341794401,
1.13295007,
];
const A_HALF_GOLDEN: [f32; 48] = [
-1.95309675,
-0.106032044,
-1.12723637,
-1.15716851,
1.95611238,
0.0175241828,
-1.85523832,
3.58392668,
1.32248962,
-2.0307951,
1.57083809,
-0.678562164,
-1.04460394,
-0.578685462,
-0.923409402,
-0.24875319,
0.280913889,
1.27480114,
-0.646233439,
0.940256894,
-0.606794596,
0.96797061,
0.0689299703,
1.12761223,
-1.24981856,
-1.2279743,
1.61927629,
-0.914031684,
-1.17003334,
1.57842183,
-3.02034569,
-0.514553428,
0.748880684,
-1.16426277,
-0.67561686,
-0.457690895,
-1.94772732,
0.925445795,
-1.28252268,
-0.844323993,
-1.51393533,
2.83296156,
1.16487026,
-0.302999437,
1.01954293,
1.14220178,
1.09462118,
1.08204854,
];
const AB_GOLDEN: [f32; 48] = [
0.218488634,
-1.14548278,
-0.832244158,
0.547200024,
-1.21331131,
1.20635152,
-0.0414975584,
-0.383256972,
0.303418756,
1.25630832,
-0.656019807,
0.641156316,
-1.11815083,
0.023001194,
0.149316788,
0.73910141,
-2.91507649,
-0.539863348,
-0.0998753458,
-0.655889273,
-1.65785527,
-0.370646477,
-0.215931475,
-0.292460054,
-1.99574506,
1.58345556,
-1.40718293,
0.145563021,
1.00802922,
2.88666725,
3.89027452,
0.748188913,
-1.55383968,
-0.871917307,
0.683666408,
-0.218239486,
0.712480366,
-0.322091937,
0.772425592,
-0.186063498,
-0.685799718,
-0.147196501,
2.36823893,
-0.336556911,
1.07483685,
0.118239105,
-0.546918273,
0.52361685,
];
const B_GOLDEN: [f32; 48] = [
-0.103938207,
-0.457222104,
-0.941307843,
-0.645068645,
-0.382739455,
-0.627984405,
-1.25409985,
0.0547681004,
2.35386777,
-0.943867266,
2.02643752,
1.08659089,
-0.869225919,
-0.964745045,
-0.239215657,
-0.0702250451,
0.344786733,
0.671282232,
0.016603034,
0.9461115,
1.2614044,
0.358712792,
0.344278932,
0.763074517,
0.119232848,
-1.70412064,
2.14636588,
-0.861565888,
0.413673908,
0.517497182,
-2.41648054,
0.81807375,
2.21736312,
0.0818412825,
-0.126288608,
1.29192376,
-1.29586041,
0.0498075821,
-1.01258016,
-0.805461645,
0.191048965,
1.71336377,
0.0372688174,
-1.38677239,
1.31535792,
0.72845459,
0.109508976,
0.806996822,
];
#[test]
fn the_base_matches_llama_cpp_on_all_three_paths() {
assert_all_three_paths_match(BASE, &BASE_GOLDEN);
}
#[test]
fn adapter_a_matches_llama_cpp_on_all_three_paths() {
let d = adapted(&[("lora_a", 1.0)]);
assert_decoder_matches_on_all_three_paths(&d, &A_GOLDEN, GRAPH_TOL, "lora_a");
}
#[test]
fn adapter_a_at_half_scale_matches_llama_cpp() {
let d = adapted(&[("lora_a", 0.5)]);
assert_decoder_matches_on_all_three_paths(&d, &A_HALF_GOLDEN, GRAPH_TOL, "lora_a:0.5");
}
#[test]
fn adapter_b_matches_llama_cpp_on_all_three_paths() {
let d = adapted(&[("lora_b", 1.0)]);
assert_decoder_matches_on_all_three_paths(&d, &B_GOLDEN, GRAPH_TOL, "lora_b");
}
#[test]
fn two_adapters_at_two_scales_match_llama_cpp() {
let d = adapted(&[("lora_a", 1.0), ("lora_b", 0.75)]);
assert_decoder_matches_on_all_three_paths(&d, &AB_GOLDEN, GRAPH_TOL, "lora_a:1+lora_b:0.75");
}
#[test]
fn report_kl_against_llama_cpp() {
for (label, adapters, golden) in [
("base", &[][..], &BASE_GOLDEN),
("lora_a", &[("lora_a", 1.0)][..], &A_GOLDEN),
("lora_a:0.5", &[("lora_a", 0.5)][..], &A_HALF_GOLDEN),
("lora_b", &[("lora_b", 1.0)][..], &B_GOLDEN),
(
"lora_a:1 + lora_b:0.75",
&[("lora_a", 1.0), ("lora_b", 0.75)][..],
&AB_GOLDEN,
),
] {
let got = prefill(&adapted(adapters));
let kl = kl_vs_golden(&got, golden);
let worst = worst_vs(&got, golden);
println!("| `{label}` | {kl:.2e} | {worst:.2e} |");
assert!(kl < 1e-8, "{label}: KL {kl}");
}
}
#[test]
fn scale_zero_is_byte_identical_to_no_adapter() {
let plain = load_graph_fixture(BASE);
let before = prefill(&plain);
assert_eq!(before, prefill(&adapted(&[("lora_a", 0.0)])));
assert_eq!(
before,
prefill(&adapted(&[("lora_a", 0.0), ("lora_b", 0.0)]))
);
}
#[test]
fn attaching_then_zeroing_returns_the_pre_attach_logits() {
let mut d = load_graph_fixture(BASE);
let before = prefill(&d);
d.attach_lora(&base_file(), adapter("lora_a"), 1.0).unwrap();
assert!(
worst_vs(&prefill(&d), &before) > 1e-2,
"the adapter must be visible at 1.0"
);
d.set_lora_scales(&[(0, 0.0)]).unwrap();
assert_eq!(prefill(&d), before);
}
#[test]
fn scales_change_at_apply_time_without_a_reload() {
let d = adapted(&[("lora_a", 1.0), ("lora_b", 1.0)]);
assert_eq!(d.lora_scales(), vec![1.0, 1.0]);
d.set_lora_scales(&[(0, 0.5)]).unwrap();
assert_eq!(
d.lora_scales(),
vec![0.5, 0.0],
"an unlisted adapter goes to 0"
);
assert_decoder_matches_on_all_three_paths(&d, &A_HALF_GOLDEN, GRAPH_TOL, "a:0.5 via set");
d.set_lora_scales(&[(0, 1.0), (1, 0.75)]).unwrap();
assert_decoder_matches_on_all_three_paths(&d, &AB_GOLDEN, GRAPH_TOL, "a:1,b:0.75 via set");
let err = d.set_lora_scales(&[(2, 1.0)]).unwrap_err();
assert!(err.contains("out of range"), "{err}");
assert_eq!(
d.lora_scales(),
vec![1.0, 0.75],
"a refused set changes nothing"
);
}
#[test]
fn the_attached_list_is_what_the_server_will_report() {
let d = adapted(&[("lora_a", 1.0), ("lora_b", 0.25)]);
assert!(d.lora_attached());
assert_eq!(d.lora_adapters.len(), 2);
let a = &d.lora_adapters[0];
assert!(a.path.ends_with("lora_a_tiny.gguf"));
assert_eq!(a.alpha, 8.0);
assert_eq!(
a.n_tensors, 16,
"7 projections x 2 layers + embedding + head"
);
assert_eq!(a.scale(), 1.0);
let b = &d.lora_adapters[1];
assert_eq!(b.alpha, 2.0);
assert_eq!(b.n_tensors, 4);
assert_eq!(b.scale(), 0.25);
assert!(!load_graph_fixture(BASE).lora_attached());
}
#[test]
fn the_adapter_decorates_exactly_the_matrices_it_names() {
let d = adapted(&[("lora_b", 1.0)]);
for (il, layer) in d.layers.iter().enumerate() {
assert_eq!(
layer.attn.q_proj.lora().map(|l| l.len()),
Some(1),
"blk.{il} q"
);
assert_eq!(
layer.attn.v_proj.lora().map(|l| l.len()),
Some(1),
"blk.{il} v"
);
assert!(
layer.attn.k_proj.lora().is_none(),
"blk.{il} k is not in lora_b"
);
assert!(
layer.attn.o_proj.lora().is_none(),
"blk.{il} o is not in lora_b"
);
}
assert!(d.embedding.lora().is_none());
assert!(d.output_head.lora().is_none());
let d = adapted(&[("lora_a", 1.0), ("lora_b", 1.0)]);
assert_eq!(
d.layers[0].attn.q_proj.lora().map(|l| l.len()),
Some(2),
"both on Q"
);
assert_eq!(
d.layers[0].attn.k_proj.lora().map(|l| l.len()),
Some(1),
"only A on K"
);
assert_eq!(d.embedding.lora().map(|l| l.len()), Some(1));
assert_eq!(d.output_head.lora().map(|l| l.len()), Some(1));
assert_eq!(
d.embedding.rows(),
48,
"the adapted matrix keeps the base's shape"
);
assert_eq!(d.embedding.cols(), 24);
}
#[test]
fn the_embedding_pair_is_the_flipped_one() {
let d = adapted(&[("lora_a", 1.0)]);
let a = adapter("lora_a");
let pair = &a.pairs["token_embd.weight"];
let base = load_graph_fixture(BASE);
let rank = pair.a.cols();
for t in [3usize, 7, 11] {
let want: Vec<f32> = (0..24)
.map(|e| {
base.embedding.dequant_row(t)[e]
+ 2.0
* (0..rank)
.map(|k| pair.b.data[e * rank + k] * pair.a.data[t * rank + k])
.sum::<f32>()
})
.collect();
let got = d.embedding.dequant_row(t);
for (g, w) in got.iter().zip(&want) {
assert!((g - w).abs() < 1e-5, "token {t}: {g} vs {w}");
}
}
}
#[test]
fn dropping_alpha_over_rank_diverges_from_llama_cpp() {
let mut a = adapter("lora_a");
a.alpha = 0.0;
let mut d = load_graph_fixture(BASE);
d.attach_lora(&base_file(), a, 1.0).unwrap();
let worst = worst_vs(&prefill(&d), &A_GOLDEN);
assert!(worst > 1e-2, "alpha/rank moved the output by only {worst}");
let mut a = adapter("lora_a");
a.alpha = 0.0;
let mut d = load_graph_fixture(BASE);
d.attach_lora(&base_file(), a, 2.0).unwrap();
assert_decoder_matches_on_all_three_paths(&d, &A_GOLDEN, GRAPH_TOL, "alpha folded into scale");
}
fn attach_err(base: &str, adapter_name: &str) -> LoraError {
let mut d = load_graph_fixture(base);
let file = GgufFile::open(graph_fixture_path(base)).unwrap();
match d.attach_lora(&file, adapter(adapter_name), 1.0) {
Ok(_) => panic!("{adapter_name} on {base} must be refused"),
Err(e) => e,
}
}
#[test]
fn an_adapter_for_another_architecture_is_refused() {
let err = attach_err("arcee", "lora_b");
assert!(matches!(err, LoraError::ArchMismatch { .. }), "{err}");
let msg = err.to_string();
assert!(msg.contains("\"llama\""), "{msg}");
assert!(msg.contains("\"arcee\""), "{msg}");
}
#[test]
fn an_adapter_naming_a_tensor_the_base_lacks_is_refused() {
let err = attach_err("lora_base_tied", "lora_a");
assert!(matches!(err, LoraError::NotInBase { .. }), "{err}");
assert!(
err.to_string()
.contains("'output.weight' does not exist in base model"),
"{err}"
);
}
#[test]
fn the_embedding_pair_on_a_tied_head_is_refused_like_llama_cpp_aborts() {
let mut d = load_graph_fixture(BASE);
d.attach_lora(&base_file(), adapter("lora_e"), 1.0).unwrap();
let err = attach_err("lora_base_tied", "lora_e");
assert!(matches!(err, LoraError::TiedHead { .. }), "{err}");
assert!(err.to_string().contains("ggml_can_mul_mat"), "{err}");
}
#[test]
fn a_pair_of_the_wrong_shape_is_refused_naming_every_dimension() {
let mut a = adapter("lora_b");
let pair = a.pairs.get_mut("blk.0.attn_q.weight").unwrap();
pair.b.data.truncate(12 * 2);
pair.b.shape = [12, 2];
let mut d = load_graph_fixture(BASE);
let err = d.attach_lora(&base_file(), a, 1.0).unwrap_err();
assert!(matches!(err, LoraError::Shape { .. }), "{err}");
let msg = err.to_string();
assert!(msg.contains("incorrect shape"), "{msg}");
assert!(msg.contains("[24 x 24]"), "{msg}");
assert!(msg.contains("[12, 2]"), "{msg}");
}
#[test]
fn an_untransposed_lora_a_is_refused() {
let mut a = adapter("lora_b");
let pair = a.pairs.get_mut("blk.1.attn_v.weight").unwrap();
pair.a.shape = [3, 16];
let mut d = load_graph_fixture(BASE);
let err = d.attach_lora(&base_file(), a, 1.0).unwrap_err();
assert!(matches!(err, LoraError::Shape { .. }), "{err}");
let mut a = adapter("lora_b");
let pair = a.pairs.get_mut("blk.1.attn_v.weight").unwrap();
pair.a.data.extend(std::iter::repeat_n(0.0, 24));
pair.a.shape = [3, 24];
let mut d = load_graph_fixture(BASE);
let err = d.attach_lora(&base_file(), a, 1.0).unwrap_err();
assert!(matches!(err, LoraError::NotTransposed { .. }), "{err}");
assert!(err.to_string().contains("not transposed"), "{err}");
}
#[test]
fn an_adapter_on_routed_experts_is_refused_by_name() {
let mut a = adapter("lora_b");
a.arch = "qwen3moe".to_string();
let pair = a.pairs.remove("blk.0.attn_q.weight").unwrap();
a.pairs.clear();
a.pairs
.insert("blk.0.ffn_gate_exps.weight".to_string(), pair);
let base = graph_fixture_path("qwen3moe");
let file = GgufFile::open(&base).unwrap();
let mut d = load_graph_fixture("qwen3moe");
let err = d.attach_lora(&file, a, 1.0).unwrap_err();
assert!(matches!(err, LoraError::RoutedExperts { .. }), "{err}");
assert!(err.to_string().contains("build_lora_mm_id"), "{err}");
}
#[test]
fn a_refused_attach_leaves_the_earlier_adapter_in_place() {
let mut d = adapted(&[("lora_a", 1.0)]);
let tied = GgufFile::open(graph_fixture_path("lora_base_tied")).unwrap();
assert!(d.attach_lora(&tied, adapter("lora_e"), 1.0).is_err());
assert_eq!(d.lora_adapters.len(), 1);
assert_decoder_matches_on_all_three_paths(&d, &A_GOLDEN, GRAPH_TOL, "lora_a after a refusal");
}