From 63e2a74ad5286c5e31682dbdbaa70a4595f6651a Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sun, 26 Jul 2026 09:10:18 +0200 Subject: [PATCH] Manage the KV cache disc usage from preferences and stats Co-Authored-By: Claude Opus 5 --- .../down.sql | 4 + .../up.sql | 4 + src/app.rs | 145 +++++++++++++- src/app/preferences.rs | 36 ++++ src/app/projects.rs | 37 +++- src/app/view/preferences.rs | 51 +++++ src/app/view/stats.rs | 133 +++++++++++++ src/database.rs | 56 +++++- src/engine.rs | 23 ++- src/engine/kvstore.rs | 24 +-- src/engine/metal.rs | 6 + src/metrics.rs | 179 +++++++++++++++++- src/schema.rs | 4 + src/settings.rs | 127 ++++++++++++- 14 files changed, 794 insertions(+), 35 deletions(-) create mode 100644 migrations/20260726140000_add_kv_cache_preferences/down.sql create mode 100644 migrations/20260726140000_add_kv_cache_preferences/up.sql diff --git a/migrations/20260726140000_add_kv_cache_preferences/down.sql b/migrations/20260726140000_add_kv_cache_preferences/down.sql new file mode 100644 index 0000000..844b9aa --- /dev/null +++ b/migrations/20260726140000_add_kv_cache_preferences/down.sql @@ -0,0 +1,4 @@ +ALTER TABLE preferences DROP COLUMN kv_continued_interval_tokens; +ALTER TABLE preferences DROP COLUMN kv_cold_max_tokens; +ALTER TABLE preferences DROP COLUMN kv_min_tokens; +ALTER TABLE preferences DROP COLUMN kv_budget_gib; diff --git a/migrations/20260726140000_add_kv_cache_preferences/up.sql b/migrations/20260726140000_add_kv_cache_preferences/up.sql new file mode 100644 index 0000000..d914ae1 --- /dev/null +++ b/migrations/20260726140000_add_kv_cache_preferences/up.sql @@ -0,0 +1,4 @@ +ALTER TABLE preferences ADD COLUMN kv_budget_gib BIGINT; +ALTER TABLE preferences ADD COLUMN kv_min_tokens INTEGER; +ALTER TABLE preferences ADD COLUMN kv_cold_max_tokens INTEGER; +ALTER TABLE preferences ADD COLUMN kv_continued_interval_tokens INTEGER; diff --git a/src/app.rs b/src/app.rs index 8546e98..152602e 100644 --- a/src/app.rs +++ b/src/app.rs @@ -15,21 +15,21 @@ use preferences::{parse_optional_gib, parse_streaming_cache}; use crate::database::{AppPreferences, Database, ProjectWithSessions, SessionState, StoredMessage}; #[cfg(target_os = "macos")] use crate::engine::ChatTurn; -use crate::metrics::{Metrics, MetricsSnapshot}; +use crate::metrics::{KvCacheReport, Metrics, MetricsSnapshot}; use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; #[cfg(target_os = "macos")] use crate::runtime::{ActiveGeneration, CheckpointTarget, GenerationEvent, GenerationService}; use crate::settings::{ - DiagnosticPreferences, ExecutionPreferences, GIB, GenerationPreferences, ReasoningMode, - RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, - StreamingCacheBudget, + DEFAULT_KV_BUDGET_GIB, DiagnosticPreferences, ExecutionPreferences, GIB, GenerationPreferences, + KvCachePreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, + SteeringPreferences, StreamingCacheBudget, }; use iced::widget::{markdown, scrollable, text_input}; use iced::{Size, Subscription, Task, keyboard, mouse, window}; use rfd::AsyncFileDialog; use std::collections::{HashMap, HashSet, VecDeque}; use std::fs; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::mpsc::{self, TryRecvError}; use std::sync::{Arc, Mutex, RwLock}; @@ -38,6 +38,8 @@ use std::time::{Duration, Instant}; const APP_ID: &str = "de.rfc1437.ds4server"; const METRICS_SAMPLE_INTERVAL: Duration = Duration::from_millis(200); +/// Disc scans are cheap but not free; the explorer does not need 5 Hz. +const CACHE_SCAN_INTERVAL: Duration = Duration::from_secs(2); pub(super) const MIN_SIDEBAR_WIDTH: i32 = 180; pub(super) const MAX_SIDEBAR_WIDTH: i32 = 520; @@ -82,6 +84,9 @@ pub(crate) struct App { metrics: Arc, pub(super) metrics_snapshot: MetricsSnapshot, pub(super) metrics_history: VecDeque, + /// Disc usage of the KV cache directories, rescanned while Stats is open. + pub(super) kv_cache_report: KvCacheReport, + last_cache_scan: Instant, last_http_requests: u64, #[cfg(target_os = "macos")] generation_service: Option, @@ -162,6 +167,10 @@ pub(crate) enum Message { PreferenceSteeringAttnChanged(String), PreferenceSimulatedMemoryChanged(String), PreferenceExpertProfileChanged(String), + PreferenceKvBudgetChanged(String), + PreferenceKvMinTokensChanged(String), + PreferenceKvColdMaxChanged(String), + PreferenceKvContinuedIntervalChanged(String), ResetPreferences, SavePreferences, DownloadArtifact(ManagedArtifactId), @@ -202,6 +211,8 @@ pub(crate) enum Message { DragSidebar(f32), EndSidebarDrag, MetricsTick, + DiscardCacheEntry(PathBuf), + ClearTransientCache, } impl App { @@ -214,6 +225,9 @@ impl App { Ok(draft) => draft, Err(error) => return Self::failed(error, main_window), }; + // Before the counters are taken, so they measure what is + // actually reachable. + sweep_orphan_checkpoints(&kv_cache_path(), &projects); let context_limit = preferences.context_tokens.max(0) as u32; // Reopen on the project we left, with a fresh draft chat. let last_project = preferences @@ -264,6 +278,8 @@ impl App { metrics, metrics_snapshot, metrics_history: VecDeque::with_capacity(120), + kv_cache_report: KvCacheReport::default(), + last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL, last_http_requests: 0, #[cfg(target_os = "macos")] generation_service, @@ -346,6 +362,8 @@ impl App { metrics, metrics_snapshot, metrics_history: VecDeque::with_capacity(120), + kv_cache_report: KvCacheReport::default(), + last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL, last_http_requests: 0, #[cfg(target_os = "macos")] generation_service, @@ -422,7 +440,12 @@ impl App { } } Message::ShowChat => self.detail_tab = DetailTab::Chat, - Message::ShowStats => self.detail_tab = DetailTab::Stats, + Message::ShowStats => { + self.detail_tab = DetailTab::Stats; + self.scan_kv_cache(); + } + Message::DiscardCacheEntry(path) => self.discard_cache_entry(&path), + Message::ClearTransientCache => self.clear_transient_cache(), // The sidebar starts at the window's left edge, so the cursor's x is // the width the user is asking for. Only the final width is stored. Message::StartSidebarDrag => self.sidebar_drag = true, @@ -652,6 +675,22 @@ impl App { self.preference_draft.expert_profile_path = value; self.preference_error = None; } + Message::PreferenceKvBudgetChanged(value) => { + self.preference_draft.kv_budget_gib = value; + self.preference_error = None; + } + Message::PreferenceKvMinTokensChanged(value) => { + self.preference_draft.kv_min_tokens = value; + self.preference_error = None; + } + Message::PreferenceKvColdMaxChanged(value) => { + self.preference_draft.kv_cold_max_tokens = value; + self.preference_error = None; + } + Message::PreferenceKvContinuedIntervalChanged(value) => { + self.preference_draft.kv_continued_interval_tokens = value; + self.preference_error = None; + } Message::ResetPreferences => { self.preference_draft.reset(); self.preference_error = None; @@ -801,6 +840,7 @@ impl App { self.tokens_per_second = None; } self.reload_projects(); + self.finish_cache_change(); } Err(error) => self.error = Some(error), } @@ -936,6 +976,7 @@ impl App { match database.delete_session(session_id) { Ok(()) => { let _ = fs::remove_file(session_checkpoint_path(session_id)); + self.finish_cache_change(); if self.session_menu == Some(session_id) { self.session_menu = None; } @@ -1050,6 +1091,56 @@ impl App { self.metrics_history.pop_front(); } self.metrics_snapshot = snapshot; + if self.detail_tab == DetailTab::Stats + && self.last_cache_scan.elapsed() >= CACHE_SCAN_INTERVAL + { + self.scan_kv_cache(); + } + } + + /// Removes one cache file. Checkpoints are disposable: the next turn on that + /// conversation prefills again. + fn discard_cache_entry(&mut self, path: &Path) { + if let Err(error) = fs::remove_file(path) { + self.error = Some(format!("Could not delete the checkpoint: {error}")); + return; + } + let _ = fs::remove_file(path.with_extension("meta")); + self.finish_cache_change(); + } + + /// Empties the transient store only. Session checkpoints belong to their + /// session and are removed with it. + fn clear_transient_cache(&mut self) { + let Ok(files) = fs::read_dir(transient_cache_path()) else { + return; + }; + for file in files.flatten() { + if file.metadata().is_ok_and(|metadata| metadata.is_file()) + && let Err(error) = fs::remove_file(file.path()) + { + self.error = Some(format!("Could not clear the transient cache: {error}")); + break; + } + } + self.finish_cache_change(); + } + + fn finish_cache_change(&mut self) { + self.metrics.rescan_cache(&kv_cache_path()); + self.metrics_snapshot = self.metrics.snapshot(); + self.scan_kv_cache(); + } + + fn scan_kv_cache(&mut self) { + let budget = self + .preferences + .runtime() + .map_or(DEFAULT_KV_BUDGET_GIB * GIB, |runtime| { + runtime.kv_cache.settings().budget_bytes + }); + self.kv_cache_report = crate::metrics::kv_cache_report(&kv_cache_path(), budget); + self.last_cache_scan = Instant::now(); } } @@ -1142,16 +1233,50 @@ fn scroll_chat_to_end() -> Task { scrollable::snap_to(chat_scroll_id(), scrollable::RelativeOffset::END) } +/// Deletes session checkpoints whose session is gone. Deleting a session or a +/// project removes its checkpoint directly; this catches the cases that path +/// misses — an interrupted or failed delete, or files left by an older +/// database — so no unreachable checkpoint keeps occupying disc. +fn sweep_orphan_checkpoints(directory: &Path, projects: &[ProjectWithSessions]) -> bool { + let known = projects + .iter() + .flat_map(|item| item.sessions.iter().map(|session| session.id)) + .collect::>(); + let Ok(files) = fs::read_dir(directory) else { + return false; + }; + let mut removed = false; + for file in files.flatten() { + let path = file.path(); + if path.extension().is_none_or(|value| value != "bin") { + continue; + } + // Only numbered session checkpoints are ours to judge; anything else in + // the directory is left alone. + let orphan = path + .file_stem() + .and_then(|value| value.to_str()) + .and_then(|value| value.parse::().ok()) + .is_some_and(|id| !known.contains(&id)); + if orphan { + removed |= fs::remove_file(&path).is_ok(); + } + } + removed +} + +fn kv_cache_path() -> PathBuf { + application_support_path().join("kv-cache") +} + fn session_checkpoint_path(session_id: i32) -> PathBuf { - application_support_path() - .join("kv-cache") - .join(format!("{session_id}.bin")) + kv_cache_path().join(format!("{session_id}.bin")) } /// Content-addressed KV cache for turns that belong to no stored session: the /// HTTP endpoint and the app's own one-shot requests share it. fn transient_cache_path() -> PathBuf { - application_support_path().join("kv-cache").join("http") + kv_cache_path().join("http") } pub(crate) fn app_icon() -> window::Icon { diff --git a/src/app/preferences.rs b/src/app/preferences.rs index 9ed3721..5319ba1 100644 --- a/src/app/preferences.rs +++ b/src/app/preferences.rs @@ -37,6 +37,10 @@ pub(super) struct PreferenceDraft { pub(super) directional_steering_attn: String, pub(super) simulated_used_memory_gib: String, pub(super) expert_profile_path: String, + pub(super) kv_budget_gib: String, + pub(super) kv_min_tokens: String, + pub(super) kv_cold_max_tokens: String, + pub(super) kv_continued_interval_tokens: String, } impl PreferenceDraft { @@ -121,6 +125,22 @@ impl PreferenceDraft { .simulated_used_memory_gib .map_or_else(String::new, |value| value.to_string()), expert_profile_path: runtime.diagnostics.expert_profile_path.unwrap_or_default(), + kv_budget_gib: runtime + .kv_cache + .budget_gib + .map_or_else(String::new, |value| value.to_string()), + kv_min_tokens: runtime + .kv_cache + .min_tokens + .map_or_else(String::new, |value| value.to_string()), + kv_cold_max_tokens: runtime + .kv_cache + .cold_max_tokens + .map_or_else(String::new, |value| value.to_string()), + kv_continued_interval_tokens: runtime + .kv_cache + .continued_interval_tokens + .map_or_else(String::new, |value| value.to_string()), }) } @@ -182,6 +202,10 @@ impl PreferenceDraft { self.directional_steering_attn.clear(); self.simulated_used_memory_gib.clear(); self.expert_profile_path.clear(); + self.kv_budget_gib.clear(); + self.kv_min_tokens.clear(); + self.kv_cold_max_tokens.clear(); + self.kv_continued_interval_tokens.clear(); } pub(super) fn execution(&self) -> Result { @@ -241,6 +265,18 @@ impl PreferenceDraft { )?, expert_profile_path: optional_text(&self.expert_profile_path), }, + kv_cache: KvCachePreferences { + budget_gib: parse_optional_gib("KV cache budget", &self.kv_budget_gib)?, + min_tokens: parse_optional_u32("KV cache minimum tokens", &self.kv_min_tokens)?, + cold_max_tokens: parse_optional_u32( + "KV cache cold maximum", + &self.kv_cold_max_tokens, + )?, + continued_interval_tokens: parse_optional_u32( + "KV cache continued interval", + &self.kv_continued_interval_tokens, + )?, + }, }) } } diff --git a/src/app/projects.rs b/src/app/projects.rs index f87de3b..3315bfb 100644 --- a/src/app/projects.rs +++ b/src/app/projects.rs @@ -132,7 +132,12 @@ impl App { pub(super) fn reload_projects(&mut self) { if let Some(database) = &mut self.database { match database.load_projects() { - Ok(projects) => self.projects = projects, + Ok(projects) => { + self.projects = projects; + if super::sweep_orphan_checkpoints(&super::kv_cache_path(), &self.projects) { + self.finish_cache_change(); + } + } Err(error) => self.error = Some(error), } } @@ -190,6 +195,36 @@ mod tests { assert_eq!(draft_title(&projects, 99), "Session 1"); } + #[test] + fn sweeping_removes_checkpoints_of_sessions_that_no_longer_exist() { + let directory = + std::env::temp_dir().join(format!("ds4-server-sweep-{}", std::process::id())); + std::fs::create_dir_all(&directory).unwrap(); + for name in ["1.bin", "2.bin", "notes.bin", "7.bin"] { + std::fs::write(directory.join(name), b"payload").unwrap(); + } + let projects = vec![ProjectWithSessions { + project: project(1, "First"), + sessions: vec![session(1, 1), session(2, 1)], + }]; + + assert!(super::super::sweep_orphan_checkpoints( + &directory, &projects + )); + + assert!(directory.join("1.bin").exists()); + assert!(directory.join("2.bin").exists()); + // Not a session checkpoint, so not ours to delete. + assert!(directory.join("notes.bin").exists()); + assert!(!directory.join("7.bin").exists()); + // Nothing left to sweep on the next pass. + assert!(!super::super::sweep_orphan_checkpoints( + &directory, &projects + )); + + std::fs::remove_dir_all(directory).unwrap(); + } + fn project(id: i32, name: &str) -> Project { Project { id, diff --git a/src/app/view/preferences.rs b/src/app/view/preferences.rs index 2057daf..9cffc83 100644 --- a/src/app/view/preferences.rs +++ b/src/app/view/preferences.rs @@ -479,12 +479,63 @@ impl App { ] .spacing(10), ); + let kv_cache_group = preference_group( + "KV CACHE", + column![ + preference_input_row( + "Disc budget (GiB)", + "How much disc the reusable prompt checkpoints of the endpoint and one-shot requests may occupy. When a new checkpoint does not fit, the least valuable older ones are deleted; they cost only a prefill to rebuild. Blank uses DS4's 4 GiB.", + text_input("4", &self.preference_draft.kv_budget_gib) + .on_input(Message::PreferenceKvBudgetChanged), + ), + preference_input_row( + "Minimum tokens", + "Shortest prompt worth keeping a checkpoint for. Below this the prefill is cheaper than the disc traffic of storing and loading it.", + text_input("512", &self.preference_draft.kv_min_tokens) + .on_input(Message::PreferenceKvMinTokensChanged), + ), + preference_input_row( + "Cold maximum tokens", + "Largest first prompt of a conversation that is still stored. Very long cold prompts write huge checkpoints that are rarely asked for a second time; 0 stops storing cold prompts entirely.", + text_input("30000", &self.preference_draft.kv_cold_max_tokens) + .on_input(Message::PreferenceKvColdMaxChanged), + ), + preference_input_row( + "Continued interval tokens", + "How far a conversation has to grow before the next checkpoint of it is written. Wider spacing writes less and keeps fewer near-identical copies; 0 stores no continued checkpoints at all.", + text_input("10000", &self.preference_draft.kv_continued_interval_tokens) + .on_input(Message::PreferenceKvContinuedIntervalChanged), + ), + text("Session checkpoints are not covered by the budget: they belong to their session and go away with it. Blank values keep the DS4 defaults.") + .size(12), + text( + self.preference_draft + .runtime() + .map_or_else( + |_| "Effective KV cache settings will appear after valid values are entered.".to_owned(), + |runtime| { + let settings = runtime.kv_cache.settings(); + format!( + "Transient store: budget {} GiB • minimum {} • cold max {} • continued every {}", + settings.budget_bytes / GIB, + settings.min_tokens, + if settings.cold_max_tokens == 0 { "off".to_owned() } else { settings.cold_max_tokens.to_string() }, + if settings.continued_interval_tokens == 0 { "off".to_owned() } else { settings.continued_interval_tokens.to_string() }, + ) + }, + ) + ) + .size(12), + ] + .spacing(10), + ); let mut fields = column![ model_group, endpoint_group, generation_group, execution_group, acceleration_group, + kv_cache_group, steering_group, ] .spacing(12); diff --git a/src/app/view/stats.rs b/src/app/view/stats.rs index 150e5ce..acde7e0 100644 --- a/src/app/view/stats.rs +++ b/src/app/view/stats.rs @@ -319,6 +319,7 @@ impl App { .spacing(9) .into(), ); + let disc = stats_panel("KV CACHE DISC USAGE", self.cache_explorer()); let server = stats_panel( "LOCAL SERVER", column![ @@ -353,6 +354,7 @@ impl App { throughput, kv_io, requests, + disc, row![model, runtime].spacing(10), row![cache, server].spacing(10), text("Counters are published by the runtime with relaxed atomics and sampled by the UI every 200 ms.") @@ -368,4 +370,135 @@ impl App { .height(Length::Fill) .into() } + + /// Disc usage of both cache buckets, split by age. The bar is the age + /// distribution of what is stored; the budget governs the transient store + /// only and is reported on its own row. + fn cache_explorer(&self) -> Element<'_, Message> { + let usage = &self.kv_cache_report; + let total = usage.total_bytes(); + let capacity = total.max(1); + let mut bar = row![].height(14).spacing(2); + let mut legend = column![].spacing(6); + for (index, (label, _)) in crate::metrics::CACHE_AGE_BUCKETS.iter().enumerate() { + let bytes = usage.age_bytes[index]; + if bytes == 0 { + continue; + } + let color = AGE_COLORS[index]; + bar = bar.push( + container(Space::new(Length::Fill, Length::Fill)) + .width(Length::FillPortion(portion(bytes, capacity))) + .style(move |_| chart_bar_style(color)), + ); + legend = legend.push( + row![ + container(Space::new(9, 9)).style(move |_| chart_bar_style(color)), + text(*label).size(12).color(muted_text()), + Space::with_width(Length::Fill), + text(format!( + "{} · {:.0}%", + format_bytes(bytes), + bytes as f32 / total.max(1) as f32 * 100.0 + )) + .size(12), + ] + .spacing(8) + .align_y(Alignment::Center), + ); + } + if total == 0 { + bar = bar.push( + container(Space::new(Length::Fill, Length::Fill)) + .style(|_| chart_bar_style(muted_text().scale_alpha(0.25))), + ); + } + + let mut rows = column![].spacing(7); + for entry in &usage.entries { + let path = entry.path.clone(); + rows = rows.push( + row![ + text(entry.name.clone()).size(12), + Space::with_width(Length::Fill), + text(format!( + "{} · {} old", + format_bytes(entry.bytes), + format_duration(entry.age_seconds as f64) + )) + .size(12) + .color(muted_text()), + action_button(text("Discard").size(12)) + .on_press(Message::DiscardCacheEntry(path)), + ] + .spacing(10) + .align_y(Alignment::Center), + ); + } + + column![ + metric_row( + "Stored", + format!( + "{} · {} files", + format_bytes(total), + usage.session_files + usage.transient_files + ), + ), + container(bar).height(14).width(Length::Fill), + legend, + metric_row( + "Sessions", + format!( + "{} · {} files · not evicted", + format_bytes(usage.session_bytes), + usage.session_files + ), + ), + metric_row( + "Transient", + format!( + "{} of {} budget · {} files · {:.0}% full", + format_bytes(usage.transient_bytes), + format_bytes(usage.budget_bytes), + usage.transient_files, + usage.transient_bytes as f32 / usage.budget_bytes.max(1) as f32 * 100.0, + ), + ), + Space::with_height(2), + text(if usage.entries.is_empty() { + "No checkpoints stored yet." + } else { + "OLDEST CHECKPOINTS" + }) + .size(10) + .color(muted_text()), + rows, + row![ + text("Discarding a checkpoint costs one prefill to rebuild it.") + .size(11) + .color(muted_text()), + Space::with_width(Length::Fill), + action_button(text("Clear transient cache").size(12)) + .on_press(Message::ClearTransientCache), + ] + .spacing(10) + .align_y(Alignment::Center), + ] + .spacing(9) + .into() + } +} + +/// Age buckets from green (fresh) to grey (past every hit half-life). +const AGE_COLORS: [Color; 5] = [ + Color::from_rgb(0.28, 0.69, 0.44), + Color::from_rgb(0.33, 0.67, 1.0), + Color::from_rgb(0.62, 0.47, 1.0), + Color::from_rgb(0.94, 0.71, 0.27), + Color::from_rgb(0.45, 0.45, 0.48), +]; + +fn portion(bytes: u64, capacity: u64) -> u16 { + ((bytes.saturating_mul(1000) / capacity.max(1)) as u16).max(1) } diff --git a/src/database.rs b/src/database.rs index bb9ad2f..6bac420 100644 --- a/src/database.rs +++ b/src/database.rs @@ -6,8 +6,8 @@ use std::path::Path; use crate::schema::{messages, preferences, projects, sessions}; use crate::settings::{ - DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, ReasoningMode, - RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, + DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, KvCachePreferences, + ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, StreamingCacheBudget, }; @@ -58,6 +58,10 @@ pub struct AppPreferences { /// Project the app reopens on. Cleared when that project goes away. pub last_project_id: Option, pub sidebar_width: i32, + pub kv_budget_gib: Option, + pub kv_min_tokens: Option, + pub kv_cold_max_tokens: Option, + pub kv_continued_interval_tokens: Option, } impl Default for AppPreferences { @@ -103,6 +107,10 @@ impl Default for AppPreferences { sidebar_collapsed: false, last_project_id: None, sidebar_width: 276, + kv_budget_gib: None, + kv_min_tokens: None, + kv_cold_max_tokens: None, + kv_continued_interval_tokens: None, } } } @@ -207,6 +215,28 @@ impl AppPreferences { .map_err(|_| "Saved simulated memory is invalid.")?, expert_profile_path: self.expert_profile_path.clone(), }, + kv_cache: KvCachePreferences { + budget_gib: self + .kv_budget_gib + .map(u64::try_from) + .transpose() + .map_err(|_| "Saved KV cache budget is invalid.")?, + min_tokens: self + .kv_min_tokens + .map(u32::try_from) + .transpose() + .map_err(|_| "Saved KV cache minimum tokens is invalid.")?, + cold_max_tokens: self + .kv_cold_max_tokens + .map(u32::try_from) + .transpose() + .map_err(|_| "Saved KV cache cold maximum is invalid.")?, + continued_interval_tokens: self + .kv_continued_interval_tokens + .map(u32::try_from) + .transpose() + .map_err(|_| "Saved KV cache continued interval is invalid.")?, + }, }) } } @@ -250,6 +280,10 @@ struct PreferenceChanges<'a> { endpoint_port: i32, endpoint_enabled: bool, endpoint_cors: bool, + kv_budget_gib: Option, + kv_min_tokens: Option, + kv_cold_max_tokens: Option, + kv_continued_interval_tokens: Option, } #[derive(Clone, Debug, Identifiable, Queryable, Selectable)] @@ -534,6 +568,18 @@ impl Database { endpoint_port, endpoint_enabled, endpoint_cors, + kv_budget_gib: runtime + .kv_cache + .budget_gib + .map(i64::try_from) + .transpose() + .map_err(|_| "KV cache budget is too large.")?, + kv_min_tokens: runtime.kv_cache.min_tokens.map(|value| value as i32), + kv_cold_max_tokens: runtime.kv_cache.cold_max_tokens.map(|value| value as i32), + kv_continued_interval_tokens: runtime + .kv_cache + .continued_interval_tokens + .map(|value| value as i32), }) .returning(AppPreferences::as_returning()) .get_result(&mut self.connection) @@ -832,6 +878,12 @@ mod tests { simulated_used_memory_gib: Some(8), expert_profile_path: Some("/tmp/experts.json".into()), }, + kv_cache: KvCachePreferences { + budget_gib: Some(16), + min_tokens: Some(256), + cold_max_tokens: Some(0), + continued_interval_tokens: Some(4_096), + }, ..RuntimePreferences::default() }; database diff --git a/src/engine.rs b/src/engine.rs index 1c2c8bb..9ef6b9a 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -313,6 +313,9 @@ impl Model { pub(crate) struct Generator { executor: metal::Executor, checkpoint: Option, + /// Token frontier of the last transient store, so continued checkpoints are + /// spaced like ds4's `continued_last_store_tokens`. + last_store_tokens: u32, metrics: Arc, } @@ -358,6 +361,7 @@ impl Generator { Ok(Self { executor, checkpoint: None, + last_store_tokens: 0, metrics, }) } @@ -414,7 +418,7 @@ impl Generator { let history = messages .split_last() .map_or(messages, |(_, history)| history); - let store = KvStore::open(directory)?; + let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?; let history_key = conversation_key(&settings.system_prompt, settings.reasoning_mode, history); let history_tag: [u8; 32] = Sha256::digest(&history_key).into(); @@ -423,14 +427,17 @@ impl Generator { if let Some(entry) = store.find(&history_key, self.executor.context()) { if self.select_checkpoint(&entry.checkpoint, entry.tag)? { store.touch(&entry)?; + self.last_store_tokens = entry.tokens; previous_checkpoint = Some(entry.checkpoint); } else { store.discard(&entry); + self.last_store_tokens = 0; previous_checkpoint = None; } } else { self.executor.reset()?; self.checkpoint = None; + self.last_store_tokens = 0; previous_checkpoint = None; self.metrics.kv_lookup(KvLookup::Miss); } @@ -450,6 +457,17 @@ impl Generator { let completed_key = conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed); let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into(); + if !settings.kv_cache.stores( + self.executor.position(), + history.is_empty(), + self.last_store_tokens, + ) { + // A gated store still leaves the finished conversation in the + // live KV, so mark it and drop the stale file association. + self.executor.note_checkpoint_tag(completed_tag); + self.checkpoint = None; + return Ok(output); + } let completed_checkpoint = store.checkpoint_path(&completed_key); self.save_checkpoint(&completed_checkpoint, completed_tag)?; let retained = store.record( @@ -468,6 +486,9 @@ impl Generator { .then(|| std::fs::metadata(&completed_checkpoint).ok()) .flatten() .map_or(0, |item| item.len()); + if retained { + self.last_store_tokens = self.executor.position(); + } self.checkpoint = retained.then_some(completed_checkpoint); } Ok(output) diff --git a/src/engine/kvstore.rs b/src/engine/kvstore.rs index 104cc40..07dd9de 100644 --- a/src/engine/kvstore.rs +++ b/src/engine/kvstore.rs @@ -7,7 +7,6 @@ use std::time::{SystemTime, UNIX_EPOCH}; const INDEX_MAGIC: &[u8; 8] = b"DS4KVIDX"; const INDEX_VERSION: u32 = 1; -const DEFAULT_BUDGET_BYTES: u64 = 4 * 1024 * 1024 * 1024; const MAX_KEY_BYTES: u64 = 256 * 1024 * 1024; const HIT_HALF_LIFE_SECONDS: f64 = 6.0 * 60.0 * 60.0; @@ -58,23 +57,19 @@ pub(super) struct KvStore { } impl KvStore { - pub(super) fn open(directory: &Path) -> Result { + pub(super) fn open(directory: &Path, budget_bytes: u64) -> Result { fs::create_dir_all(directory).map_err(|error| { format!( "Could not create transient KV cache directory {}: {error}", directory.display() ) })?; - Ok(Self { + let store = Self { directory: directory.to_owned(), - budget_bytes: DEFAULT_BUDGET_BYTES, - }) - } - - #[cfg(test)] - fn with_budget(directory: &Path, budget_bytes: u64) -> Result { - let mut store = Self::open(directory)?; - store.budget_bytes = budget_bytes; + budget_bytes, + }; + // A lowered budget takes effect at once, like ds4's eviction on open. + store.evict(None); Ok(store) } @@ -316,6 +311,7 @@ fn read_u64(file: &mut File) -> Result { #[cfg(test)] mod tests { use super::*; + use crate::settings::GIB; fn temporary_directory(name: &str) -> PathBuf { let path = std::env::temp_dir().join(format!( @@ -348,7 +344,7 @@ mod tests { #[test] fn longest_compatible_prefix_wins_and_hits_survive_reopen() { let directory = temporary_directory("prefix"); - let store = KvStore::open(&directory).unwrap(); + let store = KvStore::open(&directory, GIB).unwrap(); add(&store, b"system", 100, 32); let expected = add(&store, b"system-user-one", 200, 32); @@ -356,7 +352,7 @@ mod tests { assert_eq!(found.checkpoint, expected.checkpoint); store.touch(&found).unwrap(); assert_eq!( - KvStore::open(&directory) + KvStore::open(&directory, GIB) .unwrap() .find(b"system-user-one-more", 4096) .unwrap() @@ -371,7 +367,7 @@ mod tests { #[test] fn budget_evicts_the_least_valuable_unprotected_entry() { let directory = temporary_directory("evict"); - let store = KvStore::with_budget(&directory, 300).unwrap(); + let store = KvStore::open(&directory, 300).unwrap(); let first = add(&store, b"first", 10, 32); let second = add(&store, b"second", 100, 32); let third = add(&store, b"third", 200, 32); diff --git a/src/engine/metal.rs b/src/engine/metal.rs index 8c6f34d..0808805 100644 --- a/src/engine/metal.rs +++ b/src/engine/metal.rs @@ -801,6 +801,12 @@ impl Executor { self.checkpoint_tag } + /// Marks what the live KV state represents when no checkpoint file is + /// written for it, so the next turn can still continue in memory. + pub(super) fn note_checkpoint_tag(&mut self, tag: [u8; 32]) { + self.checkpoint_tag = tag; + } + fn encode_token(&mut self, token: u32) -> Result<(), String> { let shape = self.model.shape; let map = self.model.main.map_ptr().cast(); diff --git a/src/metrics.rs b/src/metrics.rs index 2fdfb65..c94a8a4 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -1,8 +1,8 @@ use crate::model::ModelChoice; use std::fs; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU16, AtomicU32, AtomicU64, Ordering}; -use std::time::{Duration, Instant}; +use std::time::{Duration, Instant, SystemTime}; #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] #[repr(u8)] @@ -584,6 +584,22 @@ impl Metrics { ) } + /// Re-reads the cache directories after files are removed outside the + /// runtime, so the counters do not drift away from the disc. + pub(crate) fn rescan_cache(&self, root: &Path) { + let usage = cache_usage(root); + self.kv_files.store(usage.files, Ordering::Relaxed); + self.kv_bytes.store(usage.bytes, Ordering::Relaxed); + self.local_kv_files + .store(usage.local_files, Ordering::Relaxed); + self.local_kv_bytes + .store(usage.local_bytes, Ordering::Relaxed); + self.http_kv_files + .store(usage.http_files, Ordering::Relaxed); + self.http_kv_bytes + .store(usage.http_bytes, Ordering::Relaxed); + } + fn record_checkpoint(&self, source: WorkSource, previous_bytes: Option, bytes: u64) { if bytes == 0 { return; @@ -661,6 +677,122 @@ fn model_name(value: u8) -> &'static str { } } +/// Age buckets for the disc usage explorer. The six hour boundary is one +/// eviction hit half-life, so the distribution shows what is about to lose its +/// reuse credit rather than just what is old. +pub(crate) const CACHE_AGE_BUCKETS: [(&str, u64); 5] = [ + ("Under 1h", 3_600), + ("1–6h", 6 * 3_600), + ("6–24h", 24 * 3_600), + ("1–7d", 7 * 86_400), + ("Older", u64::MAX), +]; +/// Rows shown in the explorer, oldest first — the order eviction works through. +const CACHE_ENTRY_ROWS: usize = 12; + +#[derive(Clone, Debug)] +pub(crate) struct KvCacheEntry { + pub(crate) path: PathBuf, + pub(crate) name: String, + pub(crate) bytes: u64, + pub(crate) age_seconds: u64, +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct KvCacheReport { + pub(crate) budget_bytes: u64, + pub(crate) session_bytes: u64, + pub(crate) session_files: u64, + pub(crate) transient_bytes: u64, + pub(crate) transient_files: u64, + /// Bytes per [`CACHE_AGE_BUCKETS`] entry, both buckets combined. + pub(crate) age_bytes: [u64; CACHE_AGE_BUCKETS.len()], + pub(crate) entries: Vec, +} + +impl KvCacheReport { + pub(crate) fn total_bytes(&self) -> u64 { + self.session_bytes + self.transient_bytes + } +} + +/// Scans the KV cache directories for the stats explorer. The budget applies to +/// the transient store only; session checkpoints are user-owned and are freed +/// with their session. +pub(crate) fn kv_cache_report(root: &Path, budget_bytes: u64) -> KvCacheReport { + let mut report = KvCacheReport { + budget_bytes, + ..KvCacheReport::default() + }; + let now = SystemTime::now(); + for (directory, transient) in [(root.to_path_buf(), false), (root.join("http"), true)] { + let Ok(files) = fs::read_dir(&directory) else { + continue; + }; + for file in files.flatten() { + let path = file.path(); + let Ok(metadata) = file.metadata() else { + continue; + }; + if !metadata.is_file() { + continue; + } + let bytes = metadata.len(); + if transient { + report.transient_bytes += bytes; + report.transient_files += 1; + } else { + report.session_bytes += bytes; + report.session_files += 1; + } + let age_seconds = metadata + .modified() + .ok() + .and_then(|modified| now.duration_since(modified).ok()) + .map_or(0, |age| age.as_secs()); + let bucket = CACHE_AGE_BUCKETS + .iter() + .position(|(_, limit)| age_seconds < *limit) + .unwrap_or(CACHE_AGE_BUCKETS.len() - 1); + report.age_bytes[bucket] += bytes; + // Index files are counted in the totals but folded into their + // checkpoint's row. + if path.extension().is_some_and(|value| value == "meta") { + continue; + } + report.entries.push(KvCacheEntry { + name: entry_name(&path, transient), + bytes: bytes.saturating_add( + fs::metadata(path.with_extension("meta")).map_or(0, |index| index.len()), + ), + path, + age_seconds, + }); + } + } + report.entries.sort_by(|left, right| { + right + .age_seconds + .cmp(&left.age_seconds) + .then_with(|| right.bytes.cmp(&left.bytes)) + .then_with(|| left.name.cmp(&right.name)) + }); + report.entries.truncate(CACHE_ENTRY_ROWS); + report +} + +fn entry_name(path: &Path, transient: bool) -> String { + let stem = path + .file_stem() + .and_then(|value| value.to_str()) + .unwrap_or("unknown"); + if transient { + format!("Transient {}", &stem[..stem.len().min(12)]) + } else { + format!("Session {stem}") + } +} + #[derive(Default)] struct CacheUsage { files: u64, @@ -752,4 +884,47 @@ mod tests { assert_eq!(snapshot.local_kv_bytes, 4_096); assert_eq!(snapshot.http_model_requests, 1); } + + #[test] + fn cache_report_splits_buckets_and_folds_index_files_into_their_checkpoint() { + let root = std::env::temp_dir().join(format!( + "ds4-server-cache-report-{}-{:?}", + std::process::id(), + Instant::now() + )); + fs::create_dir_all(root.join("http")).unwrap(); + fs::write(root.join("7.bin"), vec![0; 300]).unwrap(); + fs::write(root.join("http/abc.bin"), vec![0; 100]).unwrap(); + fs::write(root.join("http/abc.meta"), vec![0; 40]).unwrap(); + // The largest file is also the oldest, so size cannot pass for age. + let aged = fs::File::options() + .write(true) + .open(root.join("http/abc.bin")) + .unwrap(); + aged.set_modified(SystemTime::now() - Duration::from_secs(2 * 86_400)) + .unwrap(); + + let report = kv_cache_report(&root, 4 * 1024); + assert_eq!(report.session_bytes, 300); + assert_eq!(report.session_files, 1); + assert_eq!(report.transient_bytes, 140); + assert_eq!(report.transient_files, 2); + // Fresh files in the first bucket, the two-day-old one in "1–7d". + assert_eq!(report.age_bytes[0], 340); + assert_eq!(report.age_bytes[3], 100); + assert_eq!( + ( + report.age_bytes[1], + report.age_bytes[2], + report.age_bytes[4] + ), + (0, 0, 0) + ); + // Oldest first, and the index file rides along with its checkpoint. + assert_eq!(report.entries.len(), 2); + assert_eq!(report.entries[0].bytes, 140); + assert_eq!(report.entries[1].name, "Session 7"); + + fs::remove_dir_all(root).unwrap(); + } } diff --git a/src/schema.rs b/src/schema.rs index bdeb944..c3f1623 100644 --- a/src/schema.rs +++ b/src/schema.rs @@ -61,6 +61,10 @@ diesel::table! { sidebar_collapsed -> Bool, last_project_id -> Nullable, sidebar_width -> Integer, + kv_budget_gib -> Nullable, + kv_min_tokens -> Nullable, + kv_cold_max_tokens -> Nullable, + kv_continued_interval_tokens -> Nullable, } } diff --git a/src/settings.rs b/src/settings.rs index 82eb644..4e395a5 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -11,6 +11,11 @@ const THINK_MAX_MIN_CONTEXT: i32 = 393_216; const MAX_CPU_THREADS: u32 = 32; const MAX_MTP_DRAFT_TOKENS: i32 = 16; pub(crate) const GIB: u64 = 1024 * 1024 * 1024; +/// DS4 disk KV cache defaults, from `ds4_kvstore.h` and `--kv-disk-space-mb`. +pub(crate) const DEFAULT_KV_BUDGET_GIB: u64 = 4; +const DEFAULT_KV_MIN_TOKENS: u32 = 512; +const DEFAULT_KV_COLD_MAX_TOKENS: u32 = 30_000; +const DEFAULT_KV_CONTINUED_INTERVAL_TOKENS: u32 = 10_000; #[derive(Clone, Debug, PartialEq)] pub(crate) struct SpeculativePreferences { @@ -159,6 +164,61 @@ pub(crate) struct EngineSsdSettings { pub(crate) preload_experts: u32, } +/// Disk KV cache management, mirroring ds4's `--kv-disk-space-mb` budget and +/// its store-side gates. Unset values keep the DS4 defaults. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub(crate) struct KvCachePreferences { + pub(crate) budget_gib: Option, + pub(crate) min_tokens: Option, + pub(crate) cold_max_tokens: Option, + pub(crate) continued_interval_tokens: Option, +} + +impl KvCachePreferences { + pub(crate) fn validate(&self) -> Result<(), String> { + if let Some(gib) = self.budget_gib { + validate_gib("KV cache budget", gib)?; + } + Ok(()) + } + + pub(crate) fn settings(&self) -> KvCacheSettings { + KvCacheSettings { + budget_bytes: self.budget_gib.unwrap_or(DEFAULT_KV_BUDGET_GIB) * GIB, + min_tokens: self.min_tokens.unwrap_or(DEFAULT_KV_MIN_TOKENS), + cold_max_tokens: self.cold_max_tokens.unwrap_or(DEFAULT_KV_COLD_MAX_TOKENS), + continued_interval_tokens: self + .continued_interval_tokens + .unwrap_or(DEFAULT_KV_CONTINUED_INTERVAL_TOKENS), + } + } +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct KvCacheSettings { + pub(crate) budget_bytes: u64, + pub(crate) min_tokens: u32, + /// Largest cold (first-prompt) checkpoint that may be stored. 0 stores none. + pub(crate) cold_max_tokens: u32, + /// Token distance between continued checkpoints. 0 stores none. + pub(crate) continued_interval_tokens: u32, +} + +impl KvCacheSettings { + /// ds4's store gate: too short, too large a cold prompt, or too close to the + /// previous continued frontier all skip the store. + pub(crate) fn stores(&self, tokens: u32, cold: bool, last_store_tokens: u32) -> bool { + if tokens < self.min_tokens { + return false; + } + if cold { + return self.cold_max_tokens > 0 && tokens <= self.cold_max_tokens; + } + self.continued_interval_tokens > 0 + && tokens >= last_store_tokens.saturating_add(self.continued_interval_tokens) + } +} + #[derive(Clone, Debug, Default, PartialEq)] pub(crate) struct SteeringPreferences { pub(crate) file: Option, @@ -245,6 +305,7 @@ pub(crate) struct RuntimePreferences { pub(crate) ssd: SsdPreferences, pub(crate) steering: SteeringPreferences, pub(crate) diagnostics: DiagnosticPreferences, + pub(crate) kv_cache: KvCachePreferences, } impl RuntimePreferences { @@ -254,6 +315,7 @@ impl RuntimePreferences { self.ssd.validate(model)?; self.steering.validate(model)?; self.diagnostics.validate()?; + self.kv_cache.validate()?; if self.ssd.enabled && self.speculative.dspark_enabled { return Err("SSD streaming is not compatible with DSpark support.".into()); } @@ -434,9 +496,14 @@ impl GenerationPreferences { Ok(()) } - pub(crate) fn turn_settings(&self, model: ModelChoice) -> TurnSettings { + pub(crate) fn turn_settings( + &self, + model: ModelChoice, + kv_cache: KvCacheSettings, + ) -> TurnSettings { let glm = model == ModelChoice::Glm52; TurnSettings { + kv_cache, context_tokens: self.context_tokens, max_generated_tokens: self.max_generated_tokens, system_prompt: self.system_prompt.clone(), @@ -485,6 +552,9 @@ fn validate_gib(name: &str, gib: u64) -> Result<(), String> { #[derive(Clone, Debug, PartialEq)] pub(crate) struct TurnSettings { + /// Disc cache policy for this request. Kept out of [`EngineSettings`] so + /// changing a cache knob never reloads the model. + pub(crate) kv_cache: KvCacheSettings, pub(crate) context_tokens: i32, pub(crate) max_generated_tokens: i32, pub(crate) system_prompt: String, @@ -512,7 +582,7 @@ pub(crate) fn effective_settings( generation.validate()?; Ok(EffectiveSettings { engine: runtime.engine_settings(model, generation.context_tokens, models_path)?, - turn: generation.turn_settings(model), + turn: generation.turn_settings(model, runtime.kv_cache.settings()), }) } @@ -524,13 +594,14 @@ mod tests { #[test] fn effective_settings_preserve_unset_model_defaults() { let defaults = GenerationPreferences::default(); - let deepseek = defaults.turn_settings(ModelChoice::DeepSeekV4Flash); + let cache = KvCachePreferences::default().settings(); + let deepseek = defaults.turn_settings(ModelChoice::DeepSeekV4Flash, cache); assert_eq!( (deepseek.temperature, deepseek.top_p, deepseek.min_p), (1.0, 1.0, 0.05) ); - let glm = defaults.turn_settings(ModelChoice::Glm52); + let glm = defaults.turn_settings(ModelChoice::Glm52, cache); assert_eq!((glm.temperature, glm.top_p, glm.min_p), (1.0, 0.95, 0.0)); let explicit = GenerationPreferences { @@ -539,7 +610,7 @@ mod tests { reasoning_mode: ReasoningMode::Max, ..defaults }; - let effective = explicit.turn_settings(ModelChoice::Glm52); + let effective = explicit.turn_settings(ModelChoice::Glm52, cache); assert_eq!((effective.top_p, effective.min_p), (0.4, 0.2)); assert_eq!(effective.reasoning_mode, ReasoningMode::High); } @@ -595,6 +666,39 @@ mod tests { assert!(glm.validate(ModelChoice::DeepSeekV4Pro).is_err()); } + #[test] + fn kv_cache_gates_follow_the_ds4_store_rules() { + let defaults = KvCachePreferences::default().settings(); + assert_eq!(defaults.budget_bytes, DEFAULT_KV_BUDGET_GIB * GIB); + + // Too short to be worth a checkpoint at all. + assert!(!defaults.stores(511, true, 0)); + assert!(defaults.stores(512, true, 0)); + // Cold prompts above the cold maximum are skipped, continued ones are + // spaced by the interval regardless of length. + assert!(!defaults.stores(30_001, true, 0)); + assert!(defaults.stores(30_001, false, 20_000)); + assert!(!defaults.stores(29_999, false, 20_000)); + + let off = KvCachePreferences { + cold_max_tokens: Some(0), + continued_interval_tokens: Some(0), + ..KvCachePreferences::default() + } + .settings(); + assert!(!off.stores(1_000, true, 0)); + assert!(!off.stores(1_000_000, false, 0)); + + assert!( + KvCachePreferences { + budget_gib: Some(0), + ..KvCachePreferences::default() + } + .validate() + .is_err() + ); + } + #[test] fn runtime_settings_preserve_automatic_values_and_cross_field_rules() { let runtime = RuntimePreferences { @@ -684,6 +788,10 @@ mod tests { "--expert-profile", "--glm-mtp", "--glm-mtp-timing", + "--kv-cache-cold-max-tokens", + "--kv-cache-continued-interval-tokens", + "--kv-cache-min-tokens", + "--kv-disk-space-mb", "--min-p", "--model", "--mtp", @@ -735,6 +843,15 @@ mod tests { "--imatrix-max-tokens", "--imatrix-out", "--inspect", + // Token-boundary trimming belongs to ds4's token-prefix lookup; our + // store keys on the rendered conversation instead. + "--kv-cache-boundary-align-tokens", + "--kv-cache-boundary-trim-tokens", + // Checkpoints already bind the model file identity, so a different + // quantization can never load. + "--kv-cache-reject-different-quant", + // The cache directory is fixed inside the application container. + "--kv-disk-dir", "--layers", "--listen", "--logprobs-top-k",