mod data;
use data::{
BLOCKS_BF16, IMATRIX_F32, LLAMA_CPP_Q4_K_IMATRIX_GOLDEN, LLAMA_CPP_Q5_K_IMATRIX_GOLDEN,
LLAMA_CPP_Q6_K_IMATRIX_GOLDEN,
};
pub(crate) fn blocks() -> Vec<f32> {
BLOCKS_BF16
.iter()
.map(|&b| f32::from_bits(u32::from(b) << 16))
.collect()
}
pub(crate) fn imatrix() -> Vec<f32> {
IMATRIX_F32.iter().map(|&b| f32::from_bits(b)).collect()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{encode_row_q4_k, encode_row_q5_k, encode_row_q6_k};
fn check(name: &str, block_bytes: usize, got: &[u8], want: &[u8]) {
assert_eq!(got.len(), want.len(), "{name}: length");
for (b, (g, w)) in got
.chunks(block_bytes)
.zip(want.chunks(block_bytes))
.enumerate()
{
assert_eq!(
g, w,
"{name}: super-block {b} disagrees with llama-quantize --imatrix"
);
}
}
#[test]
fn q4_k_with_an_imatrix_matches_llama_quantize_imatrix() {
let x = blocks();
let qw = imatrix();
let mut got = Vec::new();
for (row, w) in x.chunks(256).zip(qw.chunks(256)) {
encode_row_q4_k(row, Some(w), &mut got).unwrap();
}
check("Q4_K", 144, &got, &LLAMA_CPP_Q4_K_IMATRIX_GOLDEN);
}
#[test]
fn q5_k_with_an_imatrix_matches_llama_quantize_imatrix() {
let x = blocks();
let qw = imatrix();
let mut got = Vec::new();
for (row, w) in x.chunks(256).zip(qw.chunks(256)) {
encode_row_q5_k(row, Some(w), &mut got).unwrap();
}
check("Q5_K", 176, &got, &LLAMA_CPP_Q5_K_IMATRIX_GOLDEN);
}
#[test]
fn q6_k_with_an_imatrix_matches_llama_quantize_imatrix() {
let x = blocks();
let qw = imatrix();
let mut got = Vec::new();
for (row, w) in x.chunks(256).zip(qw.chunks(256)) {
encode_row_q6_k(row, Some(w), &mut got).unwrap();
}
check("Q6_K", 210, &got, &LLAMA_CPP_Q6_K_IMATRIX_GOLDEN);
}
#[test]
fn the_imatrix_actually_changes_every_format() {
let x = blocks();
let qw = imatrix();
for (name, enc) in [
(
"Q4_K",
encode_row_q4_k as fn(&[f32], Option<&[f32]>, &mut Vec<u8>) -> Option<()>,
),
("Q5_K", encode_row_q5_k),
("Q6_K", encode_row_q6_k),
] {
let (mut plain, mut weighted) = (Vec::new(), Vec::new());
enc(&x[..1024], None, &mut plain).unwrap();
enc(&x[..1024], Some(&qw[..1024]), &mut weighted).unwrap();
assert_ne!(
plain, weighted,
"{name}: the imatrix had no effect on this row"
);
}
}
}