Manage the KV cache disc usage from preferences and stats

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

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<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)]
 pub(crate) struct SteeringPreferences {
     pub(crate) file: Option<String>,
@@ -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",