rigspark_core/
plan.rs

1use crate::{
2    advice::{Throughput, throughput},
3    catalog::{CatalogModel, PerfDataset},
4    sizing::{
5        GpuVendor, HEADROOM, Hardware, KvCacheType, SizingRequest, ValidationError, evaluate,
6    },
7};
8use serde::Serialize;
9
10#[derive(Debug, Serialize)]
11#[serde(rename_all = "camelCase")]
12pub struct PathPlan {
13    pub path: &'static str,
14    pub fits: bool,
15    #[serde(skip_serializing_if = "Option::is_none")]
16    pub quant: Option<String>,
17    pub required_bytes: Option<f64>,
18    pub usable_bytes: f64,
19    /// Extra usable memory needed before the smallest quantization fits.
20    pub shortfall_bytes: Option<f64>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub reason: Option<&'static str>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub gpu_count: Option<usize>,
25    pub throughput: Throughput,
26    pub throughput_basis: &'static str,
27}
28
29#[derive(Debug, Serialize)]
30#[serde(rename_all = "camelCase")]
31pub struct Plan {
32    pub model: String,
33    pub context: Option<f64>,
34    pub backend: String,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub kv_cache_type: Option<&'static str>,
37    pub unified_memory: bool,
38    /// False when the requested context's KV cache cannot be sized from sourced geometry.
39    pub kv_cache_known: bool,
40    pub paths: Vec<PathPlan>,
41    pub recommended: Option<&'static str>,
42}
43
44fn discrete(vendor: &GpuVendor) -> bool {
45    matches!(
46        vendor,
47        GpuVendor::Nvidia | GpuVendor::Amd | GpuVendor::Intel
48    )
49}
50
51struct Inputs<'a> {
52    model: &'a CatalogModel,
53    dataset: &'a PerfDataset,
54    context: Option<f64>,
55    backend: &'a str,
56}
57
58fn plan_path(
59    inputs: &Inputs,
60    path: &'static str,
61    hardware: Hardware,
62    modeled: bool,
63    gpu_count: Option<usize>,
64) -> Result<PathPlan, ValidationError> {
65    let Inputs {
66        model,
67        dataset,
68        context,
69        backend,
70    } = *inputs;
71    let sized = evaluate(&SizingRequest {
72        model: model.sizing(),
73        hardware: hardware.clone(),
74        context,
75    })?;
76    let fit = sized.fit;
77    let shortfall = match (fit.fits, fit.reason, fit.required_bytes) {
78        (false, Some("vram-bound" | "ram-bound"), Some(required)) => Some(
79            (required / (1.0 - HEADROOM) - fit.usable_bytes)
80                .max(0.0)
81                .ceil(),
82        ),
83        _ => None,
84    };
85    let (estimate, basis) = match (&fit.quant, modeled) {
86        (Some(quant), true) => {
87            let estimate = throughput(model, quant, &hardware, dataset, backend)?;
88            let basis = if estimate.known {
89                "estimated"
90            } else {
91                "unknown: no sourced throughput class for this hardware and backend"
92            };
93            (estimate, basis)
94        }
95        (Some(_), false) => (
96            Throughput::unknown(),
97            "unknown: split across devices is not modeled; measure with llmup bench",
98        ),
99        (None, _) => (Throughput::unknown(), "unknown: does not fit"),
100    };
101    Ok(PathPlan {
102        path,
103        fits: fit.fits,
104        quant: fit.quant.map(|quant| quant.name),
105        required_bytes: fit.required_bytes,
106        usable_bytes: fit.usable_bytes,
107        shortfall_bytes: shortfall,
108        reason: fit.reason,
109        gpu_count,
110        throughput: estimate,
111        throughput_basis: basis,
112    })
113}
114
115/// Evaluates every execution path this hardware offers for one model, without network access.
116pub fn plan(
117    model: &CatalogModel,
118    hardware: &Hardware,
119    dataset: &PerfDataset,
120    context: Option<f64>,
121    backend: &str,
122) -> Result<Plan, ValidationError> {
123    plan_with_cache(model, hardware, dataset, context, backend, None)
124}
125
126/// Like [`plan`], sizing the KV cache at `kv_cache` precision instead of f16.
127pub fn plan_with_cache(
128    model: &CatalogModel,
129    hardware: &Hardware,
130    dataset: &PerfDataset,
131    context: Option<f64>,
132    backend: &str,
133    kv_cache: Option<KvCacheType>,
134) -> Result<Plan, ValidationError> {
135    if kv_cache.is_some() && context.is_none() {
136        return Err(ValidationError("a KV cache type requires a context".into()));
137    }
138    let resized = kv_cache.map(|kind| model.with_kv_cache(kind)).transpose()?;
139    let model = resized.as_ref().unwrap_or(model);
140    let inputs = Inputs {
141        model,
142        dataset,
143        context,
144        backend,
145    };
146    let unified = hardware.is_unified();
147    let mut paths = Vec::new();
148    let gpus: Vec<_> = hardware
149        .gpu
150        .iter()
151        .filter(|gpu| discrete(&gpu.vendor) && gpu.vram_bytes > 0.0)
152        .cloned()
153        .collect();
154    if unified {
155        paths.push(plan_path(&inputs, "unified", hardware.clone(), true, None)?);
156    } else if let Some(largest) = gpus
157        .iter()
158        .max_by(|left, right| left.vram_bytes.total_cmp(&right.vram_bytes))
159    {
160        let mut single = hardware.clone();
161        single.gpu = vec![largest.clone()];
162        paths.push(plan_path(&inputs, "gpu", single, true, Some(1))?);
163        let peers: Vec<_> = gpus
164            .iter()
165            .filter(|gpu| gpu.vendor == largest.vendor)
166            .collect();
167        if peers.len() > 1 {
168            let mut pooled = hardware.clone();
169            let mut device = largest.clone();
170            device.vram_bytes = peers.iter().map(|gpu| gpu.vram_bytes).sum();
171            pooled.gpu = vec![device];
172            paths.push(plan_path(
173                &inputs,
174                "multi-gpu",
175                pooled,
176                false,
177                Some(peers.len()),
178            )?);
179        }
180        let mut offload = hardware.clone();
181        offload.gpu.clear();
182        offload.unified_memory = Some(false);
183        offload.free_ram_bytes += gpus.iter().map(|gpu| gpu.vram_bytes).sum::<f64>();
184        paths.push(plan_path(
185            &inputs,
186            "cpu-offload",
187            offload,
188            false,
189            Some(gpus.len()),
190        )?);
191    }
192    let mut cpu = hardware.clone();
193    cpu.gpu.clear();
194    cpu.unified_memory = Some(false);
195    paths.push(plan_path(&inputs, "cpu", cpu, true, None)?);
196    let recommended = paths.iter().find(|path| path.fits).map(|path| path.path);
197    Ok(Plan {
198        model: model.id.clone(),
199        context,
200        backend: backend.into(),
201        kv_cache_type: kv_cache.map(|kind| KvCacheType::label(Some(kind))),
202        unified_memory: unified,
203        kv_cache_known: context.is_none() || model.kv_bytes_per_token.is_some(),
204        paths,
205        recommended,
206    })
207}
208
209fn gib(bytes: f64) -> String {
210    format!("{:.1} GiB", bytes / 1073741824.0)
211}
212
213pub fn format_plan(plan: &Plan) -> String {
214    let context = plan.context.map_or_else(
215        || "default context".to_owned(),
216        |tokens| format!("{tokens} tokens"),
217    );
218    let context = match plan.kv_cache_type {
219        Some(kind) => format!("{context}, KV {kind}"),
220        None => context,
221    };
222    let mut text = format!(
223        "Plan for {} ({context}, {} backend{})\n\n",
224        plan.model,
225        plan.backend,
226        if plan.unified_memory {
227            ", unified memory"
228        } else {
229            ""
230        }
231    );
232    if !plan.kv_cache_known {
233        text.push_str(
234            "KV cache at this context is unknown (no sourced attention geometry); sizes cover weights only.\n\n",
235        );
236    }
237    for path in &plan.paths {
238        let label = match path.gpu_count {
239            Some(1) => format!("{} (1 GPU)", path.path),
240            Some(count) if path.path != "cpu-offload" => format!("{} ({count} GPUs)", path.path),
241            _ => path.path.to_owned(),
242        };
243        let fit = if path.fits {
244            format!(
245                "fits {} using {} of {}",
246                path.quant.as_deref().unwrap_or("?"),
247                path.required_bytes
248                    .map_or_else(|| "unknown".to_owned(), gib),
249                gib(path.usable_bytes)
250            )
251        } else if let Some(shortfall) = path.shortfall_bytes {
252            format!("does not fit; needs {} more usable memory", gib(shortfall))
253        } else {
254            format!("does not fit ({})", path.reason.unwrap_or("unknown"))
255        };
256        let speed = if path.throughput.known {
257            format!(
258                "{}-{} tok/s estimated",
259                path.throughput.low_tok_per_sec, path.throughput.high_tok_per_sec
260            )
261        } else if path.fits && path.throughput_basis.contains("not modeled") {
262            "speed unknown (not modeled; measure with llmup bench)".to_owned()
263        } else if path.fits {
264            "speed unknown (no sourced throughput data)".to_owned()
265        } else {
266            String::new()
267        };
268        text.push_str(&format!("  {label:<22} {fit}"));
269        if !speed.is_empty() {
270            text.push_str(&format!("; {speed}"));
271        }
272        text.push('\n');
273    }
274    text.push_str(&match plan.recommended {
275        Some(path) => format!("\nRecommended path: {path}\n"),
276        None => "\nNo path fits this model at this context.\n".to_owned(),
277    });
278    text
279}
rigspark_core/plan.rs

rigspark_core/
plan.rs
