Add prefix-aware KV session cache

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<bool, String> {
         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<u8> {
+    fn text(output: &mut Vec<u8>, 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]