Manage the KV cache disc usage from preferences and stats

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
Georg Bauer
2026-07-26 09:10:18 +02:00
parent b182e7007c
commit 63e2a74ad5
14 changed files with 794 additions and 35 deletions

View File

@@ -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;

View File

@@ -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;

View File

@@ -15,21 +15,21 @@ use preferences::{parse_optional_gib, parse_streaming_cache};
use crate::database::{AppPreferences, Database, ProjectWithSessions, SessionState, StoredMessage}; use crate::database::{AppPreferences, Database, ProjectWithSessions, SessionState, StoredMessage};
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
use crate::engine::ChatTurn; use crate::engine::ChatTurn;
use crate::metrics::{Metrics, MetricsSnapshot}; use crate::metrics::{KvCacheReport, Metrics, MetricsSnapshot};
use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice}; use crate::model::{self, DownloadOutcome, DownloadProgress, ManagedArtifactId, ModelChoice};
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
use crate::runtime::{ActiveGeneration, CheckpointTarget, GenerationEvent, GenerationService}; use crate::runtime::{ActiveGeneration, CheckpointTarget, GenerationEvent, GenerationService};
use crate::settings::{ use crate::settings::{
DiagnosticPreferences, ExecutionPreferences, GIB, GenerationPreferences, ReasoningMode, DEFAULT_KV_BUDGET_GIB, DiagnosticPreferences, ExecutionPreferences, GIB, GenerationPreferences,
RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, KvCachePreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences,
StreamingCacheBudget, SteeringPreferences, StreamingCacheBudget,
}; };
use iced::widget::{markdown, scrollable, text_input}; use iced::widget::{markdown, scrollable, text_input};
use iced::{Size, Subscription, Task, keyboard, mouse, window}; use iced::{Size, Subscription, Task, keyboard, mouse, window};
use rfd::AsyncFileDialog; use rfd::AsyncFileDialog;
use std::collections::{HashMap, HashSet, VecDeque}; use std::collections::{HashMap, HashSet, VecDeque};
use std::fs; use std::fs;
use std::path::PathBuf; use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::mpsc::{self, TryRecvError}; use std::sync::mpsc::{self, TryRecvError};
use std::sync::{Arc, Mutex, RwLock}; use std::sync::{Arc, Mutex, RwLock};
@@ -38,6 +38,8 @@ use std::time::{Duration, Instant};
const APP_ID: &str = "de.rfc1437.ds4server"; const APP_ID: &str = "de.rfc1437.ds4server";
const METRICS_SAMPLE_INTERVAL: Duration = Duration::from_millis(200); 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 MIN_SIDEBAR_WIDTH: i32 = 180;
pub(super) const MAX_SIDEBAR_WIDTH: i32 = 520; pub(super) const MAX_SIDEBAR_WIDTH: i32 = 520;
@@ -82,6 +84,9 @@ pub(crate) struct App {
metrics: Arc<Metrics>, metrics: Arc<Metrics>,
pub(super) metrics_snapshot: MetricsSnapshot, pub(super) metrics_snapshot: MetricsSnapshot,
pub(super) metrics_history: VecDeque<MetricsPoint>, pub(super) metrics_history: VecDeque<MetricsPoint>,
/// 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, last_http_requests: u64,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
generation_service: Option<GenerationService>, generation_service: Option<GenerationService>,
@@ -162,6 +167,10 @@ pub(crate) enum Message {
PreferenceSteeringAttnChanged(String), PreferenceSteeringAttnChanged(String),
PreferenceSimulatedMemoryChanged(String), PreferenceSimulatedMemoryChanged(String),
PreferenceExpertProfileChanged(String), PreferenceExpertProfileChanged(String),
PreferenceKvBudgetChanged(String),
PreferenceKvMinTokensChanged(String),
PreferenceKvColdMaxChanged(String),
PreferenceKvContinuedIntervalChanged(String),
ResetPreferences, ResetPreferences,
SavePreferences, SavePreferences,
DownloadArtifact(ManagedArtifactId), DownloadArtifact(ManagedArtifactId),
@@ -202,6 +211,8 @@ pub(crate) enum Message {
DragSidebar(f32), DragSidebar(f32),
EndSidebarDrag, EndSidebarDrag,
MetricsTick, MetricsTick,
DiscardCacheEntry(PathBuf),
ClearTransientCache,
} }
impl App { impl App {
@@ -214,6 +225,9 @@ impl App {
Ok(draft) => draft, Ok(draft) => draft,
Err(error) => return Self::failed(error, main_window), 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; let context_limit = preferences.context_tokens.max(0) as u32;
// Reopen on the project we left, with a fresh draft chat. // Reopen on the project we left, with a fresh draft chat.
let last_project = preferences let last_project = preferences
@@ -264,6 +278,8 @@ impl App {
metrics, metrics,
metrics_snapshot, metrics_snapshot,
metrics_history: VecDeque::with_capacity(120), metrics_history: VecDeque::with_capacity(120),
kv_cache_report: KvCacheReport::default(),
last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL,
last_http_requests: 0, last_http_requests: 0,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
generation_service, generation_service,
@@ -346,6 +362,8 @@ impl App {
metrics, metrics,
metrics_snapshot, metrics_snapshot,
metrics_history: VecDeque::with_capacity(120), metrics_history: VecDeque::with_capacity(120),
kv_cache_report: KvCacheReport::default(),
last_cache_scan: Instant::now() - CACHE_SCAN_INTERVAL,
last_http_requests: 0, last_http_requests: 0,
#[cfg(target_os = "macos")] #[cfg(target_os = "macos")]
generation_service, generation_service,
@@ -422,7 +440,12 @@ impl App {
} }
} }
Message::ShowChat => self.detail_tab = DetailTab::Chat, 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 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. // the width the user is asking for. Only the final width is stored.
Message::StartSidebarDrag => self.sidebar_drag = true, Message::StartSidebarDrag => self.sidebar_drag = true,
@@ -652,6 +675,22 @@ impl App {
self.preference_draft.expert_profile_path = value; self.preference_draft.expert_profile_path = value;
self.preference_error = None; 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 => { Message::ResetPreferences => {
self.preference_draft.reset(); self.preference_draft.reset();
self.preference_error = None; self.preference_error = None;
@@ -801,6 +840,7 @@ impl App {
self.tokens_per_second = None; self.tokens_per_second = None;
} }
self.reload_projects(); self.reload_projects();
self.finish_cache_change();
} }
Err(error) => self.error = Some(error), Err(error) => self.error = Some(error),
} }
@@ -936,6 +976,7 @@ impl App {
match database.delete_session(session_id) { match database.delete_session(session_id) {
Ok(()) => { Ok(()) => {
let _ = fs::remove_file(session_checkpoint_path(session_id)); let _ = fs::remove_file(session_checkpoint_path(session_id));
self.finish_cache_change();
if self.session_menu == Some(session_id) { if self.session_menu == Some(session_id) {
self.session_menu = None; self.session_menu = None;
} }
@@ -1050,6 +1091,56 @@ impl App {
self.metrics_history.pop_front(); self.metrics_history.pop_front();
} }
self.metrics_snapshot = snapshot; 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<Message> {
scrollable::snap_to(chat_scroll_id(), scrollable::RelativeOffset::END) 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::<HashSet<_>>();
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::<i32>().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 { fn session_checkpoint_path(session_id: i32) -> PathBuf {
application_support_path() kv_cache_path().join(format!("{session_id}.bin"))
.join("kv-cache")
.join(format!("{session_id}.bin"))
} }
/// Content-addressed KV cache for turns that belong to no stored session: the /// 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. /// HTTP endpoint and the app's own one-shot requests share it.
fn transient_cache_path() -> PathBuf { 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 { pub(crate) fn app_icon() -> window::Icon {

View File

@@ -37,6 +37,10 @@ pub(super) struct PreferenceDraft {
pub(super) directional_steering_attn: String, pub(super) directional_steering_attn: String,
pub(super) simulated_used_memory_gib: String, pub(super) simulated_used_memory_gib: String,
pub(super) expert_profile_path: 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 { impl PreferenceDraft {
@@ -121,6 +125,22 @@ impl PreferenceDraft {
.simulated_used_memory_gib .simulated_used_memory_gib
.map_or_else(String::new, |value| value.to_string()), .map_or_else(String::new, |value| value.to_string()),
expert_profile_path: runtime.diagnostics.expert_profile_path.unwrap_or_default(), 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.directional_steering_attn.clear();
self.simulated_used_memory_gib.clear(); self.simulated_used_memory_gib.clear();
self.expert_profile_path.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<ExecutionPreferences, String> { pub(super) fn execution(&self) -> Result<ExecutionPreferences, String> {
@@ -241,6 +265,18 @@ impl PreferenceDraft {
)?, )?,
expert_profile_path: optional_text(&self.expert_profile_path), 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,
)?,
},
}) })
} }
} }

View File

@@ -132,7 +132,12 @@ impl App {
pub(super) fn reload_projects(&mut self) { pub(super) fn reload_projects(&mut self) {
if let Some(database) = &mut self.database { if let Some(database) = &mut self.database {
match database.load_projects() { 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), Err(error) => self.error = Some(error),
} }
} }
@@ -190,6 +195,36 @@ mod tests {
assert_eq!(draft_title(&projects, 99), "Session 1"); 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 { fn project(id: i32, name: &str) -> Project {
Project { Project {
id, id,

View File

@@ -479,12 +479,63 @@ impl App {
] ]
.spacing(10), .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![ let mut fields = column![
model_group, model_group,
endpoint_group, endpoint_group,
generation_group, generation_group,
execution_group, execution_group,
acceleration_group, acceleration_group,
kv_cache_group,
steering_group, steering_group,
] ]
.spacing(12); .spacing(12);

View File

@@ -319,6 +319,7 @@ impl App {
.spacing(9) .spacing(9)
.into(), .into(),
); );
let disc = stats_panel("KV CACHE DISC USAGE", self.cache_explorer());
let server = stats_panel( let server = stats_panel(
"LOCAL SERVER", "LOCAL SERVER",
column![ column![
@@ -353,6 +354,7 @@ impl App {
throughput, throughput,
kv_io, kv_io,
requests, requests,
disc,
row![model, runtime].spacing(10), row![model, runtime].spacing(10),
row![cache, server].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.") 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) .height(Length::Fill)
.into() .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)
} }

View File

@@ -6,8 +6,8 @@ use std::path::Path;
use crate::schema::{messages, preferences, projects, sessions}; use crate::schema::{messages, preferences, projects, sessions};
use crate::settings::{ use crate::settings::{
DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, ReasoningMode, DiagnosticPreferences, ExecutionPreferences, GenerationPreferences, KvCachePreferences,
RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences, ReasoningMode, RuntimePreferences, SpeculativePreferences, SsdPreferences, SteeringPreferences,
StreamingCacheBudget, StreamingCacheBudget,
}; };
@@ -58,6 +58,10 @@ pub struct AppPreferences {
/// Project the app reopens on. Cleared when that project goes away. /// Project the app reopens on. Cleared when that project goes away.
pub last_project_id: Option<i32>, pub last_project_id: Option<i32>,
pub sidebar_width: i32, pub sidebar_width: i32,
pub kv_budget_gib: Option<i64>,
pub kv_min_tokens: Option<i32>,
pub kv_cold_max_tokens: Option<i32>,
pub kv_continued_interval_tokens: Option<i32>,
} }
impl Default for AppPreferences { impl Default for AppPreferences {
@@ -103,6 +107,10 @@ impl Default for AppPreferences {
sidebar_collapsed: false, sidebar_collapsed: false,
last_project_id: None, last_project_id: None,
sidebar_width: 276, 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.")?, .map_err(|_| "Saved simulated memory is invalid.")?,
expert_profile_path: self.expert_profile_path.clone(), 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_port: i32,
endpoint_enabled: bool, endpoint_enabled: bool,
endpoint_cors: bool, endpoint_cors: bool,
kv_budget_gib: Option<i64>,
kv_min_tokens: Option<i32>,
kv_cold_max_tokens: Option<i32>,
kv_continued_interval_tokens: Option<i32>,
} }
#[derive(Clone, Debug, Identifiable, Queryable, Selectable)] #[derive(Clone, Debug, Identifiable, Queryable, Selectable)]
@@ -534,6 +568,18 @@ impl Database {
endpoint_port, endpoint_port,
endpoint_enabled, endpoint_enabled,
endpoint_cors, 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()) .returning(AppPreferences::as_returning())
.get_result(&mut self.connection) .get_result(&mut self.connection)
@@ -832,6 +878,12 @@ mod tests {
simulated_used_memory_gib: Some(8), simulated_used_memory_gib: Some(8),
expert_profile_path: Some("/tmp/experts.json".into()), 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() ..RuntimePreferences::default()
}; };
database database

View File

@@ -313,6 +313,9 @@ impl Model {
pub(crate) struct Generator { pub(crate) struct Generator {
executor: metal::Executor, executor: metal::Executor,
checkpoint: Option<PathBuf>, checkpoint: Option<PathBuf>,
/// 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<Metrics>, metrics: Arc<Metrics>,
} }
@@ -358,6 +361,7 @@ impl Generator {
Ok(Self { Ok(Self {
executor, executor,
checkpoint: None, checkpoint: None,
last_store_tokens: 0,
metrics, metrics,
}) })
} }
@@ -414,7 +418,7 @@ impl Generator {
let history = messages let history = messages
.split_last() .split_last()
.map_or(messages, |(_, history)| history); .map_or(messages, |(_, history)| history);
let store = KvStore::open(directory)?; let store = KvStore::open(directory, settings.kv_cache.budget_bytes)?;
let history_key = let history_key =
conversation_key(&settings.system_prompt, settings.reasoning_mode, history); conversation_key(&settings.system_prompt, settings.reasoning_mode, history);
let history_tag: [u8; 32] = Sha256::digest(&history_key).into(); 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 let Some(entry) = store.find(&history_key, self.executor.context()) {
if self.select_checkpoint(&entry.checkpoint, entry.tag)? { if self.select_checkpoint(&entry.checkpoint, entry.tag)? {
store.touch(&entry)?; store.touch(&entry)?;
self.last_store_tokens = entry.tokens;
previous_checkpoint = Some(entry.checkpoint); previous_checkpoint = Some(entry.checkpoint);
} else { } else {
store.discard(&entry); store.discard(&entry);
self.last_store_tokens = 0;
previous_checkpoint = None; previous_checkpoint = None;
} }
} else { } else {
self.executor.reset()?; self.executor.reset()?;
self.checkpoint = None; self.checkpoint = None;
self.last_store_tokens = 0;
previous_checkpoint = None; previous_checkpoint = None;
self.metrics.kv_lookup(KvLookup::Miss); self.metrics.kv_lookup(KvLookup::Miss);
} }
@@ -450,6 +457,17 @@ impl Generator {
let completed_key = let completed_key =
conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed); conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed);
let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into(); 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); let completed_checkpoint = store.checkpoint_path(&completed_key);
self.save_checkpoint(&completed_checkpoint, completed_tag)?; self.save_checkpoint(&completed_checkpoint, completed_tag)?;
let retained = store.record( let retained = store.record(
@@ -468,6 +486,9 @@ impl Generator {
.then(|| std::fs::metadata(&completed_checkpoint).ok()) .then(|| std::fs::metadata(&completed_checkpoint).ok())
.flatten() .flatten()
.map_or(0, |item| item.len()); .map_or(0, |item| item.len());
if retained {
self.last_store_tokens = self.executor.position();
}
self.checkpoint = retained.then_some(completed_checkpoint); self.checkpoint = retained.then_some(completed_checkpoint);
} }
Ok(output) Ok(output)

View File

@@ -7,7 +7,6 @@ use std::time::{SystemTime, UNIX_EPOCH};
const INDEX_MAGIC: &[u8; 8] = b"DS4KVIDX"; const INDEX_MAGIC: &[u8; 8] = b"DS4KVIDX";
const INDEX_VERSION: u32 = 1; const INDEX_VERSION: u32 = 1;
const DEFAULT_BUDGET_BYTES: u64 = 4 * 1024 * 1024 * 1024;
const MAX_KEY_BYTES: u64 = 256 * 1024 * 1024; const MAX_KEY_BYTES: u64 = 256 * 1024 * 1024;
const HIT_HALF_LIFE_SECONDS: f64 = 6.0 * 60.0 * 60.0; const HIT_HALF_LIFE_SECONDS: f64 = 6.0 * 60.0 * 60.0;
@@ -58,23 +57,19 @@ pub(super) struct KvStore {
} }
impl KvStore { impl KvStore {
pub(super) fn open(directory: &Path) -> Result<Self, String> { pub(super) fn open(directory: &Path, budget_bytes: u64) -> Result<Self, String> {
fs::create_dir_all(directory).map_err(|error| { fs::create_dir_all(directory).map_err(|error| {
format!( format!(
"Could not create transient KV cache directory {}: {error}", "Could not create transient KV cache directory {}: {error}",
directory.display() directory.display()
) )
})?; })?;
Ok(Self { let store = Self {
directory: directory.to_owned(), directory: directory.to_owned(),
budget_bytes: DEFAULT_BUDGET_BYTES, budget_bytes,
}) };
} // A lowered budget takes effect at once, like ds4's eviction on open.
store.evict(None);
#[cfg(test)]
fn with_budget(directory: &Path, budget_bytes: u64) -> Result<Self, String> {
let mut store = Self::open(directory)?;
store.budget_bytes = budget_bytes;
Ok(store) Ok(store)
} }
@@ -316,6 +311,7 @@ fn read_u64(file: &mut File) -> Result<u64, String> {
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
use crate::settings::GIB;
fn temporary_directory(name: &str) -> PathBuf { fn temporary_directory(name: &str) -> PathBuf {
let path = std::env::temp_dir().join(format!( let path = std::env::temp_dir().join(format!(
@@ -348,7 +344,7 @@ mod tests {
#[test] #[test]
fn longest_compatible_prefix_wins_and_hits_survive_reopen() { fn longest_compatible_prefix_wins_and_hits_survive_reopen() {
let directory = temporary_directory("prefix"); let directory = temporary_directory("prefix");
let store = KvStore::open(&directory).unwrap(); let store = KvStore::open(&directory, GIB).unwrap();
add(&store, b"system", 100, 32); add(&store, b"system", 100, 32);
let expected = add(&store, b"system-user-one", 200, 32); let expected = add(&store, b"system-user-one", 200, 32);
@@ -356,7 +352,7 @@ mod tests {
assert_eq!(found.checkpoint, expected.checkpoint); assert_eq!(found.checkpoint, expected.checkpoint);
store.touch(&found).unwrap(); store.touch(&found).unwrap();
assert_eq!( assert_eq!(
KvStore::open(&directory) KvStore::open(&directory, GIB)
.unwrap() .unwrap()
.find(b"system-user-one-more", 4096) .find(b"system-user-one-more", 4096)
.unwrap() .unwrap()
@@ -371,7 +367,7 @@ mod tests {
#[test] #[test]
fn budget_evicts_the_least_valuable_unprotected_entry() { fn budget_evicts_the_least_valuable_unprotected_entry() {
let directory = temporary_directory("evict"); 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 first = add(&store, b"first", 10, 32);
let second = add(&store, b"second", 100, 32); let second = add(&store, b"second", 100, 32);
let third = add(&store, b"third", 200, 32); let third = add(&store, b"third", 200, 32);

View File

@@ -801,6 +801,12 @@ impl Executor {
self.checkpoint_tag 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> { fn encode_token(&mut self, token: u32) -> Result<(), String> {
let shape = self.model.shape; let shape = self.model.shape;
let map = self.model.main.map_ptr().cast(); let map = self.model.main.map_ptr().cast();

View File

@@ -1,8 +1,8 @@
use crate::model::ModelChoice; use crate::model::ModelChoice;
use std::fs; use std::fs;
use std::path::Path; use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU16, AtomicU32, AtomicU64, Ordering}; 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)] #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
#[repr(u8)] #[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<u64>, bytes: u64) { fn record_checkpoint(&self, source: WorkSource, previous_bytes: Option<u64>, bytes: u64) {
if bytes == 0 { if bytes == 0 {
return; 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),
("16h", 6 * 3_600),
("624h", 24 * 3_600),
("17d", 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<KvCacheEntry>,
}
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)] #[derive(Default)]
struct CacheUsage { struct CacheUsage {
files: u64, files: u64,
@@ -752,4 +884,47 @@ mod tests {
assert_eq!(snapshot.local_kv_bytes, 4_096); assert_eq!(snapshot.local_kv_bytes, 4_096);
assert_eq!(snapshot.http_model_requests, 1); 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 "17d".
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();
}
} }

View File

@@ -61,6 +61,10 @@ diesel::table! {
sidebar_collapsed -> Bool, sidebar_collapsed -> Bool,
last_project_id -> Nullable<Integer>, last_project_id -> Nullable<Integer>,
sidebar_width -> Integer, sidebar_width -> Integer,
kv_budget_gib -> Nullable<BigInt>,
kv_min_tokens -> Nullable<Integer>,
kv_cold_max_tokens -> Nullable<Integer>,
kv_continued_interval_tokens -> Nullable<Integer>,
} }
} }

View File

@@ -11,6 +11,11 @@ const THINK_MAX_MIN_CONTEXT: i32 = 393_216;
const MAX_CPU_THREADS: u32 = 32; const MAX_CPU_THREADS: u32 = 32;
const MAX_MTP_DRAFT_TOKENS: i32 = 16; const MAX_MTP_DRAFT_TOKENS: i32 = 16;
pub(crate) const GIB: u64 = 1024 * 1024 * 1024; 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)] #[derive(Clone, Debug, PartialEq)]
pub(crate) struct SpeculativePreferences { pub(crate) struct SpeculativePreferences {
@@ -159,6 +164,61 @@ pub(crate) struct EngineSsdSettings {
pub(crate) preload_experts: u32, 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<u64>,
pub(crate) min_tokens: Option<u32>,
pub(crate) cold_max_tokens: Option<u32>,
pub(crate) continued_interval_tokens: Option<u32>,
}
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)] #[derive(Clone, Debug, Default, PartialEq)]
pub(crate) struct SteeringPreferences { pub(crate) struct SteeringPreferences {
pub(crate) file: Option<String>, pub(crate) file: Option<String>,
@@ -245,6 +305,7 @@ pub(crate) struct RuntimePreferences {
pub(crate) ssd: SsdPreferences, pub(crate) ssd: SsdPreferences,
pub(crate) steering: SteeringPreferences, pub(crate) steering: SteeringPreferences,
pub(crate) diagnostics: DiagnosticPreferences, pub(crate) diagnostics: DiagnosticPreferences,
pub(crate) kv_cache: KvCachePreferences,
} }
impl RuntimePreferences { impl RuntimePreferences {
@@ -254,6 +315,7 @@ impl RuntimePreferences {
self.ssd.validate(model)?; self.ssd.validate(model)?;
self.steering.validate(model)?; self.steering.validate(model)?;
self.diagnostics.validate()?; self.diagnostics.validate()?;
self.kv_cache.validate()?;
if self.ssd.enabled && self.speculative.dspark_enabled { if self.ssd.enabled && self.speculative.dspark_enabled {
return Err("SSD streaming is not compatible with DSpark support.".into()); return Err("SSD streaming is not compatible with DSpark support.".into());
} }
@@ -434,9 +496,14 @@ impl GenerationPreferences {
Ok(()) 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; let glm = model == ModelChoice::Glm52;
TurnSettings { TurnSettings {
kv_cache,
context_tokens: self.context_tokens, context_tokens: self.context_tokens,
max_generated_tokens: self.max_generated_tokens, max_generated_tokens: self.max_generated_tokens,
system_prompt: self.system_prompt.clone(), system_prompt: self.system_prompt.clone(),
@@ -485,6 +552,9 @@ fn validate_gib(name: &str, gib: u64) -> Result<(), String> {
#[derive(Clone, Debug, PartialEq)] #[derive(Clone, Debug, PartialEq)]
pub(crate) struct TurnSettings { 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) context_tokens: i32,
pub(crate) max_generated_tokens: i32, pub(crate) max_generated_tokens: i32,
pub(crate) system_prompt: String, pub(crate) system_prompt: String,
@@ -512,7 +582,7 @@ pub(crate) fn effective_settings(
generation.validate()?; generation.validate()?;
Ok(EffectiveSettings { Ok(EffectiveSettings {
engine: runtime.engine_settings(model, generation.context_tokens, models_path)?, 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] #[test]
fn effective_settings_preserve_unset_model_defaults() { fn effective_settings_preserve_unset_model_defaults() {
let defaults = GenerationPreferences::default(); 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!( assert_eq!(
(deepseek.temperature, deepseek.top_p, deepseek.min_p), (deepseek.temperature, deepseek.top_p, deepseek.min_p),
(1.0, 1.0, 0.05) (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)); assert_eq!((glm.temperature, glm.top_p, glm.min_p), (1.0, 0.95, 0.0));
let explicit = GenerationPreferences { let explicit = GenerationPreferences {
@@ -539,7 +610,7 @@ mod tests {
reasoning_mode: ReasoningMode::Max, reasoning_mode: ReasoningMode::Max,
..defaults ..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.top_p, effective.min_p), (0.4, 0.2));
assert_eq!(effective.reasoning_mode, ReasoningMode::High); assert_eq!(effective.reasoning_mode, ReasoningMode::High);
} }
@@ -595,6 +666,39 @@ mod tests {
assert!(glm.validate(ModelChoice::DeepSeekV4Pro).is_err()); 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] #[test]
fn runtime_settings_preserve_automatic_values_and_cross_field_rules() { fn runtime_settings_preserve_automatic_values_and_cross_field_rules() {
let runtime = RuntimePreferences { let runtime = RuntimePreferences {
@@ -684,6 +788,10 @@ mod tests {
"--expert-profile", "--expert-profile",
"--glm-mtp", "--glm-mtp",
"--glm-mtp-timing", "--glm-mtp-timing",
"--kv-cache-cold-max-tokens",
"--kv-cache-continued-interval-tokens",
"--kv-cache-min-tokens",
"--kv-disk-space-mb",
"--min-p", "--min-p",
"--model", "--model",
"--mtp", "--mtp",
@@ -735,6 +843,15 @@ mod tests {
"--imatrix-max-tokens", "--imatrix-max-tokens",
"--imatrix-out", "--imatrix-out",
"--inspect", "--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", "--layers",
"--listen", "--listen",
"--logprobs-top-k", "--logprobs-top-k",