rigspark_core/
sizing.rs

1use serde::{Deserialize, Serialize};
2use thiserror::Error;
3
4const SAFE_INTEGER: f64 = 9_007_199_254_740_991.0;
5const OS_RESERVE: f64 = 2.0 * 1024.0 * 1024.0 * 1024.0;
6pub const HEADROOM: f64 = 0.15;
7
8#[derive(Debug, Error)]
9#[error("{0}")]
10pub struct ValidationError(pub String);
11
12#[derive(Debug, Clone, Deserialize, Serialize)]
13#[serde(rename_all = "lowercase")]
14pub enum Architecture {
15    Dense,
16    Moe,
17}
18
19#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
20#[serde(rename_all = "lowercase")]
21pub enum CpuArch {
22    X64,
23    Arm64,
24}
25
26#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
27#[serde(rename_all = "lowercase")]
28pub enum Platform {
29    Darwin,
30    Linux,
31    Win32,
32}
33
34#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)]
35#[serde(rename_all = "lowercase")]
36pub enum GpuVendor {
37    Apple,
38    Nvidia,
39    Amd,
40    Intel,
41    None,
42}
43
44#[derive(Debug, Clone, Deserialize, Serialize)]
45#[serde(rename_all = "camelCase", deny_unknown_fields)]
46pub struct Gpu {
47    pub vendor: GpuVendor,
48    pub vram_bytes: f64,
49}
50
51#[derive(Debug, Clone, Deserialize, Serialize)]
52#[serde(rename_all = "camelCase", deny_unknown_fields)]
53pub struct Hardware {
54    pub arch: CpuArch,
55    pub platform: Platform,
56    pub total_ram_bytes: f64,
57    pub free_ram_bytes: f64,
58    pub free_disk_bytes: f64,
59    pub gpu: Vec<Gpu>,
60    /// Shared CPU/GPU memory pool (Apple Silicon, AMD APUs, NVIDIA Grace/Jetson SoCs).
61    #[serde(default, skip_serializing_if = "Option::is_none")]
62    pub unified_memory: Option<bool>,
63}
64
65impl Hardware {
66    /// Detected topology when known; otherwise only Apple Silicon Macs are assumed unified.
67    pub fn is_unified(&self) -> bool {
68        self.unified_memory
69            .unwrap_or(self.arch == CpuArch::Arm64 && self.platform == Platform::Darwin)
70    }
71}
72
73#[derive(Debug, Clone, Deserialize, Serialize)]
74#[serde(rename_all = "camelCase", deny_unknown_fields)]
75pub struct Quantization {
76    pub name: String,
77    pub disk_bytes: f64,
78    pub min_ram_bytes: f64,
79    pub min_vram_bytes: f64,
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pub sha256: Option<String>,
82    #[serde(skip_serializing_if = "Option::is_none")]
83    pub digest_verified: Option<bool>,
84}
85
86#[derive(Debug, Deserialize)]
87#[serde(rename_all = "camelCase", deny_unknown_fields)]
88pub struct Model {
89    pub id: String,
90    pub params: String,
91    pub architecture: Architecture,
92    pub context_length: f64,
93    pub kv_bytes_per_token: Option<f64>,
94    pub quantizations: Vec<Quantization>,
95}
96
97#[derive(Debug, Deserialize)]
98#[serde(rename_all = "camelCase", deny_unknown_fields)]
99pub struct SizingRequest {
100    pub model: Model,
101    pub hardware: Hardware,
102    pub context: Option<f64>,
103}
104
105#[derive(Debug, Serialize)]
106#[serde(rename_all = "camelCase")]
107pub struct FitResult {
108    pub fits: bool,
109    #[serde(skip_serializing_if = "Option::is_none")]
110    pub quant: Option<Quantization>,
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub reason: Option<&'static str>,
113    pub usable_bytes: f64,
114    pub required_bytes: Option<f64>,
115}
116
117#[derive(Debug, Serialize)]
118#[serde(rename_all = "camelCase")]
119pub struct SizingResult {
120    pub fit: FitResult,
121    pub memory_kind: &'static str,
122    pub usable_bytes: f64,
123    pub weights: Vec<f64>,
124    pub required: Vec<f64>,
125    pub at_context: Vec<Option<f64>>,
126    pub max_context: Vec<Option<f64>>,
127}
128
129fn invalid(message: &str) -> ValidationError {
130    ValidationError(message.into())
131}
132
133fn safe_number(value: f64, minimum: f64, label: &str) -> Result<(), ValidationError> {
134    if !value.is_finite() || value < minimum || value > SAFE_INTEGER || value.fract() != 0.0 {
135        return Err(invalid(&format!(
136            "{label} must be a safe integer >= {minimum}"
137        )));
138    }
139    Ok(())
140}
141
142pub fn parse_param_count(label: &str) -> Result<f64, ValidationError> {
143    let split = label
144        .len()
145        .checked_sub(1)
146        .ok_or_else(|| invalid("invalid parameter label"))?;
147    if !label.is_ascii() {
148        return Err(invalid("invalid parameter label"));
149    }
150    let (number, unit) = label.split_at(split);
151    let mut parts = number.split('.');
152    let integer = parts.next().unwrap_or("");
153    let fraction = parts.next();
154    if integer.is_empty()
155        || !integer.bytes().all(|byte| byte.is_ascii_digit())
156        || fraction.is_some_and(|value| {
157            value.is_empty() || !value.bytes().all(|byte| byte.is_ascii_digit())
158        })
159        || parts.next().is_some()
160    {
161        return Err(invalid("invalid parameter label"));
162    }
163    let multiplier = match unit {
164        "M" => 1e6,
165        "B" => 1e9,
166        "T" => 1e12,
167        _ => return Err(invalid("invalid parameter unit")),
168    };
169    let count = number
170        .parse::<f64>()
171        .map_err(|_| invalid("invalid parameter label"))?
172        * multiplier;
173    if !count.is_finite() || count <= 0.0 || count > SAFE_INTEGER {
174        return Err(invalid("parameter count outside safe range"));
175    }
176    Ok(count)
177}
178
179pub fn quant_bits(name: &str) -> Option<f64> {
180    let name = name.to_ascii_lowercase();
181    [
182        ("iq1", 1.9),
183        ("iq2", 2.4),
184        ("iq3", 3.4),
185        ("iq4", 4.3),
186        ("q2", 2.8),
187        ("q3", 3.5),
188        ("q4", 4.7),
189        ("q5", 5.6),
190        ("q6", 6.6),
191        ("q8", 8.5),
192        ("f16", 16.5),
193        ("fp16", 16.5),
194        ("bf16", 16.5),
195        ("f32", 32.5),
196        ("fp32", 32.5),
197    ]
198    .into_iter()
199    .find_map(|(prefix, bits)| name.starts_with(prefix).then_some(bits))
200}
201
202pub fn kv_cache_bytes(per_token: u64, tokens: u64) -> Result<u64, ValidationError> {
203    if per_token == 0 {
204        return Err(invalid("KV bytes per token must be positive"));
205    }
206    let bytes = per_token
207        .checked_mul(tokens)
208        .ok_or_else(|| invalid("KV size overflow"))?;
209    if bytes > SAFE_INTEGER as u64 {
210        return Err(invalid("KV size exceeds safe-integer range"));
211    }
212    Ok(bytes)
213}
214
215/// KV cache element types shared by llama.cpp and Ollama.
216#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Deserialize, Serialize)]
217pub enum KvCacheType {
218    #[default]
219    #[serde(rename = "f16")]
220    F16,
221    #[serde(rename = "q8_0")]
222    Q8_0,
223    #[serde(rename = "q4_0")]
224    Q4_0,
225}
226
227impl KvCacheType {
228    pub fn name(self) -> &'static str {
229        match self {
230            Self::F16 => "f16",
231            Self::Q8_0 => "q8_0",
232            Self::Q4_0 => "q4_0",
233        }
234    }
235    pub fn parse(name: &str) -> Option<Self> {
236        [Self::F16, Self::Q8_0, Self::Q4_0]
237            .into_iter()
238            .find(|kind| kind.name() == name)
239    }
240    /// Label for advice output; the implicit default keeps its historical `fp16` spelling.
241    pub fn label(kind: Option<Self>) -> &'static str {
242        match kind {
243            None | Some(Self::F16) => "fp16",
244            Some(kind) => kind.name(),
245        }
246    }
247    /// (bytes per block, elements per block) from ggml's block_q8_0 and block_q4_0 layouts.
248    fn block(self) -> (u64, u64) {
249        match self {
250            Self::F16 => (2, 1),
251            Self::Q8_0 => (34, 32),
252            Self::Q4_0 => (18, 32),
253        }
254    }
255    fn bytes_for(self, elements: u64) -> Option<u64> {
256        let (bytes, per_block) = self.block();
257        elements.div_ceil(per_block).checked_mul(bytes)
258    }
259}
260
261/// Re-sizes a catalog f16 KV rate (K and V halves) for the requested cache types.
262pub fn typed_kv_bytes_per_token(
263    f16_per_token: u64,
264    k: KvCacheType,
265    v: KvCacheType,
266) -> Result<u64, ValidationError> {
267    if f16_per_token == 0 || !f16_per_token.is_multiple_of(4) {
268        return Err(invalid(
269            "KV bytes per token must split into f16 K and V halves",
270        ));
271    }
272    let elements = f16_per_token / 4;
273    k.bytes_for(elements)
274        .zip(v.bytes_for(elements))
275        .and_then(|(k, v)| k.checked_add(v))
276        .filter(|bytes| *bytes <= SAFE_INTEGER as u64)
277        .ok_or_else(|| invalid("KV size exceeds safe-integer range"))
278}
279
280fn validate(request: &SizingRequest) -> Result<(), ValidationError> {
281    let model = &request.model;
282    if model.id.is_empty() || model.id.len() > 256 {
283        return Err(invalid("invalid model id length"));
284    }
285    parse_param_count(&model.params)?;
286    safe_number(model.context_length, 1.0, "contextLength")?;
287    if let Some(context) = request.context {
288        safe_number(context, 1.0, "context")?;
289        if context > 10_000_000.0 {
290            return Err(invalid("context exceeds ceiling"));
291        }
292    }
293    if let Some(rate) = model.kv_bytes_per_token {
294        safe_number(rate, 1.0, "kvBytesPerToken")?;
295    }
296    for (name, value) in [
297        ("totalRamBytes", request.hardware.total_ram_bytes),
298        ("freeRamBytes", request.hardware.free_ram_bytes),
299        ("freeDiskBytes", request.hardware.free_disk_bytes),
300    ] {
301        safe_number(value, 0.0, name)?;
302    }
303    if request.hardware.gpu.len() > 128 || model.quantizations.len() > 1024 {
304        return Err(invalid("input collection too large"));
305    }
306    for gpu in &request.hardware.gpu {
307        safe_number(gpu.vram_bytes, 0.0, "vramBytes")?;
308    }
309    for quant in &model.quantizations {
310        if quant.name.is_empty() || quant.name.len() > 100 {
311            return Err(invalid("invalid quantization name"));
312        }
313        safe_number(quant.disk_bytes, 1.0, "diskBytes")?;
314        safe_number(quant.min_ram_bytes, 0.0, "minRamBytes")?;
315        safe_number(quant.min_vram_bytes, 0.0, "minVramBytes")?;
316        if quant
317            .sha256
318            .as_ref()
319            .is_some_and(|sha| sha.len() != 64 || !sha.bytes().all(|byte| byte.is_ascii_hexdigit()))
320        {
321            return Err(invalid("invalid sha256"));
322        }
323    }
324    Ok(())
325}
326
327pub fn memory_capacity(hardware: &Hardware) -> (&'static str, f64) {
328    let unified = hardware.is_unified();
329    let vram = hardware
330        .gpu
331        .iter()
332        .map(|gpu| gpu.vram_bytes)
333        .fold(0.0_f64, f64::max);
334    let memory_kind = if !unified && vram > 0.0 {
335        "vram"
336    } else {
337        "ram"
338    };
339    let usable_bytes = if memory_kind == "vram" {
340        vram
341    } else {
342        ((if unified {
343            hardware.total_ram_bytes
344        } else {
345            hardware.free_ram_bytes
346        }) - OS_RESERVE)
347            .max(0.0)
348    };
349    (memory_kind, usable_bytes)
350}
351pub fn evaluate(request: &SizingRequest) -> Result<SizingResult, ValidationError> {
352    validate(request)?;
353    let model = &request.model;
354    let hardware = &request.hardware;
355    let (memory_kind, usable_bytes) = memory_capacity(hardware);
356    let budget = usable_bytes * (1.0 - HEADROOM);
357    let params = parse_param_count(&model.params)?;
358    let mut weights = Vec::new();
359    let mut required = Vec::new();
360    let mut at_context = Vec::new();
361    let mut max_context = Vec::new();
362    let mut best: Option<(usize, f64)> = None;
363    let mut smallest: Option<(f64, bool)> = None;
364    for (index, quant) in model.quantizations.iter().enumerate() {
365        let bits = quant_bits(&quant.name);
366        if bits.is_none() && matches!(model.architecture, Architecture::Moe) {
367            return Err(invalid("cannot size MoE with unknown quantization"));
368        }
369        let weight = quant
370            .disk_bytes
371            .max(bits.map_or(0.0, |bits| (params * bits / 8.0).ceil()));
372        let legacy = weight + (weight * 0.15).ceil();
373        safe_number(legacy, 1.0, "required memory")?;
374        let slack = (weight * 0.05).ceil();
375        let context_bytes = match (model.kv_bytes_per_token, request.context) {
376            (Some(rate), Some(context)) => {
377                let explicit = weight + kv_cache_bytes(rate as u64, context as u64)? as f64 + slack;
378                safe_number(explicit, 1.0, "context memory")?;
379                Some(legacy.max(explicit))
380            }
381            _ => None,
382        };
383        let maximum = model.kv_bytes_per_token.map(|rate| {
384            if legacy > budget {
385                0.0
386            } else {
387                ((budget - weight - slack) / rate).floor()
388            }
389        });
390        let footprint = context_bytes.unwrap_or(legacy);
391        let memory_fits = footprint <= budget;
392        if memory_fits && quant.disk_bytes <= hardware.free_disk_bytes {
393            let better = match best {
394                None => true,
395                Some((best_index, best_bytes)) => {
396                    footprint > best_bytes
397                        || (footprint == best_bytes
398                            && bits
399                                .zip(quant_bits(&model.quantizations[best_index].name))
400                                .is_some_and(|(current, prior)| current > prior))
401                }
402            };
403            if better {
404                best = Some((index, footprint));
405            }
406        }
407        if smallest.is_none_or(|(bytes, _)| footprint < bytes) {
408            smallest = Some((footprint, memory_fits));
409        }
410        weights.push(weight);
411        required.push(legacy);
412        at_context.push(context_bytes);
413        max_context.push(maximum);
414    }
415    let fit = if request
416        .context
417        .is_some_and(|context| context > model.context_length)
418    {
419        FitResult {
420            fits: false,
421            quant: None,
422            reason: Some("context-bound"),
423            usable_bytes,
424            required_bytes: None,
425        }
426    } else if let Some((index, bytes)) = best {
427        FitResult {
428            fits: true,
429            quant: Some(model.quantizations[index].clone()),
430            reason: None,
431            usable_bytes,
432            required_bytes: Some(bytes),
433        }
434    } else {
435        let reason = if smallest.is_some_and(|(_, fits)| fits) {
436            "disk-bound"
437        } else if memory_kind == "vram" {
438            "vram-bound"
439        } else {
440            "ram-bound"
441        };
442        FitResult {
443            fits: false,
444            quant: None,
445            reason: Some(reason),
446            usable_bytes,
447            required_bytes: smallest.map(|(bytes, _)| bytes),
448        }
449    };
450    Ok(SizingResult {
451        fit,
452        memory_kind,
453        usable_bytes,
454        weights,
455        required,
456        at_context,
457        max_context,
458    })
459}
rigspark_core/sizing.rs

rigspark_core/
sizing.rs
