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)]
45#[cfg_attr(feature = "ts-bindings", derive(ts_rs::TS), ts(export))]
46pub enum KvCacheType {
47    F32,
48    F16,
49    Bf16,
50    Q8_0,
51    Q5_1,
52    Q5_0,
53    Q4_1,
54    Q4_0,
55}
56
57impl KvCacheType {
58    /// Every accepted type, in the order they should be offered to a user.
59    ///
60    /// The one list: [`FromStr`] parses against it, its error message names
61    /// it, and the CLI builds `--cache-type-k`/`-v`'s possible values from
62    /// it. Adding a variant here is enough to make it parse, appear in
63    /// `--help` and appear in shell completions.
64    pub const ALL: &'static [Self] = &[
65        Self::F32,
66        Self::F16,
67        Self::Bf16,
68        Self::Q8_0,
69        Self::Q5_1,
70        Self::Q5_0,
71        Self::Q4_1,
72        Self::Q4_0,
73    ];
74
75    /// The value passed on the command line (matches llama.cpp's own
76    /// `ggml_type_name`).
77    #[must_use]
78    pub const fn as_llama_arg(self) -> &'static str {
79        match self {
80            Self::F32 => "f32",
81            Self::F16 => "f16",
82            Self::Bf16 => "bf16",
83            Self::Q8_0 => "q8_0",
84            Self::Q5_1 => "q5_1",
85            Self::Q5_0 => "q5_0",
86            Self::Q4_1 => "q4_1",
87            Self::Q4_0 => "q4_0",
88        }
89    }
90
91    /// `(block_bytes, block_elems)` — the on-disk/in-memory layout ggml uses
92    /// for this type. Unquantized types are a trivial one-element block;
93    /// quantized types pack `block_elems` values into `block_bytes` bytes
94    /// (a shared scale/min plus packed sub-byte values).
95    #[must_use]
96    pub const fn block_layout(self) -> (u64, u64) {
97        match self {
98            Self::F32 => (4, 1),
99            Self::F16 | Self::Bf16 => (2, 1),
100            Self::Q8_0 => (34, 32),
101            Self::Q5_1 => (24, 32),
102            Self::Q5_0 => (22, 32),
103            Self::Q4_1 => (20, 32),
104            Self::Q4_0 => (18, 32),
105        }
106    }
107
108    /// Estimated bytes to store `elems` elements at this type, rounding up
109    /// to whole blocks (a partial trailing block still costs a full block).
110    #[must_use]
111    pub const fn bytes_for_elems(self, elems: u64) -> u64 {
112        let (block_bytes, block_elems) = self.block_layout();
113        elems.div_ceil(block_elems) * block_bytes
114    }
115}
116
117impl FromStr for KvCacheType {
118    type Err = String;
119
120    fn from_str(s: &str) -> Result<Self, Self::Err> {
121        let normalized = s.trim().to_ascii_lowercase();
122        Self::ALL
123            .iter()
124            .find(|t| t.as_llama_arg() == normalized)
125            .copied()
126            .ok_or_else(|| {
127                let expected: Vec<&str> = Self::ALL.iter().map(|t| t.as_llama_arg()).collect();
128                format!(
129                    "unknown KV cache type {normalized:?} (expected one of: {})",
130                    expected.join(", ")
131                )
132            })
133    }
134}
135
136/// Default K cache type: quantized to roughly halve the KV footprint versus
137/// `f16`, doubling how much conversation history the RAM/disk prompt caches
138/// can hold.
139pub const DEFAULT_CACHE_TYPE_K: KvCacheType = KvCacheType::Q8_0;
140
141/// Default V cache type. Same rationale as [`DEFAULT_CACHE_TYPE_K`]; see the
142/// module docs on the Flash Attention requirement for quantized V.
143pub const DEFAULT_CACHE_TYPE_V: KvCacheType = KvCacheType::Q8_0;
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn as_llama_arg_matches_ggml_type_names() {
151        assert_eq!(KvCacheType::F16.as_llama_arg(), "f16");
152        assert_eq!(KvCacheType::Q8_0.as_llama_arg(), "q8_0");
153    }
154
155    #[test]
156    fn from_str_round_trips_through_as_llama_arg() {
157        for t in [
158            KvCacheType::F32,
159            KvCacheType::F16,
160            KvCacheType::Bf16,
161            KvCacheType::Q8_0,
162            KvCacheType::Q5_1,
163            KvCacheType::Q5_0,
164            KvCacheType::Q4_1,
165            KvCacheType::Q4_0,
166        ] {
167            assert_eq!(KvCacheType::from_str(t.as_llama_arg()), Ok(t));
168        }
169    }
170
171    #[test]
172    fn from_str_is_case_insensitive_and_trims() {
173        assert_eq!(KvCacheType::from_str(" Q8_0 "), Ok(KvCacheType::Q8_0));
174        assert_eq!(KvCacheType::from_str("F16"), Ok(KvCacheType::F16));
175    }
176
177    #[test]
178    fn from_str_rejects_unknown_type() {
179        assert!(KvCacheType::from_str("q2_k").is_err());
180    }
181
182    #[test]
183    fn f16_bytes_for_elems_is_two_bytes_per_element() {
184        assert_eq!(KvCacheType::F16.bytes_for_elems(1000), 2000);
185    }
186
187    #[test]
188    fn q8_0_bytes_for_elems_matches_ggml_block_layout() {
189        // Exactly one block: 32 elements -> 34 bytes.
190        assert_eq!(KvCacheType::Q8_0.bytes_for_elems(32), 34);
191        // Two full blocks: 64 elements -> 68 bytes.
192        assert_eq!(KvCacheType::Q8_0.bytes_for_elems(64), 68);
193    }
194
195    #[test]
196    fn bytes_for_elems_rounds_up_a_partial_trailing_block() {
197        // 33 elements needs 2 blocks of 32 at q8_0 -> 68 bytes, not 34+.
198        assert_eq!(KvCacheType::Q8_0.bytes_for_elems(33), 68);
199    }
200
201    #[test]
202    fn q8_0_is_smaller_than_f16_for_the_same_element_count() {
203        let elems = 65_536;
204        assert!(KvCacheType::Q8_0.bytes_for_elems(elems) < KvCacheType::F16.bytes_for_elems(elems));
205    }
206
207    #[test]
208    fn cache_ram_setting_default_is_auto() {
209        assert_eq!(CacheRamSetting::default(), CacheRamSetting::Auto);
210    }
211}
gglib_core/cache_config.rs

gglib_core/
cache_config.rs
