ik_llama_cpp_2/quantize.rs
1//! Model quantization write-path over ik's `llama_model_quantize`.
2//!
3//! The anchor crate (`utilityai/llama-cpp-rs`, `llama-cpp-2`) ships no quantize
4//! wrapper, so this is implemented fresh in the crate's module-local style
5//! (own `thiserror` error, builder over the raw params struct, `as_raw` escape
6//! hatch), mirroring [`crate::gguf`] and [`crate::model::params`].
7//!
8//! Quantizing a source GGUF (usually f16/bf16/higher precision) produces a
9//! smaller GGUF at `output`. ik_llama.cpp extends the upstream params struct
10//! with many fork-specific per-tensor override fields (`attn_*_type`,
11//! `ffn_*_type`, `token_embedding_type`, …) plus extra flags; those keep the
12//! values from `llama_model_quantize_default_params()` unless the caller reaches
13//! through [`LlamaQuantizeParams::as_raw_mut`].
14
15use std::ffi::{CString, NulError};
16use std::path::Path;
17
18use ik_llama_cpp_sys as sys;
19
20/// Errors returned by [`quantize`].
21#[derive(Debug, thiserror::Error)]
22pub enum QuantizeError {
23 /// An input or output path was not valid UTF-8 (cannot become a C string).
24 #[error("quantize path was not valid UTF-8")]
25 InvalidPath,
26 /// A path contained an interior NUL byte.
27 #[error("quantize path contained an interior NUL byte")]
28 Nul(#[from] NulError),
29 /// `llama_model_quantize` returned a non-zero status code.
30 #[error("llama_model_quantize failed with status {0}")]
31 Quantize(u32),
32}
33
34/// Common `llama_ftype` target types, re-exported from the sys crate.
35///
36/// These are the raw `ik_llama_cpp_sys::LLAMA_FTYPE_*` constants (each a
37/// `llama_ftype`, i.e. a `u32`); the full set (84 variants) lives in the sys
38/// crate and any of them may be passed to [`LlamaQuantizeParams::with_ftype`].
39/// The list below is a curated subset covering the standard llama.cpp quants
40/// plus ik_llama.cpp's SOTA "K"/"KS"/"KT" and repacked "R4/R8" families.
41pub mod ftype {
42 pub use ik_llama_cpp_sys::llama_ftype;
43
44 // --- Standard llama.cpp float / legacy / K-quants ---
45 pub use ik_llama_cpp_sys::{
46 LLAMA_FTYPE_ALL_F32, LLAMA_FTYPE_MOSTLY_BF16, LLAMA_FTYPE_MOSTLY_F16,
47 LLAMA_FTYPE_MOSTLY_Q2_K, LLAMA_FTYPE_MOSTLY_Q2_K_S, LLAMA_FTYPE_MOSTLY_Q3_K_L,
48 LLAMA_FTYPE_MOSTLY_Q3_K_M, LLAMA_FTYPE_MOSTLY_Q3_K_S, LLAMA_FTYPE_MOSTLY_Q4_0,
49 LLAMA_FTYPE_MOSTLY_Q4_1, LLAMA_FTYPE_MOSTLY_Q4_K_M, LLAMA_FTYPE_MOSTLY_Q4_K_S,
50 LLAMA_FTYPE_MOSTLY_Q5_0, LLAMA_FTYPE_MOSTLY_Q5_1, LLAMA_FTYPE_MOSTLY_Q5_K_M,
51 LLAMA_FTYPE_MOSTLY_Q5_K_S, LLAMA_FTYPE_MOSTLY_Q6_K, LLAMA_FTYPE_MOSTLY_Q8_0,
52 };
53
54 // --- Standard llama.cpp i-quants ---
55 pub use ik_llama_cpp_sys::{
56 LLAMA_FTYPE_MOSTLY_IQ1_M, LLAMA_FTYPE_MOSTLY_IQ1_S, LLAMA_FTYPE_MOSTLY_IQ2_M,
57 LLAMA_FTYPE_MOSTLY_IQ2_S, LLAMA_FTYPE_MOSTLY_IQ2_XS, LLAMA_FTYPE_MOSTLY_IQ2_XXS,
58 LLAMA_FTYPE_MOSTLY_IQ3_M, LLAMA_FTYPE_MOSTLY_IQ3_S, LLAMA_FTYPE_MOSTLY_IQ3_XXS,
59 LLAMA_FTYPE_MOSTLY_IQ4_NL, LLAMA_FTYPE_MOSTLY_IQ4_XS,
60 };
61
62 // --- ik_llama.cpp SOTA quants (the reason this fork exists) ---
63 pub use ik_llama_cpp_sys::{
64 LLAMA_FTYPE_MOSTLY_IQ1_KT, LLAMA_FTYPE_MOSTLY_IQ2_K, LLAMA_FTYPE_MOSTLY_IQ2_KL,
65 LLAMA_FTYPE_MOSTLY_IQ2_KS, LLAMA_FTYPE_MOSTLY_IQ2_KT, LLAMA_FTYPE_MOSTLY_IQ3_K,
66 LLAMA_FTYPE_MOSTLY_IQ3_KL, LLAMA_FTYPE_MOSTLY_IQ3_KS, LLAMA_FTYPE_MOSTLY_IQ3_KT,
67 LLAMA_FTYPE_MOSTLY_IQ4_K, LLAMA_FTYPE_MOSTLY_IQ4_KS, LLAMA_FTYPE_MOSTLY_IQ4_KSS,
68 LLAMA_FTYPE_MOSTLY_IQ4_KT, LLAMA_FTYPE_MOSTLY_IQ5_K, LLAMA_FTYPE_MOSTLY_IQ5_KS,
69 LLAMA_FTYPE_MOSTLY_IQ6_K, LLAMA_FTYPE_MOSTLY_Q8_KV,
70 };
71}
72
73/// Parameters controlling a [`quantize`] run.
74///
75/// Starts from `llama_model_quantize_default_params()`; the builder methods
76/// override the handful of fields common quantization runs need. ik's
77/// per-tensor override fields (`output_tensor_type`, `token_embedding_type`,
78/// `attn_q_type`, `attn_k_type`, `attn_v_type`, `attn_qkv_type`,
79/// `attn_output_type`, `ffn_gate_type`, `ffn_down_type`, `ffn_up_type`,
80/// `ffn_gate_inp_type`, `extra_output_type`, `per_layer_token_embedding_type`)
81/// and the extra flags (`only_copy`, `pure_`, `keep_split`,
82/// `ignore_imatrix_rules`, `only_repack`, `dry_run`, `partial_requant`) plus the
83/// opaque pointers (`imatrix`, `kv_overrides`, `custom_quants`,
84/// `repack_pattern`, `user_data`) keep their default values unless set through
85/// [`as_raw_mut`](LlamaQuantizeParams::as_raw_mut).
86#[derive(Debug, Clone)]
87pub struct LlamaQuantizeParams {
88 pub(crate) raw: sys::llama_model_quantize_params,
89}
90
91impl Default for LlamaQuantizeParams {
92 fn default() -> Self {
93 // SAFETY: returns a fully-initialized POD struct by value; pointer
94 // fields are set to null by the C side.
95 Self {
96 raw: unsafe { sys::llama_model_quantize_default_params() },
97 }
98 }
99}
100
101impl LlamaQuantizeParams {
102 /// Target quantization format (e.g. [`ftype::LLAMA_FTYPE_MOSTLY_Q4_K_M`] or
103 /// an ik SOTA quant like [`ftype::LLAMA_FTYPE_MOSTLY_IQ4_K`]).
104 #[must_use]
105 pub fn with_ftype(mut self, ftype: sys::llama_ftype) -> Self {
106 self.raw.ftype = ftype;
107 self
108 }
109
110 /// Number of threads to use (0 lets the C side pick a default).
111 #[must_use]
112 pub fn with_n_threads(mut self, n_threads: i32) -> Self {
113 self.raw.nthread = n_threads;
114 self
115 }
116
117 /// Allow requantizing tensors that are already quantized (default false).
118 #[must_use]
119 pub fn with_allow_requantize(mut self, allow_requantize: bool) -> Self {
120 self.raw.allow_requantize = allow_requantize;
121 self
122 }
123
124 /// Also quantize the `output.weight` tensor (default false keeps it larger
125 /// for quality).
126 #[must_use]
127 pub fn with_quantize_output_tensor(mut self, quantize_output_tensor: bool) -> Self {
128 self.raw.quantize_output_tensor = quantize_output_tensor;
129 self
130 }
131
132 /// Immutable access to the raw params (advanced/escape hatch).
133 #[must_use]
134 pub fn as_raw(&self) -> &sys::llama_model_quantize_params {
135 &self.raw
136 }
137
138 /// Mutable access to the raw params, for setting ik's extended per-tensor
139 /// override fields, flags, or the opaque `imatrix`/`kv_overrides` pointers
140 /// that are not surfaced by dedicated builders.
141 #[must_use]
142 pub fn as_raw_mut(&mut self) -> &mut sys::llama_model_quantize_params {
143 &mut self.raw
144 }
145}
146
147/// Quantize the GGUF model at `input`, writing the result to `output`.
148///
149/// Wraps `llama_model_quantize(fname_inp, fname_out, *const params)`; a non-zero
150/// return code becomes [`QuantizeError::Quantize`].
151///
152/// # Errors
153/// - [`QuantizeError::InvalidPath`] if a path is not valid UTF-8.
154/// - [`QuantizeError::Nul`] if a path contains an interior NUL byte.
155/// - [`QuantizeError::Quantize`] if the underlying C call reports failure.
156pub fn quantize(
157 input: &Path,
158 output: &Path,
159 params: &LlamaQuantizeParams,
160) -> Result<(), QuantizeError> {
161 let inp = CString::new(input.to_str().ok_or(QuantizeError::InvalidPath)?)?;
162 let out = CString::new(output.to_str().ok_or(QuantizeError::InvalidPath)?)?;
163
164 // SAFETY: both C strings outlive the call; `¶ms.raw` is a valid,
165 // fully-initialized `*const llama_model_quantize_params`.
166 let ret = unsafe { sys::llama_model_quantize(inp.as_ptr(), out.as_ptr(), ¶ms.raw) };
167
168 if ret == 0 {
169 Ok(())
170 } else {
171 Err(QuantizeError::Quantize(ret))
172 }
173}
174
175#[cfg(test)]
176mod tests {
177 use super::*;
178
179 // No-model unit test: params construct and builders round-trip. Actually
180 // quantizing needs a higher-precision source GGUF we do not have in
181 // fixtures (the test model is already IQ1_S), so the write-path itself is
182 // only exercised by build/link here.
183 #[test]
184 fn params_default_and_builders_round_trip() {
185 let default = LlamaQuantizeParams::default();
186 // Default builds and exposes its raw view.
187 let _ = default.as_raw();
188
189 let params = LlamaQuantizeParams::default()
190 .with_ftype(ftype::LLAMA_FTYPE_MOSTLY_Q4_K_M)
191 .with_n_threads(4)
192 .with_allow_requantize(true)
193 .with_quantize_output_tensor(true);
194
195 assert_eq!(params.as_raw().nthread, 4);
196 assert_eq!(params.as_raw().ftype, ftype::LLAMA_FTYPE_MOSTLY_Q4_K_M);
197 assert!(params.as_raw().allow_requantize);
198 assert!(params.as_raw().quantize_output_tensor);
199 }
200
201 #[test]
202 fn ik_sota_ftype_is_exposed() {
203 // ik SOTA quant const is reachable through the re-export module.
204 let params = LlamaQuantizeParams::default().with_ftype(ftype::LLAMA_FTYPE_MOSTLY_IQ4_K);
205 assert_eq!(params.as_raw().ftype, ftype::LLAMA_FTYPE_MOSTLY_IQ4_K);
206 }
207}