Durable KV prefix cache: agent prefixes survive eviction, restart, and cold nodes

.gitignore

@@ -27,6 +27,7 @@ dist/MeshLLMFFI.xcframework.zip
 dist/native-sdk/
 dist/native-sdk-static/
 dist/llama-stage-static/
+dist/native-runtimes/
 sdk/kotlin/src/main/kotlin/uniffi/
 sdk/kotlin/example/example-jvm/src/main/kotlin/uniffi/
 sdk/node/native/

Cargo.toml

@@ -78,6 +78,7 @@ ahash = "0.8.12"
 anyhow = "1"
 blake3 = "1"
 clap = { version = "4", features = ["derive"] }
+memmap2 = "0.9"
 mesh-llm-build-info = { path = "crates/mesh-llm-build-info", version = "0.72.1" }
 mesh-llm-release-footer = { path = "crates/mesh-llm-release-footer", version = "0.72.1" }
 mesh-llm-skills = { path = "crates/mesh-llm-skills", version = "0.72.1" }
@@ -93,6 +94,13 @@ too_many_lines = "warn"
 # Do not remove this section: direct GGUF debug startup spends most of its time
 # hashing the source model via sha2/sha2-asm before native skippy model open,
 # so keep those crates near release speed in dev builds.
+# blake3 verifies every KV page read back from the disk tier, so an
+# unoptimized build hashes multi-gigabyte pages at a few hundred MB/s and
+# makes a cache restore look far more expensive than it is. Same reasoning
+# as sha2 below: cryptographic primitives are hot even in dev builds.
+[profile.dev.package.blake3]
+opt-level = 3
+
 [profile.dev.package.sha2]
 opt-level = 3
 

crates/mesh-llm-cli/src/parser/commands.rs

@@ -10,6 +10,17 @@ use mesh_llm_events::LogFormat;
 use mesh_llm_events::audit::{AuditLevel, AuditLogFormat};
 use serde::Serialize;
 
+fn parse_kv_cache_disk(value: &str) -> Result<String, String> {
+    let value = value.trim();
+    if value.eq_ignore_ascii_case("auto") || value.eq_ignore_ascii_case("off") {
+        return Ok(value.to_ascii_lowercase());
+    }
+    match value.parse::<f64>() {
+        Ok(gb) if gb.is_finite() && gb > 0.0 && gb * 1024.0 >= 1.0 => Ok(value.to_string()),
+        _ => Err("expected a positive size in GB, `auto`, or `off`".to_string()),
+    }
+}
+
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, ValueEnum)]
 pub enum BinaryFlavor {
     #[default]
@@ -671,6 +682,22 @@ pub struct Cli {
     #[arg(long)]
     pub max_vram: Option<f64>,
 
+    /// Keep agent prompt prefixes on disk so they survive a restart, capped at
+    /// this many GB for the whole node. Defaults to `auto`, sized from free
+    /// space. Use `off` to disable the disk cache.
+    #[arg(
+        long,
+        value_name = "GB|auto|off",
+        default_value = "auto",
+        value_parser = parse_kv_cache_disk
+    )]
+    pub kv_cache_disk: String,
+
+    /// Directory for the on-disk KV prefix cache. Defaults to
+    /// `~/.mesh-llm/kv-cache`.
+    #[arg(long, value_name = "PATH")]
+    pub kv_cache_disk_dir: Option<PathBuf>,
+
     /// Disable broadcasting GPU name, hostname, VRAM, and reserved bytes to peers. By default all nodes announce this hardware info.
     #[arg(long = "no-enumerate-host", hide = true)]
     pub no_enumerate_host: bool,
@@ -1196,6 +1223,33 @@ mod tests {
     use clap::{CommandFactory, Parser, error::ErrorKind};
     use mesh_llm_events::LogFormat;
 
+    #[test]
+    fn disk_cache_defaults_to_auto_and_accepts_explicit_off() {
+        let default = Cli::try_parse_from(["mesh-llm", "serve"]).expect("default CLI");
+        assert_eq!(default.kv_cache_disk, "auto");
+
+        let normalized = crate::parser::normalize_runtime_surface_args([
+            "mesh-llm",
+            "serve",
+            "--kv-cache-disk",
+            "off",
+        ]);
+        let off = Cli::try_parse_from(normalized.normalized).expect("off CLI");
+        assert_eq!(off.kv_cache_disk, "off");
+    }
+
+    #[test]
+    fn disk_cache_rejects_invalid_or_zero_budgets() {
+        for value in ["0", "-1", "tiny", "0.0001"] {
+            let argument = format!("--kv-cache-disk={value}");
+            let normalized =
+                crate::parser::normalize_runtime_surface_args(["mesh-llm", "serve", &argument]);
+            let error =
+                Cli::try_parse_from(normalized.normalized).expect_err("invalid disk cache budget");
+            assert_eq!(error.kind(), ErrorKind::ValueValidation);
+        }
+    }
+
     #[test]
     fn native_serving_plugin_deadline_rejects_zero() {
         let normalized = crate::parser::normalize_runtime_surface_args([

crates/mesh-llm-host-runtime/src/api/routes/logs/tests.rs

@@ -1,4 +1,8 @@
 use mesh_llm_events::logging::identifiers::RequestId;
+use openai_frontend::{
+    OpenAiBackendOperation, OpenAiFrontendRoute, OpenAiLifecycleContext, OpenAiLifecycleEvent,
+    OpenAiRequestMethod, OpenAiTerminalResult, OpenAiUsage, Usage,
+};
 
 use super::*;
 
@@ -25,6 +29,56 @@ fn runtime() -> (tempfile::TempDir, LoggingRuntimeState) {
     (temp, LoggingRuntimeState::initialize(&foundation, &config))
 }
 
+#[tokio::test]
+async fn disk_prefix_cached_tokens_reach_durable_usage_dto() {
+    let (_temp, state) = runtime();
+    let request_id = RequestId::new();
+    let context = OpenAiLifecycleContext::new(
+        request_id,
+        OpenAiRequestMethod::Post,
+        OpenAiFrontendRoute::ChatCompletions,
+    );
+    let operation = OpenAiBackendOperation::ChatCompletionStream;
+    let observer = state
+        .openai_lifecycle_observer()
+        .expect("OpenAI lifecycle observer");
+
+    observer.observe(&OpenAiLifecycleEvent::Admitted {
+        context: context.clone(),
+    });
+    // A disk-prefix restoration is surfaced by the backend through the
+    // OpenAI usage detail; exercise that production conversion before logging.
+    let disk_prefix_usage = Usage::new(21, 8).with_cached_tokens(13);
+    observer.observe(&OpenAiLifecycleEvent::ResponseCompleted {
+        context: context.clone(),
+        operation,
+        usage: OpenAiUsage::from(&disk_prefix_usage),
+    });
+    observer.observe(&OpenAiLifecycleEvent::StreamTerminal {
+        context,
+        result: OpenAiTerminalResult::Completed { status_code: 200 },
+    });
+    assert!(state.pump_persistence_for_test().await > 0);
+
+    let request_key = request_id.as_uuid().to_string();
+    let events = request_events(
+        &state,
+        "/api/logs/requests/id/events?limit=20",
+        &request_key,
+    )
+    .await
+    .expect("durable request events");
+    let wire = serde_json::to_value(events).expect("event DTO JSON");
+    let usage = wire["items"]
+        .as_array()
+        .expect("event items")
+        .iter()
+        .find(|event| event["kind"] == "usage_recorded")
+        .expect("durable usage record");
+
+    assert_eq!(usage["cachedPromptTokens"], 13);
+}
+
 #[tokio::test]
 async fn list_merges_active_and_durable_without_duplicate_ids() {
     let (_temp, state) = runtime();

crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs

@@ -106,13 +106,49 @@ impl FamilyPolicy {
                     max_bytes,
                     min_tokens,
                     shared_prefix_stride_tokens: 128,
-                    shared_prefix_record_limit: 2,
+                    shared_prefix_record_limit: derive_shared_prefix_record_limit(bounded_entries),
                 })
             }
         }
     }
 }
 
+/// How many prefix lengths a single request may record.
+///
+/// Recording only the two longest candidates makes cross-session
+/// sharing unreachable in practice: for an 8000-token agentic prompt
+/// the recorded lengths are `[8000, 7936]`, both in the request's
+/// tail. A second session with the same system prompt and tool
+/// schemas but a different tail probes the shared region, finds
+/// nothing, and pays a full cold prefill — even though the lookup
+/// grid already probes 62 different lengths for it.
+///
+/// A deeper ladder lets `PrefixCandidatePolicy` keep the near-tail
+/// continuation pages *and* reach down into the shared
+/// system-prompt/tool-schema region. See the ladder documentation in
+/// `skippy-cache/src/config.rs` for the selection strategy.
+///
+/// The limit is bounded by cache cardinality for the same reason
+/// `max_entries` is bounded by the cell pool: every recorded
+/// candidate pins KV cells, so a ladder deeper than the cache can
+/// hold just churns the LRU and evicts the entries it recorded a
+/// moment earlier. Recording at most a quarter of the cache per
+/// request leaves room for several distinct prompts to coexist,
+/// which is the whole point of retention.
+///
+/// Never below the historical default of 2, so no configuration
+/// records less than it did before.
+const MIN_SHARED_PREFIX_RECORD_LIMIT: u64 = 2;
+const MAX_SHARED_PREFIX_RECORD_LIMIT: u64 = 6;
+
+fn derive_shared_prefix_record_limit(max_entries: usize) -> u64 {
+    let quarter_of_cache = (max_entries as u64) / 4;
+    quarter_of_cache.clamp(
+        MIN_SHARED_PREFIX_RECORD_LIMIT,
+        MAX_SHARED_PREFIX_RECORD_LIMIT,
+    )
+}
+
 /// Cap the prefix-cache `max_entries` so resident prefixes cannot
 /// exhaust the unified KV cell pool.
 ///
@@ -925,4 +961,29 @@ mod tests {
         assert_eq!(cache.payload, StageKvCachePayload::ResidentKv);
         assert_eq!(cache.max_bytes, 3_211_264);
     }
+
+    /// A deeper record ladder is what makes cross-session prefix sharing
+    /// reachable, but it must stay bounded by what the cache can actually
+    /// hold, or it just churns the LRU.
+    #[test]
+    fn record_limit_scales_with_cache_capacity() {
+        // Tiny caches keep the historical behaviour.
+        assert_eq!(derive_shared_prefix_record_limit(1), 2);
+        assert_eq!(derive_shared_prefix_record_limit(4), 2);
+        assert_eq!(derive_shared_prefix_record_limit(8), 2);
+        // Roomier caches record a deeper ladder...
+        assert_eq!(derive_shared_prefix_record_limit(16), 4);
+        // ...but never more than a quarter of the cache, and never unbounded.
+        assert_eq!(derive_shared_prefix_record_limit(24), 6);
+        assert_eq!(derive_shared_prefix_record_limit(512), 6);
+    }
+
+    /// The limit must never drop below what shipped previously, so no
+    /// configuration records less than it did before this change.
+    #[test]
+    fn record_limit_never_regresses_below_the_historical_default() {
+        for max_entries in 0..64 {
+            assert!(derive_shared_prefix_record_limit(max_entries) >= 2);
+        }
+    }
 }

crates/mesh-llm-host-runtime/src/runtime/options.rs

@@ -70,6 +70,8 @@ pub struct RuntimeOptions {
     pub split_topology_lock: Option<PathBuf>,
     pub ctx_size: Option<u32>,
     pub max_vram: Option<f64>,
+    pub kv_cache_disk: String,
+    pub kv_cache_disk_dir: Option<PathBuf>,
     pub no_enumerate_host: bool,
     pub bin_dir: Option<PathBuf>,
     pub llama_flavor: Option<mesh_llm_system::backend::BinaryFlavor>,
@@ -147,6 +149,8 @@ impl Default for RuntimeOptions {
             split_topology_lock: None,
             ctx_size: None,
             max_vram: None,
+            kv_cache_disk: "auto".to_string(),
+            kv_cache_disk_dir: None,
             no_enumerate_host: false,
             bin_dir: None,
             llama_flavor: None,

crates/mesh-llm-host-runtime/src/runtime/run_auto.rs

@@ -604,9 +604,55 @@ pub(super) fn configure_run_auto_process_state(
     skippy_runtime::set_filtered_native_logs_enabled(true);
     bridge_skippy_native_logs(native_log_rx);
     skippy::configure_materialized_stage_cache();
+    configure_kv_disk_cache(options);
     configure_skippy_native_logging(runtime.as_ref().map(|runtime| runtime.dir()));
 }
 
+/// Propagate CLI policy through a typed process-local configuration. This does
+/// not mutate process environment; lower-level legacy environment variables
+/// remain available when skippy-server is used without this host configuration.
+fn configure_kv_disk_cache(options: &RuntimeOptions) {
+    let budget = parse_kv_cache_disk(&options.kv_cache_disk)
+        .expect("CLI validates --kv-cache-disk before runtime startup");
+    let budget = match budget {
+        KvDiskBudget::Off => skippy_server::KvDiskCacheBudget::Off,
+        KvDiskBudget::Auto => skippy_server::KvDiskCacheBudget::Auto,
+        KvDiskBudget::Mib(mib) => {
+            skippy_server::KvDiskCacheBudget::Bytes(mib.saturating_mul(1024 * 1024))
+        }
+    };
+    let config = skippy_server::KvDiskCacheConfig {
+        budget,
+        directory: options.kv_cache_disk_dir.clone(),
+    };
+    if skippy_server::configure_kv_disk_cache(config).is_err() {
+        tracing::debug!("KV disk cache policy was already configured");
+    }
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum KvDiskBudget {
+    Off,
+    Auto,
+    Mib(u64),
+}
+
+fn parse_kv_cache_disk(raw: &str) -> Option<KvDiskBudget> {
+    let value = raw.trim();
+    if value.eq_ignore_ascii_case("off") {
+        return Some(KvDiskBudget::Off);
+    }
+    if value.eq_ignore_ascii_case("auto") {
+        return Some(KvDiskBudget::Auto);
+    }
+    match value.parse::<f64>() {
+        Ok(gb) if gb.is_finite() && gb > 0.0 && gb * 1024.0 >= 1.0 => {
+            Some(KvDiskBudget::Mib((gb * 1024.0).round() as u64))
+        }
+        _ => None,
+    }
+}
+
 pub(super) fn spawn_node_benchmark_task(node: &mesh::Node, bin_dir: &Path) {
     let mem_arc = node.gpu_mem_bandwidth_gbps.clone();
     let compute_fp32_arc = node.gpu_compute_tflops_fp32.clone();
@@ -1521,3 +1567,30 @@ pub(super) async fn run_auto(ctx: RunAutoContext) -> Result<()> {
     }
     Ok(())
 }
+
+#[cfg(test)]
+mod kv_cache_disk_tests {
+    use super::{KvDiskBudget, parse_kv_cache_disk};
+
+    #[test]
+    fn sizes_are_read_as_gigabytes() {
+        assert_eq!(parse_kv_cache_disk("8"), Some(KvDiskBudget::Mib(8192)));
+        assert_eq!(parse_kv_cache_disk(" 0.5 "), Some(KvDiskBudget::Mib(512)));
+    }
+
+    #[test]
+    fn auto_defers_to_the_free_space_policy() {
+        assert_eq!(parse_kv_cache_disk("auto"), Some(KvDiskBudget::Auto));
+        assert_eq!(parse_kv_cache_disk("AUTO"), Some(KvDiskBudget::Auto));
+        assert_eq!(parse_kv_cache_disk("off"), Some(KvDiskBudget::Off));
+    }
+
+    /// Each of these would otherwise disable the tier while looking like the
+    /// user had enabled it, which is the one outcome worth being loud about.
+    #[test]
+    fn unusable_budgets_are_rejected_rather_than_silently_ignored() {
+        for raw in ["0", "-4", "", "lots", "0.0001", "nan", "inf"] {
+            assert_eq!(parse_kv_cache_disk(raw), None, "should reject {raw:?}");
+        }
+    }
+}

crates/mesh-llm/src/lib.rs

@@ -386,6 +386,8 @@ fn runtime_options_from_cli(cli: mesh_llm_cli::Cli) -> mesh_llm_host_runtime::Ru
         split_topology_lock: cli.split_topology_lock,
         ctx_size: cli.ctx_size,
         max_vram: cli.max_vram,
+        kv_cache_disk: cli.kv_cache_disk,
+        kv_cache_disk_dir: cli.kv_cache_disk_dir,
         no_enumerate_host: cli.no_enumerate_host,
         bin_dir: cli.bin_dir,
         llama_flavor: cli.llama_flavor.map(map_binary_flavor),