gglib_core/
cache_config.rs

1//! KV-cache configuration types: quantized cache types and the host-RAM
2//! prompt cache setting.
3//!
4//! Kept as a standalone, low-complexity module (not folded into
5//! `server_config`) so cache-related config resolution has one home.
6
7use std::str::FromStr;
8
9use serde::{Deserialize, Serialize};
10
11// =============================================================================
12// Host-RAM prompt cache budget setting (`--cache-ram`)
13// =============================================================================
14
15/// How to determine the host-RAM prompt cache budget (`--cache-ram`).
16///
17/// Deliberately a two-state enum rather than `Option<u64>`: benchmark
18/// launches (which must never gain a prompt cache — it would perturb
19/// throughput measurements and RAM footprint) pass `ExplicitMb(0)`, which
20/// unambiguously disables the cache rather than leaving it to an implicit
21/// "no value" state.
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
23pub enum CacheRamSetting {
24    /// Compute a budget from system RAM, model size, and the KV estimate.
25    /// The default variant — every launch surface auto-sizes unless it
26    /// opts out.
27    #[default]
28    Auto,
29    /// Use exactly this MiB value. `0` disables the cache.
30    ExplicitMb(u64),
31}
32
33// =============================================================================
34// KV cache quantization type (`--cache-type-k` / `--cache-type-v`)
35// =============================================================================
36
37/// A llama.cpp KV-cache element type, as accepted by `--cache-type-k` /
38/// `--cache-type-v`.
39///
40/// Quantized K types are supported unconditionally by llama.cpp; quantized V
41/// types additionally require Flash Attention to be active (llama.cpp
42/// hard-errors at startup otherwise — see `resolve_kv_cache_types` in
43/// `gglib-runtime` for the escape hatches).
44#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
45pub enum KvCacheType {
46    F32,
47    F16,
48    Bf16,
49    Q8_0,
50    Q5_1,
51    Q5_0,
52    Q4_1,
53    Q4_0,
54}
55
56impl KvCacheType {
57    /// Every accepted type, in the order they should be offered to a user.
58    ///
59    /// The one list: [`FromStr`] parses against it, its error message names
60    /// it, and the CLI builds `--cache-type-k`/`-v`'s possible values from
61    /// it. Adding a variant here is enough to make it parse, appear in
62    /// `--help` and appear in shell completions.
63    pub const ALL: &'static [Self] = &[
64        Self::F32,
65        Self::F16,
66        Self::Bf16,
67        Self::Q8_0,
68        Self::Q5_1,
69        Self::Q5_0,
70        Self::Q4_1,
71        Self::Q4_0,
72    ];
73
74    /// The value passed on the command line (matches llama.cpp's own
75    /// `ggml_type_name`).
76    #[must_use]
77    pub const fn as_llama_arg(self) -> &'static str {
78        match self {
79            Self::F32 => "f32",
80            Self::F16 => "f16",
81            Self::Bf16 => "bf16",
82            Self::Q8_0 => "q8_0",
83            Self::Q5_1 => "q5_1",
84            Self::Q5_0 => "q5_0",
85            Self::Q4_1 => "q4_1",
86            Self::Q4_0 => "q4_0",
87        }
88    }
89
90    /// `(block_bytes, block_elems)` — the on-disk/in-memory layout ggml uses
91    /// for this type. Unquantized types are a trivial one-element block;
92    /// quantized types pack `block_elems` values into `block_bytes` bytes
93    /// (a shared scale/min plus packed sub-byte values).
94    #[must_use]
95    pub const fn block_layout(self) -> (u64, u64) {
96        match self {
97            Self::F32 => (4, 1),
98            Self::F16 | Self::Bf16 => (2, 1),
99            Self::Q8_0 => (34, 32),
100            Self::Q5_1 => (24, 32),
101            Self::Q5_0 => (22, 32),
102            Self::Q4_1 => (20, 32),
103            Self::Q4_0 => (18, 32),
104        }
105    }
106
107    /// Estimated bytes to store `elems` elements at this type, rounding up
108    /// to whole blocks (a partial trailing block still costs a full block).
109    #[must_use]
110    pub const fn bytes_for_elems(self, elems: u64) -> u64 {
111        let (block_bytes, block_elems) = self.block_layout();
112        elems.div_ceil(block_elems) * block_bytes
113    }
114}
115
116impl FromStr for KvCacheType {
117    type Err = String;
118
119    fn from_str(s: &str) -> Result<Self, Self::Err> {
120        let normalized = s.trim().to_ascii_lowercase();
121        Self::ALL
122            .iter()
123            .find(|t| t.as_llama_arg() == normalized)
124            .copied()
125            .ok_or_else(|| {
126                let expected: Vec<&str> = Self::ALL.iter().map(|t| t.as_llama_arg()).collect();
127                format!(
128                    "unknown KV cache type {normalized:?} (expected one of: {})",
129                    expected.join(", ")
130                )
131            })
132    }
133}
134
135/// Default K cache type: quantized to roughly halve the KV footprint versus
136/// `f16`, doubling how much conversation history the RAM/disk prompt caches
137/// can hold.
138pub const DEFAULT_CACHE_TYPE_K: KvCacheType = KvCacheType::Q8_0;
139
140/// Default V cache type. Same rationale as [`DEFAULT_CACHE_TYPE_K`]; see the
141/// module docs on the Flash Attention requirement for quantized V.
142pub const DEFAULT_CACHE_TYPE_V: KvCacheType = KvCacheType::Q8_0;
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn as_llama_arg_matches_ggml_type_names() {
150        assert_eq!(KvCacheType::F16.as_llama_arg(), "f16");
151        assert_eq!(KvCacheType::Q8_0.as_llama_arg(), "q8_0");
152    }
153
154    #[test]
155    fn from_str_round_trips_through_as_llama_arg() {
156        for t in [
157            KvCacheType::F32,
158            KvCacheType::F16,
159            KvCacheType::Bf16,
160            KvCacheType::Q8_0,
161            KvCacheType::Q5_1,
162            KvCacheType::Q5_0,
163            KvCacheType::Q4_1,
164            KvCacheType::Q4_0,
165        ] {
166            assert_eq!(KvCacheType::from_str(t.as_llama_arg()), Ok(t));
167        }
168    }
169
170    #[test]
171    fn from_str_is_case_insensitive_and_trims() {
172        assert_eq!(KvCacheType::from_str(" Q8_0 "), Ok(KvCacheType::Q8_0));
173        assert_eq!(KvCacheType::from_str("F16"), Ok(KvCacheType::F16));
174    }
175
176    #[test]
177    fn from_str_rejects_unknown_type() {
178        assert!(KvCacheType::from_str("q2_k").is_err());
179    }
180
181    #[test]
182    fn f16_bytes_for_elems_is_two_bytes_per_element() {
183        assert_eq!(KvCacheType::F16.bytes_for_elems(1000), 2000);
184    }
185
186    #[test]
187    fn q8_0_bytes_for_elems_matches_ggml_block_layout() {
188        // Exactly one block: 32 elements -> 34 bytes.
189        assert_eq!(KvCacheType::Q8_0.bytes_for_elems(32), 34);
190        // Two full blocks: 64 elements -> 68 bytes.
191        assert_eq!(KvCacheType::Q8_0.bytes_for_elems(64), 68);
192    }
193
194    #[test]
195    fn bytes_for_elems_rounds_up_a_partial_trailing_block() {
196        // 33 elements needs 2 blocks of 32 at q8_0 -> 68 bytes, not 34+.
197        assert_eq!(KvCacheType::Q8_0.bytes_for_elems(33), 68);
198    }
199
200    #[test]
201    fn q8_0_is_smaller_than_f16_for_the_same_element_count() {
202        let elems = 65_536;
203        assert!(KvCacheType::Q8_0.bytes_for_elems(elems) < KvCacheType::F16.bytes_for_elems(elems));
204    }
205
206    #[test]
207    fn cache_ram_setting_default_is_auto() {
208        assert_eq!(CacheRamSetting::default(), CacheRamSetting::Auto);
209    }
210}
gglib_core/cache_config.rs

gglib_core/
cache_config.rs
