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

rigspark_core/
sizing.rs
