use frink_gguf::GgmlType;
#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Target {
Q8_0,
Q4_K_S,
Q4_K_M,
Q5_K_S,
Q5_K_M,
Q6_K,
}
impl Target {
pub const ALL: &'static [Target] = &[
Target::Q8_0,
Target::Q4_K_S,
Target::Q4_K_M,
Target::Q5_K_S,
Target::Q5_K_M,
Target::Q6_K,
];
pub fn name(self) -> &'static str {
match self {
Target::Q8_0 => "Q8_0",
Target::Q4_K_S => "Q4_K_S",
Target::Q4_K_M => "Q4_K_M",
Target::Q5_K_S => "Q5_K_S",
Target::Q5_K_M => "Q5_K_M",
Target::Q6_K => "Q6_K",
}
}
pub fn ggml_type(self) -> GgmlType {
match self {
Target::Q8_0 => GgmlType::Q8_0,
Target::Q4_K_S | Target::Q4_K_M => GgmlType::Q4K,
Target::Q5_K_S | Target::Q5_K_M => GgmlType::Q5K,
Target::Q6_K => GgmlType::Q6K,
}
}
pub fn llama_ftype(self) -> u32 {
match self {
Target::Q8_0 => 7, Target::Q4_K_S => 14, Target::Q4_K_M => 15, Target::Q5_K_S => 16, Target::Q5_K_M => 17, Target::Q6_K => 18, }
}
pub fn fallback_note(self) -> &'static str {
match self {
Target::Q8_0 => "llama.cpp has no fallback type for Q8_0 either -- it stops here too.",
Target::Q4_K_S | Target::Q4_K_M => {
"llama.cpp answers this by changing the tensor's TYPE (Q4_K -> Q5_0, or F16 if the \
row is not a multiple of 32 either); frink can write neither, so it stops rather \
than write a file whose name says Q4_K."
}
Target::Q5_K_S | Target::Q5_K_M => {
"llama.cpp answers this by changing the tensor's TYPE (Q5_K -> Q5_1, or F16 if the \
row is not a multiple of 32 either); frink can write neither, so it stops rather \
than write a file whose name says Q5_K."
}
Target::Q6_K => {
"llama.cpp answers this by changing the tensor's TYPE (Q6_K -> Q8_0, or F16 if the \
row is not a multiple of 32 either); frink stops rather than write a tensor \
whose type disagrees with the plan it printed."
}
}
}
}
const LLAMA_CPP_TARGETS: &[(&str, Option<Target>)] = &[
("q1_0", None),
("q2_0", None),
("q4_0", None),
("q4_1", None),
("mxfp4_moe", None),
("q5_0", None),
("q5_1", None),
("iq2_xxs", None),
("iq2_xs", None),
("iq2_s", None),
("iq2_m", None),
("iq1_s", None),
("iq1_m", None),
("tq1_0", None),
("tq2_0", None),
("q2_k", None),
("q2_k_s", None),
("iq3_xxs", None),
("iq3_s", None),
("iq3_m", None),
("q3_k", None),
("iq3_xs", None),
("q3_k_s", None),
("q3_k_m", None),
("q3_k_l", None),
("iq4_nl", None),
("iq4_xs", None),
("q4_k", Some(Target::Q4_K_M)),
("q4_k_s", Some(Target::Q4_K_S)),
("q4_k_m", Some(Target::Q4_K_M)),
("q5_k", Some(Target::Q5_K_M)),
("q5_k_s", Some(Target::Q5_K_S)),
("q5_k_m", Some(Target::Q5_K_M)),
("q6_k", Some(Target::Q6_K)),
("q8_0", Some(Target::Q8_0)),
("f16", None),
("bf16", None),
("f32", None),
("copy", None),
];
#[derive(Debug, PartialEq, Eq)]
pub enum TargetRefusal {
NotWritableYet(String),
Unknown(String),
}
impl std::fmt::Display for TargetRefusal {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let writable = writable_targets();
match self {
TargetRefusal::NotWritableYet(name) => write!(
f,
"frink cannot WRITE {name} yet. It can read {name} and run it; it has no encoder \
for it.\n\
`frink quantize` writes: {writable}.\n\
The remaining K-quant and IQ encoders are each an iterative per-super-block \
scale/min fit (and, for the IQ tiers, a lattice search over a codebook). A \
min/max approximation of one produces a file that loads and generates measurably \
worse text, so frink stops here instead of writing it. Use llama.cpp's \
`llama-quantize --type {name}` for now; frink reads what it produces."
),
TargetRefusal::Unknown(name) => write!(
f,
"'{name}' is not a quantization type. `frink quantize` writes: {writable}."
),
}
}
}
pub fn writable_targets() -> String {
Target::ALL
.iter()
.map(|t| t.name())
.collect::<Vec<_>>()
.join(", ")
}
pub fn parse_target(raw: &str) -> Result<Target, TargetRefusal> {
let lower = raw.to_ascii_lowercase();
let Some((_, encoder)) = LLAMA_CPP_TARGETS.iter().find(|(n, _)| *n == lower) else {
return Err(TargetRefusal::Unknown(raw.to_string()));
};
let Some(target) = encoder else {
return Err(TargetRefusal::NotWritableYet(raw.to_string()));
};
Ok(*target)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Disposition {
Quantize(GgmlType),
Copy(&'static str),
}
const KEEP_AT_SOURCE_PRECISION: &[(&str, &str)] = &[
("_norm.weight", "norm"),
("ffn_gate_inp.weight", "expert gating"),
(
"ffn_gate_tid2eid.weight",
"token-id -> expert-id routing table",
),
("altup", "tiny"),
("laurel", "tiny"),
("per_layer_model_proj", "small"),
("position_embd.weight", "positional embedding"),
("token_types.weight", "token types"),
("ssm_conv1d", "conv1d kernel"),
("shortconv.conv.weight", "conv kernel"),
("indexer.k_proj.weight", "tiny"),
("indexer.q_proj.weight", "tiny"),
("time_mix_first.weight", "RWKV small 2D"),
("time_mix_w0.weight", "RWKV small 2D"),
("time_mix_w1.weight", "RWKV small 2D"),
("time_mix_w2.weight", "RWKV small 2D"),
("time_mix_v0.weight", "RWKV small 2D"),
("time_mix_v1.weight", "RWKV small 2D"),
("time_mix_v2.weight", "RWKV small 2D"),
("time_mix_a0.weight", "RWKV small 2D"),
("time_mix_a1.weight", "RWKV small 2D"),
("time_mix_a2.weight", "RWKV small 2D"),
("time_mix_g1.weight", "RWKV small 2D"),
("time_mix_g2.weight", "RWKV small 2D"),
("time_mix_decay_w1.weight", "RWKV small 2D"),
("time_mix_decay_w2.weight", "RWKV small 2D"),
("time_mix_lerp_fused.weight", "RWKV small 2D"),
("attn_rel_b.weight", "relative position bias"),
(".position_embd", "positional embedding"),
("sam.pos_embd", "multimodal"),
("sam.neck.", "multimodal"),
("sam.net_", "multimodal"),
(".rel_pos", "multimodal"),
(".patch_embd", "multimodal"),
(".patch_merger", "multimodal"),
("a.rvq.codebook", "audio codebook"),
("mm.a.code_embd", "audio codebook"),
];
fn ggml_n_dims(shape: &[u64]) -> usize {
for i in (1..shape.len()).rev() {
if shape[i] > 1 {
return i + 1;
}
}
1
}
pub fn allows_quantization(name: &str, shape: &[u64]) -> Option<&'static str> {
if ggml_n_dims(shape) < 2 {
return Some("1-D");
}
if !name.ends_with("weight") {
return Some("not a weight");
}
for (needle, reason) in KEEP_AT_SOURCE_PRECISION {
if name.contains(needle) {
return Some(reason);
}
}
None
}
pub fn disposition(dtype: GgmlType, chosen: GgmlType) -> Disposition {
if dtype == chosen {
return Disposition::Copy("already the target type");
}
Disposition::Quantize(chosen)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_target_with_no_encoder_is_refused_by_name_and_says_what_can_be_written() {
let err = parse_target("q3_k_m").unwrap_err();
assert_eq!(err, TargetRefusal::NotWritableYet("q3_k_m".into()));
let msg = err.to_string();
assert!(msg.contains("cannot WRITE q3_k_m"), "{msg}");
assert!(msg.contains("writes: Q8_0"), "{msg}");
assert!(msg.contains("stops here instead of writing it"), "{msg}");
}
#[test]
fn the_k_quant_mixes_no_longer_need_pure() {
for name in [
"q4_k", "q4_k_s", "q4_k_m", "q5_k", "q5_k_s", "q5_k_m", "q6_k",
] {
assert!(parse_target(name).is_ok(), "{name}");
}
}
#[test]
fn q4_k_is_the_same_target_as_q4_k_m_the_way_llama_quantize_aliases_it() {
assert_eq!(parse_target("q4_k").unwrap(), Target::Q4_K_M);
assert_eq!(parse_target("q4_k_m").unwrap(), Target::Q4_K_M);
assert_eq!(parse_target("q4_k_s").unwrap(), Target::Q4_K_S);
assert_eq!(
Target::Q4_K_S.ggml_type(),
Target::Q4_K_M.ggml_type(),
"one encoder serves both"
);
assert_ne!(
Target::Q4_K_S.llama_ftype(),
Target::Q4_K_M.llama_ftype(),
"and they must not report the same mix"
);
}
#[test]
fn every_llama_cpp_target_frink_cannot_write_refuses_as_a_gap_not_a_typo() {
for (name, encoder) in LLAMA_CPP_TARGETS {
let parsed = parse_target(name);
match encoder {
Some(t) => assert_eq!(parsed.as_ref().ok(), Some(t), "{name} should be writable"),
None => assert_eq!(
parsed.unwrap_err(),
TargetRefusal::NotWritableYet((*name).to_string()),
"{name}"
),
}
}
}
#[test]
fn the_target_enum_and_the_llama_cpp_name_table_cover_each_other() {
for t in Target::ALL {
assert!(
LLAMA_CPP_TARGETS
.iter()
.any(|(n, e)| *n == t.name().to_ascii_lowercase() && *e == Some(*t)),
"{} is not reachable from the llama-quantize name table",
t.name()
);
}
for (name, encoder) in LLAMA_CPP_TARGETS {
if let Some(t) = encoder {
assert!(
Target::ALL.contains(t),
"{name} maps to a target that is not in Target::ALL"
);
}
}
}
#[test]
fn the_writable_summary_names_every_target() {
let s = writable_targets();
for t in Target::ALL {
assert!(s.contains(t.name()), "{s} is missing {}", t.name());
}
}
#[test]
fn a_name_that_is_not_a_quant_at_all_says_so() {
let err = parse_target("q4_k_ultra").unwrap_err();
assert_eq!(err, TargetRefusal::Unknown("q4_k_ultra".into()));
assert!(err.to_string().contains("is not a quantization type"));
}
#[test]
fn target_names_are_case_insensitive_like_llama_quantize() {
assert_eq!(parse_target("q8_0").unwrap(), Target::Q8_0);
assert_eq!(parse_target("Q8_0").unwrap(), Target::Q8_0);
assert_eq!(parse_target("Q4_K_M").unwrap(), Target::Q4_K_M);
assert_eq!(
parse_target("Q8_o").unwrap_err(),
TargetRefusal::Unknown("Q8_o".into())
);
}
#[test]
fn trailing_unit_dimensions_do_not_make_a_tensor_two_dimensional() {
assert_eq!(ggml_n_dims(&[4096]), 1);
assert_eq!(ggml_n_dims(&[4096, 1]), 1);
assert_eq!(ggml_n_dims(&[4096, 1, 1]), 1);
assert_eq!(ggml_n_dims(&[4096, 11008]), 2);
assert_eq!(ggml_n_dims(&[4096, 1, 8]), 3);
}
#[test]
fn the_tensors_llama_cpp_keeps_at_source_precision_are_kept() {
let two_d = [4096u64, 4096];
let cases: &[(&str, bool)] = &[
("blk.0.attn_q.weight", true),
("blk.0.ffn_down.weight", true),
("token_embd.weight", true),
("output.weight", true),
("blk.0.attn_norm.weight", false),
("output_norm.weight", false),
("blk.0.ffn_gate_inp.weight", false),
("blk.0.ssm_conv1d.weight", false),
("position_embd.weight", false),
("token_types.weight", false),
("blk.0.altup_proj.weight", false),
("blk.0.attn_q.bias", false),
("v.patch_embd.weight", false),
];
for (name, want_quantized) in cases {
let got = allows_quantization(name, &two_d);
assert_eq!(got.is_none(), *want_quantized, "{name} -> {got:?}");
}
assert_eq!(
allows_quantization("blk.0.attn_q.weight", &[4096]),
Some("1-D")
);
}
#[test]
fn a_tensor_already_in_the_target_type_is_copied_not_re_encoded() {
assert_eq!(
disposition(GgmlType::Q8_0, GgmlType::Q8_0),
Disposition::Copy("already the target type")
);
assert_eq!(
disposition(GgmlType::F16, GgmlType::Q6K),
Disposition::Quantize(GgmlType::Q6K)
);
}
}