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 cpu_only = model
148        .quantizations
149        .iter()
150        .any(crate::sizing::Quantization::is_bitnet);
151    let mut paths = Vec::new();
152    let gpus: Vec<_> = hardware
153        .gpu
154        .iter()
155        .filter(|gpu| !cpu_only && discrete(&gpu.vendor) && gpu.vram_bytes > 0.0)
156        .cloned()
157        .collect();
158    if unified && !cpu_only {
159        paths.push(plan_path(&inputs, "unified", hardware.clone(), true, None)?);
160    } else if let Some(largest) = gpus
161        .iter()
162        .max_by(|left, right| left.vram_bytes.total_cmp(&right.vram_bytes))
163    {
164        let mut single = hardware.clone();
165        single.gpu = vec![largest.clone()];
166        paths.push(plan_path(&inputs, "gpu", single, true, Some(1))?);
167        let peers: Vec<_> = gpus
168            .iter()
169            .filter(|gpu| gpu.vendor == largest.vendor)
170            .collect();
171        if peers.len() > 1 {
172            let mut pooled = hardware.clone();
173            let mut device = largest.clone();
174            device.vram_bytes = peers.iter().map(|gpu| gpu.vram_bytes).sum();
175            pooled.gpu = vec![device];
176            paths.push(plan_path(
177                &inputs,
178                "multi-gpu",
179                pooled,
180                false,
181                Some(peers.len()),
182            )?);
183        }
184        let mut offload = hardware.clone();
185        offload.gpu.clear();
186        offload.unified_memory = Some(false);
187        offload.free_ram_bytes += gpus.iter().map(|gpu| gpu.vram_bytes).sum::<f64>();
188        paths.push(plan_path(
189            &inputs,
190            "cpu-offload",
191            offload,
192            false,
193            Some(gpus.len()),
194        )?);
195    }
196    let mut cpu = hardware.clone();
197    cpu.gpu.clear();
198    cpu.unified_memory = Some(false);
199    paths.push(plan_path(&inputs, "cpu", cpu, true, None)?);
200    let recommended = paths.iter().find(|path| path.fits).map(|path| path.path);
201    Ok(Plan {
202        model: model.id.clone(),
203        context,
204        backend: backend.into(),
205        kv_cache_type: kv_cache.map(|kind| KvCacheType::label(Some(kind))),
206        unified_memory: unified,
207        kv_cache_known: context.is_none() || model.kv_bytes_per_token.is_some(),
208        paths,
209        recommended,
210    })
211}
212
213fn gib(bytes: f64) -> String {
214    format!("{:.1} GiB", bytes / 1073741824.0)
215}
216
217pub fn format_plan(plan: &Plan) -> String {
218    let context = plan.context.map_or_else(
219        || "default context".to_owned(),
220        |tokens| format!("{tokens} tokens"),
221    );
222    let context = match plan.kv_cache_type {
223        Some(kind) => format!("{context}, KV {kind}"),
224        None => context,
225    };
226    let mut text = format!(
227        "Plan for {} ({context}, {} backend{})\n\n",
228        plan.model,
229        plan.backend,
230        if plan.unified_memory {
231            ", unified memory"
232        } else {
233            ""
234        }
235    );
236    if !plan.kv_cache_known {
237        text.push_str(
238            "KV cache at this context is unknown (no sourced attention geometry); sizes cover weights only.\n\n",
239        );
240    }
241    for path in &plan.paths {
242        let label = match path.gpu_count {
243            Some(1) => format!("{} (1 GPU)", path.path),
244            Some(count) if path.path != "cpu-offload" => format!("{} ({count} GPUs)", path.path),
245            _ => path.path.to_owned(),
246        };
247        let fit = if path.fits {
248            format!(
249                "fits {} using {} of {}",
250                path.quant.as_deref().unwrap_or("?"),
251                path.required_bytes
252                    .map_or_else(|| "unknown".to_owned(), gib),
253                gib(path.usable_bytes)
254            )
255        } else if let Some(shortfall) = path.shortfall_bytes {
256            format!("does not fit; needs {} more usable memory", gib(shortfall))
257        } else {
258            format!("does not fit ({})", path.reason.unwrap_or("unknown"))
259        };
260        let speed = if path.throughput.known {
261            format!(
262                "{}-{} tok/s estimated",
263                path.throughput.low_tok_per_sec, path.throughput.high_tok_per_sec
264            )
265        } else if path.fits && path.throughput_basis.contains("not modeled") {
266            "speed unknown (not modeled; measure with llmup bench)".to_owned()
267        } else if path.fits {
268            "speed unknown (no sourced throughput data)".to_owned()
269        } else {
270            String::new()
271        };
272        text.push_str(&format!("  {label:<22} {fit}"));
273        if !speed.is_empty() {
274            text.push_str(&format!("; {speed}"));
275        }
276        text.push('\n');
277    }
278    text.push_str(&match plan.recommended {
279        Some(path) => format!("\nRecommended path: {path}\n"),
280        None => "\nNo path fits this model at this context.\n".to_owned(),
281    });
282    text
283}
rigspark_core/plan.rs

rigspark_core/
plan.rs
