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        // OCP Microscaling: 4-bit elements plus one shared 8-bit scale per 32-element block.
192        ("mxfp4", 4.25),
193        ("iq1", 1.9),
194        ("iq2", 2.4),
195        ("iq3", 3.4),
196        ("iq4", 4.3),
197        ("q2", 2.8),
198        ("q3", 3.5),
199        ("q4", 4.7),
200        ("q5", 5.6),
201        ("q6", 6.6),
202        ("q8", 8.5),
203        ("f16", 16.5),
204        ("fp16", 16.5),
205        ("bf16", 16.5),
206        ("f32", 32.5),
207        ("fp32", 32.5),
208    ]
209    .into_iter()
210    .find_map(|(prefix, bits)| name.starts_with(prefix).then_some(bits))
211}
212
213pub fn kv_cache_bytes(per_token: u64, tokens: u64) -> Result<u64, ValidationError> {
214    if per_token == 0 {
215        return Err(invalid("KV bytes per token must be positive"));
216    }
217    let bytes = per_token
218        .checked_mul(tokens)
219        .ok_or_else(|| invalid("KV size overflow"))?;
220    if bytes > SAFE_INTEGER as u64 {
221        return Err(invalid("KV size exceeds safe-integer range"));
222    }
223    Ok(bytes)
224}
225
226/// KV cache element types shared by llama.cpp and Ollama.
227#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Hash, Deserialize, Serialize)]
228pub enum KvCacheType {
229    #[default]
230    #[serde(rename = "f16")]
231    F16,
232    #[serde(rename = "q8_0")]
233    Q8_0,
234    #[serde(rename = "q4_0")]
235    Q4_0,
236}
237
238impl KvCacheType {
239    pub fn name(self) -> &'static str {
240        match self {
241            Self::F16 => "f16",
242            Self::Q8_0 => "q8_0",
243            Self::Q4_0 => "q4_0",
244        }
245    }
246    pub fn parse(name: &str) -> Option<Self> {
247        [Self::F16, Self::Q8_0, Self::Q4_0]
248            .into_iter()
249            .find(|kind| kind.name() == name)
250    }
251    /// Label for advice output; the implicit default keeps its historical `fp16` spelling.
252    pub fn label(kind: Option<Self>) -> &'static str {
253        match kind {
254            None | Some(Self::F16) => "fp16",
255            Some(kind) => kind.name(),
256        }
257    }
258    /// (bytes per block, elements per block) from ggml's block_q8_0 and block_q4_0 layouts.
259    fn block(self) -> (u64, u64) {
260        match self {
261            Self::F16 => (2, 1),
262            Self::Q8_0 => (34, 32),
263            Self::Q4_0 => (18, 32),
264        }
265    }
266    fn bytes_for(self, elements: u64) -> Option<u64> {
267        let (bytes, per_block) = self.block();
268        elements.div_ceil(per_block).checked_mul(bytes)
269    }
270}
271
272/// Re-sizes a catalog f16 KV rate (K and V halves) for the requested cache types.
273pub fn typed_kv_bytes_per_token(
274    f16_per_token: u64,
275    k: KvCacheType,
276    v: KvCacheType,
277) -> Result<u64, ValidationError> {
278    if f16_per_token == 0 || !f16_per_token.is_multiple_of(4) {
279        return Err(invalid(
280            "KV bytes per token must split into f16 K and V halves",
281        ));
282    }
283    let elements = f16_per_token / 4;
284    k.bytes_for(elements)
285        .zip(v.bytes_for(elements))
286        .and_then(|(k, v)| k.checked_add(v))
287        .filter(|bytes| *bytes <= SAFE_INTEGER as u64)
288        .ok_or_else(|| invalid("KV size exceeds safe-integer range"))
289}
290
291fn validate(request: &SizingRequest) -> Result<(), ValidationError> {
292    let model = &request.model;
293    if model.id.is_empty() || model.id.len() > 256 {
294        return Err(invalid("invalid model id length"));
295    }
296    parse_param_count(&model.params)?;
297    safe_number(model.context_length, 1.0, "contextLength")?;
298    if let Some(context) = request.context {
299        safe_number(context, 1.0, "context")?;
300        if context > 10_000_000.0 {
301            return Err(invalid("context exceeds ceiling"));
302        }
303    }
304    if let Some(rate) = model.kv_bytes_per_token {
305        safe_number(rate, 1.0, "kvBytesPerToken")?;
306    }
307    for (name, value) in [
308        ("totalRamBytes", request.hardware.total_ram_bytes),
309        ("freeRamBytes", request.hardware.free_ram_bytes),
310        ("freeDiskBytes", request.hardware.free_disk_bytes),
311    ] {
312        safe_number(value, 0.0, name)?;
313    }
314    if request.hardware.gpu.len() > 128 || model.quantizations.len() > 1024 {
315        return Err(invalid("input collection too large"));
316    }
317    for gpu in &request.hardware.gpu {
318        safe_number(gpu.vram_bytes, 0.0, "vramBytes")?;
319    }
320    for quant in &model.quantizations {
321        if quant.name.is_empty() || quant.name.len() > 100 {
322            return Err(invalid("invalid quantization name"));
323        }
324        safe_number(quant.disk_bytes, 1.0, "diskBytes")?;
325        safe_number(quant.min_ram_bytes, 0.0, "minRamBytes")?;
326        safe_number(quant.min_vram_bytes, 0.0, "minVramBytes")?;
327        if quant
328            .sha256
329            .as_ref()
330            .is_some_and(|sha| sha.len() != 64 || !sha.bytes().all(|byte| byte.is_ascii_hexdigit()))
331        {
332            return Err(invalid("invalid sha256"));
333        }
334    }
335    Ok(())
336}
337
338pub fn memory_capacity(hardware: &Hardware) -> (&'static str, f64) {
339    let unified = hardware.is_unified();
340    let vram = hardware
341        .gpu
342        .iter()
343        .map(|gpu| gpu.vram_bytes)
344        .fold(0.0_f64, f64::max);
345    let memory_kind = if !unified && vram > 0.0 {
346        "vram"
347    } else {
348        "ram"
349    };
350    let usable_bytes = if memory_kind == "vram" {
351        vram
352    } else {
353        ((if unified {
354            hardware.total_ram_bytes
355        } else {
356            hardware.free_ram_bytes
357        }) - OS_RESERVE)
358            .max(0.0)
359    };
360    (memory_kind, usable_bytes)
361}
362pub fn evaluate(request: &SizingRequest) -> Result<SizingResult, ValidationError> {
363    validate(request)?;
364    let model = &request.model;
365    let hardware = &request.hardware;
366    let (memory_kind, usable_bytes) = memory_capacity(hardware);
367    let budget = usable_bytes * (1.0 - HEADROOM);
368    let params = parse_param_count(&model.params)?;
369    let mut weights = Vec::new();
370    let mut required = Vec::new();
371    let mut at_context = Vec::new();
372    let mut max_context = Vec::new();
373    let mut best: Option<(usize, f64)> = None;
374    let mut smallest: Option<(f64, bool)> = None;
375    for (index, quant) in model.quantizations.iter().enumerate() {
376        let bits = quant_bits(&quant.name);
377        if bits.is_none() && matches!(model.architecture, Architecture::Moe) {
378            return Err(invalid("cannot size MoE with unknown quantization"));
379        }
380        let weight = quant
381            .disk_bytes
382            .max(bits.map_or(0.0, |bits| (params * bits / 8.0).ceil()));
383        let legacy = weight + (weight * 0.15).ceil();
384        safe_number(legacy, 1.0, "required memory")?;
385        let slack = (weight * 0.05).ceil();
386        let context_bytes = match (model.kv_bytes_per_token, request.context) {
387            (Some(rate), Some(context)) => {
388                let explicit = weight + kv_cache_bytes(rate as u64, context as u64)? as f64 + slack;
389                safe_number(explicit, 1.0, "context memory")?;
390                Some(legacy.max(explicit))
391            }
392            _ => None,
393        };
394        let maximum = model.kv_bytes_per_token.map(|rate| {
395            if legacy > budget {
396                0.0
397            } else {
398                ((budget - weight - slack) / rate).floor()
399            }
400        });
401        let footprint = context_bytes.unwrap_or(legacy);
402        let memory_fits = footprint <= budget;
403        if memory_fits && quant.disk_bytes <= hardware.free_disk_bytes {
404            let better = match best {
405                None => true,
406                Some((best_index, best_bytes)) => {
407                    footprint > best_bytes
408                        || (footprint == best_bytes
409                            && bits
410                                .zip(quant_bits(&model.quantizations[best_index].name))
411                                .is_some_and(|(current, prior)| current > prior))
412                }
413            };
414            if better {
415                best = Some((index, footprint));
416            }
417        }
418        if smallest.is_none_or(|(bytes, _)| footprint < bytes) {
419            smallest = Some((footprint, memory_fits));
420        }
421        weights.push(weight);
422        required.push(legacy);
423        at_context.push(context_bytes);
424        max_context.push(maximum);
425    }
426    let fit = if request
427        .context
428        .is_some_and(|context| context > model.context_length)
429    {
430        FitResult {
431            fits: false,
432            quant: None,
433            reason: Some("context-bound"),
434            usable_bytes,
435            required_bytes: None,
436        }
437    } else if let Some((index, bytes)) = best {
438        FitResult {
439            fits: true,
440            quant: Some(model.quantizations[index].clone()),
441            reason: None,
442            usable_bytes,
443            required_bytes: Some(bytes),
444        }
445    } else {
446        let reason = if smallest.is_some_and(|(_, fits)| fits) {
447            "disk-bound"
448        } else if memory_kind == "vram" {
449            "vram-bound"
450        } else {
451            "ram-bound"
452        };
453        FitResult {
454            fits: false,
455            quant: None,
456            reason: Some(reason),
457            usable_bytes,
458            required_bytes: smallest.map(|(bytes, _)| bytes),
459        }
460    };
461    Ok(SizingResult {
462        fit,
463        memory_kind,
464        usable_bytes,
465        weights,
466        required,
467        at_context,
468        max_context,
469    })
470}
rigspark_core/sizing.rs

rigspark_core/
sizing.rs
