diff --git a/crates/goose-server/src/routes/local_inference.rs b/crates/goose-server/src/routes/local_inference.rs index 8ee6fb6f78..1ee5f903c0 100644 --- a/crates/goose-server/src/routes/local_inference.rs +++ b/crates/goose-server/src/routes/local_inference.rs @@ -48,7 +48,6 @@ pub struct LocalModelResponse { pub status: ModelDownloadStatus, pub recommended: bool, pub settings: ModelSettings, - /// Estimated inference speed tier (fast/medium/slow) pub speed_tier: SpeedTier, } diff --git a/crates/goose/src/providers/local_inference/local_model_registry.rs b/crates/goose/src/providers/local_inference/local_model_registry.rs index 7adf1c0eb4..bf96011038 100644 --- a/crates/goose/src/providers/local_inference/local_model_registry.rs +++ b/crates/goose/src/providers/local_inference/local_model_registry.rs @@ -132,7 +132,6 @@ pub struct LocalModelEntry { pub settings: ModelSettings, #[serde(default)] pub size_bytes: u64, - /// Quality rank from quantization (1-28, higher is better quality) #[serde(default = "default_quality_rank")] pub quality_rank: u8, } diff --git a/crates/goose/src/providers/local_inference/recommender.rs b/crates/goose/src/providers/local_inference/recommender.rs index 5a8db0a85b..8655492374 100644 --- a/crates/goose/src/providers/local_inference/recommender.rs +++ b/crates/goose/src/providers/local_inference/recommender.rs @@ -1,10 +1,4 @@ //! Model recommendation engine for local inference. -//! -//! Provides hardware-aware model recommendations by: -//! - Detecting available memory (GPU VRAM or system RAM) -//! - Estimating model memory requirements including KV cache -//! - Scoring models based on quality and fit -//! - Estimating inference speed use super::hf_models::HfQuantVariant; use super::local_model_registry::{ @@ -16,8 +10,6 @@ use serde::{Deserialize, Serialize}; use std::fmt; use utoipa::ToSchema; -/// Bytes per parameter for each quantization level. -/// Used to estimate parameter count from file size and for memory calculations. const QUANT_BPP: &[(&str, f64)] = &[ ("F32", 4.0), ("BF16", 2.0), @@ -55,7 +47,6 @@ const QUANT_BPP: &[(&str, f64)] = &[ ("MXFP4_MOE", 0.58), ]; -/// Get bytes-per-parameter for a quantization level. pub fn quant_bpp(quant: &str) -> f64 { let upper = quant.to_uppercase(); QUANT_BPP @@ -65,7 +56,6 @@ pub fn quant_bpp(quant: &str) -> f64 { .unwrap_or(0.58) // Default to Q4_K_M-ish } -/// Estimate parameter count in billions from file size and quantization. pub fn estimate_params_billion(size_bytes: u64, quant: &str) -> f64 { let bpp = quant_bpp(quant); if bpp == 0.0 || size_bytes == 0 { @@ -74,12 +64,6 @@ pub fn estimate_params_billion(size_bytes: u64, quant: &str) -> f64 { (size_bytes as f64) / bpp / 1e9 } -/// Estimate memory required in GB to run a model with given context length. -/// -/// Memory formula: model_weights + KV_cache + overhead -/// - model_weights = params_b * bpp (already in the file size) -/// - KV_cache ≈ 0.000008 * params_b * context_length (in GB) -/// - overhead ≈ 0.5 GB for runtime buffers pub fn estimate_memory_gb(size_bytes: u64, quant: &str, context_length: u32) -> f64 { let params_b = estimate_params_billion(size_bytes, quant); let model_mem_gb = size_bytes as f64 / 1e9; @@ -88,7 +72,6 @@ pub fn estimate_memory_gb(size_bytes: u64, quant: &str, context_length: u32) -> model_mem_gb + kv_cache_gb + overhead_gb } -/// Speed tier for user-friendly display. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, ToSchema)] #[serde(rename_all = "lowercase")] pub enum SpeedTier { @@ -97,19 +80,14 @@ pub enum SpeedTier { Slow, } -/// How the model will be executed based on available memory. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RunMode { - /// Fully loaded in GPU VRAM - fastest Gpu, - /// Partial GPU with CPU offload - slower due to memory transfers CpuOffload, - /// CPU only, no GPU acceleration CpuOnly, } impl RunMode { - /// Speed multiplier for this run mode (relative to full GPU). pub fn speed_multiplier(&self) -> f64 { match self { RunMode::Gpu => 1.0, @@ -119,20 +97,16 @@ impl RunMode { } } -/// Predict how the model will run based on memory requirements and availability. pub fn predict_run_mode(model_memory_gb: f64, vram_gb: f64, has_gpu: bool) -> RunMode { if !has_gpu { return RunMode::CpuOnly; } - // If model fits in VRAM with some headroom, it'll run fully on GPU if model_memory_gb <= vram_gb * 0.95 { RunMode::Gpu } else if model_memory_gb <= vram_gb * 2.0 { - // Model is larger than VRAM but not massively so - will offload to CPU RunMode::CpuOffload } else { - // Model is way too large for GPU - effectively CPU only RunMode::CpuOnly } } @@ -149,10 +123,6 @@ impl fmt::Display for SpeedTier { impl SpeedTier { pub fn from_tps(tps: f64) -> Self { - // Thresholds calibrated for practical usability: - // - Fast: comfortable interactive use - // - Medium: usable but noticeable delay - // - Slow: significant wait between responses if tps >= 20.0 { SpeedTier::Fast } else if tps >= 8.0 { @@ -163,7 +133,6 @@ impl SpeedTier { } } -/// Quantization speed multiplier - lower precision is faster. fn quant_speed_multiplier(quant: &str) -> f64 { let upper = quant.to_uppercase(); if upper.starts_with("F32") { @@ -189,14 +158,12 @@ fn quant_speed_multiplier(quant: &str) -> f64 { } } -/// Get the number of CPU cores available. pub fn cpu_core_count() -> usize { std::thread::available_parallelism() .map(|p| p.get()) .unwrap_or(4) } -/// Get available GPU/accelerator VRAM in GB. pub fn available_vram_gb() -> f64 { let devices = list_llama_ggml_backend_devices(); devices @@ -214,55 +181,25 @@ pub fn available_vram_gb() -> f64 { .unwrap_or(0.0) } -/// Estimate inference speed in tokens per second. -/// -/// Accounts for model size, quantization, GPU availability, core count, and run mode. -/// A large model on a small GPU will be penalized for CPU offloading. pub fn estimate_speed_tps(params_billion: f64, quant: &str, has_gpu: bool) -> f64 { - estimate_speed_tps_with_memory(params_billion, quant, has_gpu, available_vram_gb()) -} - -/// Estimate inference speed with explicit VRAM availability. -/// This allows callers to provide known memory values for more accurate estimates. -pub fn estimate_speed_tps_with_memory( - params_billion: f64, - quant: &str, - has_gpu: bool, - vram_gb: f64, -) -> f64 { - // Base speed constant by backend (empirically calibrated) - // These represent rough tok/s for a 1B parameter model - // Modern GPUs (M3, RTX 4090, etc) can do 200+ tok/s on 1B models let base_k = if has_gpu { - // GPU backends - Metal/CUDA/etc - // Calibrated for M3/M4 and modern NVIDIA cards 300.0 + } else if cfg!(target_arch = "aarch64") { + 80.0 } else { - // CPU-only is much slower - if cfg!(target_arch = "aarch64") { - 80.0 // ARM CPUs (Apple Silicon CPU-only, etc) - } else { - 60.0 // x86 CPUs - } + 60.0 }; - // Speed scales inversely with parameter count let params = params_billion.max(0.1); let base_tps = base_k / params; - // 10% bonus for systems with 8+ cores (matches llmfit) let core_mult = if cpu_core_count() >= 8 { 1.1 } else { 1.0 }; - // Estimate model memory to predict run mode - // Use a rough approximation: params * 0.6 for Q4_K_M-ish models + overhead let model_mem_gb = params_billion * quant_bpp(quant) + 1.0; - let run_mode = predict_run_mode(model_mem_gb, vram_gb, has_gpu); - - // Apply all multipliers + let run_mode = predict_run_mode(model_mem_gb, available_vram_gb(), has_gpu); base_tps * quant_speed_multiplier(quant) * core_mult * run_mode.speed_multiplier() } -/// Check if a GPU/accelerator is available for inference. pub fn has_gpu_accelerator() -> bool { let devices = list_llama_ggml_backend_devices(); devices.iter().any(|d| { @@ -275,7 +212,6 @@ pub fn has_gpu_accelerator() -> bool { }) } -/// Result of model fit analysis. #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct ModelFit { /// Whether the model fits in available memory @@ -293,8 +229,6 @@ pub struct ModelFit { } impl ModelFit { - /// Compute a combined score for ranking models. - /// Higher is better. Balances fit (headroom) with quality. pub fn score(&self) -> f64 { if !self.fits { return -1.0; @@ -314,7 +248,6 @@ impl ModelFit { } } -/// Analyze how well a model fits the available memory. pub fn analyze_model_fit( size_bytes: u64, quant: &str, @@ -343,10 +276,6 @@ pub fn analyze_model_fit( } } -/// Get available memory for inference in bytes. -/// -/// Prefers GPU/accelerator memory if available, falls back to CPU/system RAM. -/// On Apple Silicon with unified memory, this returns the available unified memory. pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 { let _ = &runtime.backend(); let devices = list_llama_ggml_backend_devices(); @@ -379,15 +308,8 @@ pub fn available_inference_memory_bytes(runtime: &InferenceRuntime) -> u64 { } } -/// Default context length used for recommendations when not specified. const DEFAULT_CONTEXT_LENGTH: u32 = 8192; -/// Recommend the best local model from the registry for the current hardware. -/// -/// Returns the model ID of the recommended model. Prefers: -/// 1. Downloaded models that fit in memory -/// 2. Highest quality model that fits -/// 3. Falls back to smallest featured model if nothing fits pub fn recommend_local_model(runtime: &InferenceRuntime) -> String { let available_memory = available_inference_memory_bytes(runtime); @@ -434,10 +356,19 @@ pub fn recommend_local_model(runtime: &InferenceRuntime) -> String { FEATURED_MODELS[0].to_string() } -/// Recommend the best quantization variant for a given model based on available memory. -/// -/// Returns the index of the recommended variant in the input slice. -/// Memory requirements (including KV cache and overhead) are calculated by analyze_model_fit. +#[cfg(test)] +fn make_variant(quant: &str, size_bytes: u64) -> HfQuantVariant { + use super::hf_models::quant_quality_rank; + HfQuantVariant { + quantization: quant.to_string(), + size_bytes, + filename: format!("model-{}.gguf", quant), + download_url: String::new(), + description: "", + quality_rank: quant_quality_rank(quant), + } +} + pub fn recommend_variant( variants: &[HfQuantVariant], available_memory_bytes: u64, @@ -472,197 +403,32 @@ pub fn recommend_variant( #[cfg(test)] mod tests { use super::*; + use test_case::test_case; - #[test] - fn test_quant_bpp() { - assert!((quant_bpp("Q4_K_M") - 0.58).abs() < 0.01); - assert!((quant_bpp("Q8_0") - 1.05).abs() < 0.01); - assert!((quant_bpp("F16") - 2.0).abs() < 0.01); - assert!((quant_bpp("unknown") - 0.58).abs() < 0.01); // Default + fn standard_variants() -> Vec { + vec![ + make_variant("Q2_K", 3_000_000_000), + make_variant("Q4_K_M", 5_000_000_000), + make_variant("Q6_K", 7_000_000_000), + make_variant("Q8_0", 9_000_000_000), + ] } - #[test] - fn test_estimate_params_billion() { - // A 4GB Q4_K_M file should be ~7B params - let params = estimate_params_billion(4_000_000_000, "Q4_K_M"); - assert!(params > 6.0 && params < 8.0); - - // A 1GB Q4_K_M file should be ~1.7B params - let params = estimate_params_billion(1_000_000_000, "Q4_K_M"); - assert!(params > 1.5 && params < 2.0); - } - - #[test] - fn test_estimate_memory_gb() { - // 4GB file with 8K context - let mem = estimate_memory_gb(4_000_000_000, "Q4_K_M", 8192); - // Should be slightly more than 4GB (file + KV cache + overhead) - assert!(mem > 4.0 && mem < 5.0); - } - - #[test] - fn test_analyze_model_fit() { - // 4GB model with 8GB available - let fit = analyze_model_fit(4_000_000_000, "Q4_K_M", 19, 8192, 8_000_000_000); - assert!(fit.fits); - assert!(fit.headroom_percent > 0.0); - - // 4GB model with 3GB available - shouldn't fit - let fit = analyze_model_fit(4_000_000_000, "Q4_K_M", 19, 8192, 3_000_000_000); - assert!(!fit.fits); - assert!(fit.headroom_percent < 0.0); - } - - #[test] - fn test_model_fit_score() { - // High quality, good headroom - let fit1 = ModelFit { - fits: true, - estimated_memory_gb: 4.0, - available_memory_gb: 8.0, - headroom_percent: 50.0, - quality_rank: 25, - params_billion: 7.0, - }; - - // Lower quality, same headroom - let fit2 = ModelFit { - fits: true, - estimated_memory_gb: 4.0, - available_memory_gb: 8.0, - headroom_percent: 50.0, - quality_rank: 10, - params_billion: 7.0, - }; - - assert!(fit1.score() > fit2.score()); - - // Doesn't fit - let fit3 = ModelFit { - fits: false, - estimated_memory_gb: 10.0, - available_memory_gb: 8.0, - headroom_percent: -25.0, - quality_rank: 25, - params_billion: 14.0, - }; - assert!(fit3.score() < 0.0); - } - - #[test] - fn test_recommend_variant() { - let variants = vec![ - HfQuantVariant { - quantization: "Q2_K".into(), - size_bytes: 2_000_000_000, - filename: "m-Q2_K.gguf".into(), - download_url: String::new(), - description: "Small", - quality_rank: 7, - }, - HfQuantVariant { - quantization: "Q4_K_M".into(), - size_bytes: 4_000_000_000, - filename: "m-Q4_K_M.gguf".into(), - download_url: String::new(), - description: "Medium", - quality_rank: 19, - }, - HfQuantVariant { - quantization: "Q8_0".into(), - size_bytes: 8_000_000_000, - filename: "m-Q8_0.gguf".into(), - download_url: String::new(), - description: "Large", - quality_rank: 25, - }, - ]; - - // Memory estimation includes KV cache + overhead: - // Q2_K (2GB file): ~2GB + 0.28GB KV + 0.5GB overhead = ~2.78GB needed - // Q4_K_M (4GB file): ~4GB + 0.45GB KV + 0.5GB overhead = ~4.95GB needed - // Q8_0 (8GB file): ~8GB + 0.50GB KV + 0.5GB overhead = ~9.0GB needed - // - // Scoring balances quality (70%) with headroom (30%), so a model with - // comfortable headroom may beat a higher-quality model that barely fits. - - // With 6GB available: Q4_K_M fits (4.95 < 6), Q8_0 doesn't (9.0 > 6) - assert_eq!(recommend_variant(&variants, 6_000_000_000), Some(1)); - - // With 10GB available: Q8_0 fits with ~10% headroom (ideal range ≤20%) - // Q8_0: headroom = 10%, headroom_norm = 1.0, quality = 25/28 = 0.89 - // score = 0.7*0.89 + 0.3*1.0 = 0.92 - // Q4_K_M: headroom = 50%, headroom_norm = 1.0 - 0.75*0.5 = 0.62, quality = 19/28 = 0.68 - // score = 0.7*0.68 + 0.3*0.62 = 0.66 - // Q8_0 wins - higher quality with good headroom - assert_eq!(recommend_variant(&variants, 10_000_000_000), Some(2)); - - // With 16GB available: Q8_0 still wins (higher quality, acceptable headroom) - // Q8_0: headroom = 44%, headroom_norm = 1.0 - 0.75*0.4 = 0.7 - // score = 0.7*0.89 + 0.3*0.7 = 0.83 - assert_eq!(recommend_variant(&variants, 16_000_000_000), Some(2)); - - // With 2GB available: nothing fits (Q2_K needs ~2.78GB) - assert_eq!(recommend_variant(&variants, 2_000_000_000), None); - - // With 4GB available: Q2_K fits (2.78 < 4) - assert_eq!(recommend_variant(&variants, 4_000_000_000), Some(0)); - } - - #[test] - fn test_predict_run_mode() { - // Model fits in VRAM - assert_eq!(predict_run_mode(4.0, 8.0, true), RunMode::Gpu); - - // Model slightly larger than VRAM - will offload - assert_eq!(predict_run_mode(10.0, 8.0, true), RunMode::CpuOffload); - - // Model way too large for VRAM - assert_eq!(predict_run_mode(20.0, 8.0, true), RunMode::CpuOnly); - - // No GPU - assert_eq!(predict_run_mode(4.0, 8.0, false), RunMode::CpuOnly); - assert_eq!(predict_run_mode(4.0, 0.0, true), RunMode::CpuOnly); - } - - #[test] - fn test_estimate_speed_tps_with_memory() { - // 1B model fitting in 16GB VRAM should be fast (full GPU) - let tps = estimate_speed_tps_with_memory(1.0, "Q4_K_M", true, 16.0); - assert!(tps > 100.0, "1B on GPU should be >100 tps, got {}", tps); - - // 7B model fitting in 16GB VRAM - let tps = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 16.0); - assert!(tps > 20.0, "7B on GPU should be >20 tps, got {}", tps); - - // 7B model on small 4GB VRAM - will need offload, slower - let tps_offload = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 4.0); - let tps_full = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 16.0); - assert!( - tps_offload < tps_full, - "Offload should be slower: {} vs {}", - tps_offload, - tps_full + // available memory → expected quant + #[test_case(4_500_000_000, Some("Q2_K") ; "tight memory only fits smallest")] + #[test_case(7_000_000_000, Some("Q4_K_M") ; "moderate memory picks mid quality")] + #[test_case(9_000_000_000, Some("Q6_K") ; "good memory picks high quality")] + #[test_case(24_000_000_000, Some("Q8_0") ; "plenty of memory picks best quality")] + #[test_case(2_000_000_000, None ; "insufficient memory recommends nothing")] + fn recommend_variant_picks_best_quality_that_fits( + available_bytes: u64, + expected_quant: Option<&str>, + ) { + let variants = standard_variants(); + let result = recommend_variant(&variants, available_bytes); + assert_eq!( + result.map(|i| variants[i].quantization.as_str()), + expected_quant, ); - - // CPU is slower than GPU - let cpu_tps = estimate_speed_tps_with_memory(7.0, "Q4_K_M", false, 0.0); - assert!(tps_full > cpu_tps, "GPU should be faster than CPU"); - - // Lower quant is faster - let q4_tps = estimate_speed_tps_with_memory(7.0, "Q4_K_M", true, 16.0); - let q8_tps = estimate_speed_tps_with_memory(7.0, "Q8_0", true, 16.0); - assert!(q4_tps > q8_tps, "Q4 should be faster than Q8"); - } - - #[test] - fn test_speed_tier() { - // New thresholds: Fast >= 20, Medium >= 8, Slow < 8 - assert_eq!(SpeedTier::from_tps(100.0), SpeedTier::Fast); - assert_eq!(SpeedTier::from_tps(20.0), SpeedTier::Fast); - assert_eq!(SpeedTier::from_tps(19.9), SpeedTier::Medium); - assert_eq!(SpeedTier::from_tps(8.0), SpeedTier::Medium); - assert_eq!(SpeedTier::from_tps(7.9), SpeedTier::Slow); - assert_eq!(SpeedTier::from_tps(3.0), SpeedTier::Slow); } }