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

rigspark_core/
sizing.rs
