From a1761fa73170ca482974a5dd3059f99579805c00 Mon Sep 17 00:00:00 2001 From: Georg Bauer Date: Sat, 25 Jul 2026 13:56:10 +0200 Subject: [PATCH] Add prefix-aware KV session cache --- src/engine.rs | 171 +++++++++++-------- src/engine/kvstore.rs | 385 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 484 insertions(+), 72 deletions(-) create mode 100644 src/engine/kvstore.rs diff --git a/src/engine.rs b/src/engine.rs index 7ffdee8..9ac0811 100644 --- a/src/engine.rs +++ b/src/engine.rs @@ -1,5 +1,7 @@ mod gguf; #[cfg(target_os = "macos")] +mod kvstore; +#[cfg(target_os = "macos")] mod metal; mod tokenizer; mod validation; @@ -10,6 +12,8 @@ use crate::model::ModelChoice; use crate::settings::TurnSettings; use crate::settings::{EngineSettings, ReasoningMode}; use gguf::{F16, F32, Gguf, I32, IQ2_XXS, Q2_K, Q4_0, Q4_K, Q5_K, Q6_K, Q8_0, Tensor, Value}; +#[cfg(target_os = "macos")] +use kvstore::{KvStore, StoreReason}; use sha2::{Digest, Sha256}; use std::path::{Path, PathBuf}; #[cfg(target_os = "macos")] @@ -409,15 +413,25 @@ impl Generator { let history = messages .split_last() .map_or(messages, |(_, history)| history); - let history_tag = - conversation_tag(&settings.system_prompt, settings.reasoning_mode, history); - let checkpoint = directory.join(format!("{}.bin", hex_tag(history_tag))); + let store = KvStore::open(directory)?; + let history_key = + conversation_key(&settings.system_prompt, settings.reasoning_mode, history); + let history_tag: [u8; 32] = Sha256::digest(&history_key).into(); + let mut previous_checkpoint = self.checkpoint.clone(); if self.executor.checkpoint_tag() != history_tag { - self.select_checkpoint(&checkpoint, history_tag)?; - if self.executor.checkpoint_tag() != history_tag { + if let Some(entry) = store.find(&history_key, self.executor.context()) { + if self.select_checkpoint(&entry.checkpoint, entry.tag)? { + store.touch(&entry)?; + previous_checkpoint = Some(entry.checkpoint); + } else { + store.discard(&entry); + previous_checkpoint = None; + } + } else { self.executor.reset()?; self.checkpoint = None; - let _ = std::fs::remove_file(&checkpoint); + previous_checkpoint = None; + self.metrics.kv_lookup(KvLookup::Miss); } } else { self.metrics.kv_lookup(KvLookup::MemoryHit); @@ -427,30 +441,34 @@ impl Generator { self.generate_inner(messages, settings, cancelled, &mut emit, &mut progress)?; let mut completed = messages.to_vec(); completed.push(output.message.clone()); - let completed_tag = - conversation_tag(&settings.system_prompt, settings.reasoning_mode, &completed); - std::fs::create_dir_all(directory).map_err(|error| { - format!( - "Could not create transient KV cache directory {}: {error}", - directory.display() - ) - })?; - let completed_checkpoint = directory.join(format!("{}.bin", hex_tag(completed_tag))); - output.previous_checkpoint_bytes = std::fs::metadata(&completed_checkpoint) - .ok() + output.previous_checkpoint_bytes = previous_checkpoint + .as_deref() + .and_then(|path| std::fs::metadata(path).ok()) .map(|item| item.len()); - self.save_checkpoint( - &completed_checkpoint, - if prompt_complete { - completed_tag - } else { - [0; 32] - }, - )?; - output.checkpoint_bytes = std::fs::metadata(&completed_checkpoint) - .map(|item| item.len()) - .unwrap_or(0); - self.checkpoint = Some(completed_checkpoint); + if prompt_complete { + let completed_key = + conversation_key(&settings.system_prompt, settings.reasoning_mode, &completed); + let completed_tag: [u8; 32] = Sha256::digest(&completed_key).into(); + let completed_checkpoint = store.checkpoint_path(&completed_key); + self.save_checkpoint(&completed_checkpoint, completed_tag)?; + let retained = store.record( + &completed_checkpoint, + &completed_key, + completed_tag, + self.executor.position(), + self.executor.context(), + if history.is_empty() { + StoreReason::Cold + } else { + StoreReason::Continued + }, + )?; + output.checkpoint_bytes = retained + .then(|| std::fs::metadata(&completed_checkpoint).ok()) + .flatten() + .map_or(0, |item| item.len()); + self.checkpoint = retained.then_some(completed_checkpoint); + } Ok(output) } @@ -458,15 +476,15 @@ impl Generator { &mut self, checkpoint: &Path, expected_tag: [u8; 32], - ) -> Result<(), String> { + ) -> Result { if self.checkpoint.as_deref() == Some(checkpoint) { - self.metrics - .kv_lookup(if self.executor.checkpoint_tag() == expected_tag { - KvLookup::MemoryHit - } else { - KvLookup::Miss - }); - return Ok(()); + let found = self.executor.checkpoint_tag() == expected_tag; + self.metrics.kv_lookup(if found { + KvLookup::MemoryHit + } else { + KvLookup::Miss + }); + return Ok(found); } self.executor.reset()?; self.metrics.kv_read_started(); @@ -476,8 +494,9 @@ impl Generator { .load_checkpoint(checkpoint, &mut |bytes| self.metrics.kv_read_bytes(bytes)); self.metrics .kv_read_finished(started.elapsed(), loaded.is_err()); + let found = matches!(loaded, Ok(true)) && self.executor.checkpoint_tag() == expected_tag; let lookup = match loaded { - Ok(true) if self.executor.checkpoint_tag() == expected_tag => KvLookup::DiskHit, + Ok(true) if found => KvLookup::DiskHit, Ok(true) | Ok(false) => KvLookup::Miss, Err(_) => { self.executor.reset()?; @@ -487,7 +506,7 @@ impl Generator { }; self.metrics.kv_lookup(lookup); self.checkpoint = Some(checkpoint.to_owned()); - Ok(()) + Ok(found) } fn save_checkpoint(&mut self, checkpoint: &Path, tag: [u8; 32]) -> Result<(), String> { @@ -871,44 +890,38 @@ fn emit_safe_text( } #[cfg(target_os = "macos")] -fn hex_tag(tag: [u8; 32]) -> String { - const HEX: &[u8; 16] = b"0123456789abcdef"; - let mut output = String::with_capacity(64); - for byte in tag { - output.push(HEX[(byte >> 4) as usize] as char); - output.push(HEX[(byte & 0x0f) as usize] as char); +fn conversation_key(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> Vec { + fn text(output: &mut Vec, value: &str) { + output.extend_from_slice(&(value.len() as u64).to_le_bytes()); + output.extend_from_slice(value.as_bytes()); + } + + let mut output = b"DS4Server chat checkpoint v2".to_vec(); + text(&mut output, system); + output.push(match reasoning { + ReasoningMode::Direct => 0, + ReasoningMode::High => 1, + ReasoningMode::Max => 2, + }); + for message in messages { + output.push(u8::from(message.user)); + output.push(u8::from(message.skip_previous_eos)); + match &message.reasoning { + Some(reasoning) => { + output.push(1); + text(&mut output, reasoning); + } + None => output.push(0), + } + output.push(u8::from(message.reasoning_complete)); + text(&mut output, &message.content); } output } #[cfg(target_os = "macos")] fn conversation_tag(system: &str, reasoning: ReasoningMode, messages: &[ChatTurn]) -> [u8; 32] { - fn text(hasher: &mut Sha256, value: &str) { - hasher.update((value.len() as u64).to_le_bytes()); - hasher.update(value.as_bytes()); - } - - let mut hasher = Sha256::new(); - hasher.update(b"DS4Server chat checkpoint v1"); - text(&mut hasher, system); - hasher.update([match reasoning { - ReasoningMode::Direct => 0, - ReasoningMode::High => 1, - ReasoningMode::Max => 2, - }]); - for message in messages { - hasher.update([u8::from(message.user)]); - match &message.reasoning { - Some(reasoning) => { - hasher.update([1]); - text(&mut hasher, reasoning); - } - None => hasher.update([0]), - } - hasher.update([u8::from(message.reasoning_complete)]); - text(&mut hasher, &message.content); - } - hasher.finalize().into() + Sha256::digest(conversation_key(system, reasoning, messages)).into() } #[cfg(target_os = "macos")] @@ -1061,7 +1074,7 @@ mod sampling_tests { #[test] fn checkpoint_tag_covers_the_canonical_chat_state() { - let messages = [ChatTurn { + let mut messages = vec![ChatTurn { user: true, skip_previous_eos: false, reasoning: None, @@ -1081,6 +1094,20 @@ mod sampling_tests { tag, conversation_tag("System", ReasoningMode::Direct, &messages) ); + messages[0].skip_previous_eos = true; + assert_ne!( + tag, + conversation_tag("System", ReasoningMode::High, &messages) + ); + let prefix = conversation_key("System", ReasoningMode::High, &messages); + messages.push(ChatTurn { + user: false, + skip_previous_eos: false, + reasoning: Some("because".into()), + reasoning_complete: true, + content: "Hi".into(), + }); + assert!(conversation_key("System", ReasoningMode::High, &messages).starts_with(&prefix)); } #[test] diff --git a/src/engine/kvstore.rs b/src/engine/kvstore.rs new file mode 100644 index 0000000..104cc40 --- /dev/null +++ b/src/engine/kvstore.rs @@ -0,0 +1,385 @@ +use sha2::{Digest, Sha256}; +use std::cmp::Ordering; +use std::fs::{self, File}; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +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; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum StoreReason { + Cold, + Continued, +} + +impl StoreReason { + fn code(self) -> u8 { + match self { + Self::Cold => 1, + Self::Continued => 2, + } + } + + fn from_code(value: u8) -> Option { + Some(match value { + 1 => Self::Cold, + 2 => Self::Continued, + _ => return None, + }) + } + + fn is_anchor(self) -> bool { + self == Self::Cold + } +} + +#[derive(Clone, Debug)] +pub(super) struct Entry { + pub(super) checkpoint: PathBuf, + pub(super) tag: [u8; 32], + pub(super) tokens: u32, + context: u32, + hits: u32, + created_at: u64, + last_used: u64, + reason: StoreReason, + key: Vec, + bytes: u64, +} + +pub(super) struct KvStore { + directory: PathBuf, + budget_bytes: u64, +} + +impl KvStore { + pub(super) fn open(directory: &Path) -> Result { + fs::create_dir_all(directory).map_err(|error| { + format!( + "Could not create transient KV cache directory {}: {error}", + directory.display() + ) + })?; + Ok(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; + Ok(store) + } + + pub(super) fn find(&self, prompt_key: &[u8], context: u32) -> Option { + self.entries() + .into_iter() + .filter(|entry| entry.context == context && prompt_key.starts_with(&entry.key)) + .max_by(|left, right| { + left.key + .len() + .cmp(&right.key.len()) + .then_with(|| left.tokens.cmp(&right.tokens)) + }) + } + + pub(super) fn checkpoint_path(&self, key: &[u8]) -> PathBuf { + self.directory.join(format!("{}.bin", key_hash(key))) + } + + pub(super) fn record( + &self, + checkpoint: &Path, + key: &[u8], + tag: [u8; 32], + tokens: u32, + context: u32, + reason: StoreReason, + ) -> Result { + let checkpoint_bytes = fs::metadata(checkpoint) + .map_err(|error| error.to_string())? + .len(); + let index_bytes = index_size(key)?; + if checkpoint_bytes.saturating_add(index_bytes) > self.budget_bytes { + let _ = fs::remove_file(checkpoint); + return Ok(false); + } + let now = unix_time(); + let entry = Entry { + checkpoint: checkpoint.to_owned(), + tag, + tokens, + context, + hits: 0, + created_at: now, + last_used: now, + reason, + key: key.to_vec(), + bytes: checkpoint_bytes.saturating_add(index_bytes), + }; + write_entry(&entry)?; + self.evict(Some(checkpoint)); + Ok(checkpoint.exists()) + } + + pub(super) fn touch(&self, entry: &Entry) -> Result<(), String> { + let mut touched = entry.clone(); + touched.hits = touched.hits.saturating_add(1); + touched.last_used = unix_time(); + write_entry(&touched) + } + + pub(super) fn discard(&self, entry: &Entry) { + let _ = fs::remove_file(&entry.checkpoint); + let _ = fs::remove_file(index_path(&entry.checkpoint)); + } + + fn entries(&self) -> Vec { + let Ok(files) = fs::read_dir(&self.directory) else { + return Vec::new(); + }; + files + .filter_map(Result::ok) + .map(|file| file.path()) + .filter(|path| path.extension().is_some_and(|value| value == "meta")) + .filter_map(|path| read_entry(&path).ok()) + .filter(|entry| entry.checkpoint.is_file()) + .collect() + } + + fn evict(&self, protected: Option<&Path>) { + let mut entries = self.entries(); + let mut total = entries.iter().map(|entry| entry.bytes).sum::(); + let now = unix_time(); + while total > self.budget_bytes { + let Some((index, _)) = entries + .iter() + .enumerate() + .filter(|(_, entry)| protected != Some(entry.checkpoint.as_path())) + .min_by(|(_, left), (_, right)| { + eviction_score(left, now) + .partial_cmp(&eviction_score(right, now)) + .unwrap_or(Ordering::Equal) + .then_with(|| left.last_used.cmp(&right.last_used)) + }) + else { + break; + }; + let entry = entries.swap_remove(index); + total = total.saturating_sub(entry.bytes); + self.discard(&entry); + } + } +} + +fn eviction_score(entry: &Entry, now: u64) -> f64 { + let age = now.saturating_sub(entry.last_used.max(entry.created_at)) as f64; + let hits = f64::from(entry.hits) * f64::exp2(-age / HIT_HALF_LIFE_SECONDS); + let anchor = if entry.reason.is_anchor() { 2.0 } else { 1.0 }; + anchor * (hits + 1.0) * f64::from(entry.tokens) / entry.bytes.max(1) as f64 +} + +fn write_entry(entry: &Entry) -> Result<(), String> { + let path = index_path(&entry.checkpoint); + let temporary = path.with_extension("meta.tmp"); + let mut file = File::create(&temporary).map_err(|error| error.to_string())?; + file.write_all(INDEX_MAGIC) + .map_err(|error| error.to_string())?; + write_u32(&mut file, INDEX_VERSION)?; + file.write_all(&entry.tag) + .map_err(|error| error.to_string())?; + write_u32(&mut file, entry.tokens)?; + write_u32(&mut file, entry.context)?; + write_u32(&mut file, entry.hits)?; + file.write_all(&[entry.reason.code(), 0, 0, 0]) + .map_err(|error| error.to_string())?; + write_u64(&mut file, entry.created_at)?; + write_u64(&mut file, entry.last_used)?; + write_u64( + &mut file, + u64::try_from(entry.key.len()).map_err(|_| "KV cache key is too large")?, + )?; + file.write_all(&entry.key) + .map_err(|error| error.to_string())?; + file.sync_all().map_err(|error| error.to_string())?; + fs::rename(temporary, path).map_err(|error| error.to_string()) +} + +fn read_entry(path: &Path) -> Result { + let mut file = File::open(path).map_err(|error| error.to_string())?; + let mut magic = [0; 8]; + file.read_exact(&mut magic) + .map_err(|error| error.to_string())?; + if &magic != INDEX_MAGIC || read_u32(&mut file)? != INDEX_VERSION { + return Err("KV cache index has an invalid signature".into()); + } + let mut tag = [0; 32]; + file.read_exact(&mut tag) + .map_err(|error| error.to_string())?; + let tokens = read_u32(&mut file)?; + let context = read_u32(&mut file)?; + let hits = read_u32(&mut file)?; + let mut reason = [0; 4]; + file.read_exact(&mut reason) + .map_err(|error| error.to_string())?; + let reason = StoreReason::from_code(reason[0]) + .ok_or_else(|| "KV cache index has an invalid reason".to_owned())?; + let created_at = read_u64(&mut file)?; + let last_used = read_u64(&mut file)?; + let key_bytes = read_u64(&mut file)?; + if tokens == 0 || context == 0 || key_bytes > MAX_KEY_BYTES { + return Err("KV cache index has invalid dimensions".into()); + } + let mut key = vec![0; usize::try_from(key_bytes).map_err(|error| error.to_string())?]; + file.read_exact(&mut key) + .map_err(|error| error.to_string())?; + let checkpoint = path.with_extension("bin"); + if checkpoint.file_stem().and_then(|value| value.to_str()) != Some(&key_hash(&key)) { + return Err("KV cache index key does not match its filename".into()); + } + let bytes = fs::metadata(&checkpoint) + .map_err(|error| error.to_string())? + .len() + .saturating_add(fs::metadata(path).map_err(|error| error.to_string())?.len()); + Ok(Entry { + checkpoint, + tag, + tokens, + context, + hits, + created_at, + last_used, + reason, + key, + bytes, + }) +} + +fn index_path(checkpoint: &Path) -> PathBuf { + checkpoint.with_extension("meta") +} + +fn index_size(key: &[u8]) -> Result { + 84_u64 + .checked_add(u64::try_from(key.len()).map_err(|error| error.to_string())?) + .ok_or_else(|| "KV cache index is too large".to_owned()) +} + +fn key_hash(key: &[u8]) -> String { + let digest: [u8; 32] = Sha256::digest(key).into(); + const HEX: &[u8; 16] = b"0123456789abcdef"; + let mut output = String::with_capacity(64); + for byte in digest { + output.push(HEX[(byte >> 4) as usize] as char); + output.push(HEX[(byte & 0x0f) as usize] as char); + } + output +} + +fn unix_time() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |duration| duration.as_secs()) +} + +fn write_u32(file: &mut File, value: u32) -> Result<(), String> { + file.write_all(&value.to_le_bytes()) + .map_err(|error| error.to_string()) +} + +fn write_u64(file: &mut File, value: u64) -> Result<(), String> { + file.write_all(&value.to_le_bytes()) + .map_err(|error| error.to_string()) +} + +fn read_u32(file: &mut File) -> Result { + let mut bytes = [0; 4]; + file.read_exact(&mut bytes) + .map_err(|error| error.to_string())?; + Ok(u32::from_le_bytes(bytes)) +} + +fn read_u64(file: &mut File) -> Result { + let mut bytes = [0; 8]; + file.read_exact(&mut bytes) + .map_err(|error| error.to_string())?; + Ok(u64::from_le_bytes(bytes)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn temporary_directory(name: &str) -> PathBuf { + let path = std::env::temp_dir().join(format!( + "ds4-server-kvstore-{name}-{}-{}", + std::process::id(), + unix_time() + )); + fs::create_dir_all(&path).unwrap(); + path + } + + fn add(store: &KvStore, key: &[u8], tokens: u32, bytes: usize) -> Entry { + let path = store.checkpoint_path(key); + fs::write(&path, vec![0; bytes]).unwrap(); + assert!( + store + .record( + &path, + key, + [tokens as u8; 32], + tokens, + 4096, + StoreReason::Continued, + ) + .unwrap() + ); + store.find(key, 4096).unwrap() + } + + #[test] + fn longest_compatible_prefix_wins_and_hits_survive_reopen() { + let directory = temporary_directory("prefix"); + let store = KvStore::open(&directory).unwrap(); + add(&store, b"system", 100, 32); + let expected = add(&store, b"system-user-one", 200, 32); + + let found = store.find(b"system-user-one-more", 4096).unwrap(); + assert_eq!(found.checkpoint, expected.checkpoint); + store.touch(&found).unwrap(); + assert_eq!( + KvStore::open(&directory) + .unwrap() + .find(b"system-user-one-more", 4096) + .unwrap() + .hits, + 1 + ); + assert!(store.find(b"system-user-one-more", 2048).is_none()); + + fs::remove_dir_all(directory).unwrap(); + } + + #[test] + fn budget_evicts_the_least_valuable_unprotected_entry() { + let directory = temporary_directory("evict"); + let store = KvStore::with_budget(&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); + + assert!(!first.checkpoint.exists()); + assert!(second.checkpoint.exists()); + assert!(third.checkpoint.exists()); + + fs::remove_dir_all(directory).unwrap(); + } +}