From ce0741eb51913cc5b7af0bfdacc0c8a78fe66eac Mon Sep 17 00:00:00 2001 From: scama Date: Mon, 14 Sep 2026 07:25:13 +1000 Subject: [PATCH] feat(skippy): wire bounded host-RAM L2 serving --- .../src/model/built_in_schema/declarations.rs | 2 +- crates/mesh-llm-config/src/wiring_status.rs | 8 +- .../src/inference/skippy/family_policy.rs | 1 + .../inference/skippy/resolver/resolution.rs | 13 +- .../src/inference/skippy/resolver/support.rs | 3 - .../src/inference/skippy/resolver/tests.rs | 2 + .../inference/skippy/resolver/translation.rs | 21 +- .../src/inference/skippy/resolver/types.rs | 1 + .../src/runtime/local_model_only.rs | 1 + .../config_schema_defaults_ui_reference.json | 7 + crates/skippy-cache/src/l2/mod.rs | 318 ++++++++++++- .../src/prompt_cli/stage_config.rs | 2 + crates/skippy-protocol/src/config.rs | 4 + .../src/binary_transport/stage_execution.rs | 1 + .../src/frontend/local_generation/tests.rs | 1 + .../token_generation/kv_restore.rs | 2 +- .../src/frontend/prefix_cache.rs | 2 + .../src/frontend/tests/support.rs | 1 + .../src/kv_integration/activation.rs | 1 + .../src/kv_integration/cache_affinity.rs | 1 + .../src/kv_integration/config.rs | 126 ++++- .../src/kv_integration/exact_state.rs | 27 +- .../src/kv_integration/identity.rs | 1 + .../src/kv_integration/l2_serving.rs | 437 ++++++++++++++++++ .../skippy-server/src/kv_integration/mod.rs | 32 ++ .../src/kv_integration/resident_prefix.rs | 1 + docs/USAGE.md | 2 +- docs/skippy/CONFIGURATION.md | 2 +- docs/skippy/PROMPT_CACHE.md | 21 + tools/xtask/data/console_print_allowlist.json | 8 +- website/src/docs/pages/config-reference.md | 2 +- 31 files changed, 1018 insertions(+), 33 deletions(-) create mode 100644 crates/skippy-server/src/kv_integration/l2_serving.rs diff --git a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs index 741bd946c2..625c3edcce 100644 --- a/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs +++ b/crates/mesh-llm-config/src/model/built_in_schema/declarations.rs @@ -342,7 +342,7 @@ fn model_fit_settings( ), basic_setting(&format!("{prefix}.kv_offload"), bool_or_auto_schema()), basic_setting(&format!("{prefix}.kv_unified"), bool_or_auto_schema()), - unwired_setting( + basic_setting( &format!("{prefix}.cache_ram_mib"), ConfigValueSchema::Integer, ), diff --git a/crates/mesh-llm-config/src/wiring_status.rs b/crates/mesh-llm-config/src/wiring_status.rs index caa7c63640..f757699335 100644 --- a/crates/mesh-llm-config/src/wiring_status.rs +++ b/crates/mesh-llm-config/src/wiring_status.rs @@ -622,10 +622,10 @@ pub const WIRING_MANIFEST: &[WiringEntry] = &[ }, WiringEntry { path: "model_fit.cache_ram_mib", - status: WiringStatus::Unwired, - owner: "PR2", - reason: "Any positive value fails at model load", - behavior: WiringBehavior::BailsDownstream, + status: WiringStatus::Wired, + owner: "n/a", + reason: "", + behavior: WiringBehavior::None, }, WiringEntry { path: "model_fit.cache_idle_slots", diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs index 7441df5e4f..ca91b85e3e 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/family_policy.rs @@ -54,6 +54,7 @@ impl FamilyPolicy { payload: StageKvCachePayload::Auto, max_entries: bounded_entries, max_bytes, + l2_max_bytes: 0, min_tokens, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: derive_shared_prefix_record_limit(bounded_entries), diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs index 834df81d14..e737503d8f 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/resolution.rs @@ -17,7 +17,7 @@ use super::types::{ BUILTIN_BATCH, BUILTIN_CTX_SIZE, BUILTIN_PARALLEL, BUILTIN_PREFILL_CHUNK_SIZE, BUILTIN_SAFETY_MARGIN_GB, BUILTIN_UBATCH, ResolvedHardwareConfig, ResolvedModelFitConfig, ResolvedMultimodalConfig, ResolvedSkippyConfig, ResolvedSkippyExecutionConfig, - ResolvedThroughputConfig, SkippyConfigResolveRequest, + ResolvedStageKvCache, ResolvedThroughputConfig, SkippyConfigResolveRequest, }; use crate::plugin::{ BoolOrAuto, ModelConfigDefaults, ModelConfigEntry, ModelFitConfig, ThroughputConfig, @@ -266,6 +266,16 @@ fn resolve_model_fit_config( .or(context.global_model_fit.and_then(|fit| fit.flash_attention)) .unwrap_or_else(|| effective_flash_attention(&cache_type_v)); let prefix_cache = resolve_prefix_cache(context.model_fit, context.global_model_fit)?; + let l2_max_bytes = pick_owned( + context.model_fit.and_then(|fit| fit.cache_ram_mib), + context.global_model_fit.and_then(|fit| fit.cache_ram_mib), + ) + .unwrap_or(0) + .checked_mul(1024 * 1024) + .ok_or_else(|| anyhow::anyhow!("model_fit.cache_ram_mib exceeds the byte range"))?; + if l2_max_bytes > 0 && matches!(prefix_cache, ResolvedStageKvCache::Disabled) { + anyhow::bail!("model_fit.cache_ram_mib requires prefix caching to be enabled"); + } Ok(ResolvedModelFitConfig { ctx_size, @@ -275,6 +285,7 @@ fn resolve_model_fit_config( cache_type_v, kv_cache_policy: kv.effective_policy, prefix_cache, + l2_max_bytes, kv_offload, kv_offload_resolved, kv_unified, diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs index 3b4cee8854..036412fdcb 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/support.rs @@ -68,9 +68,6 @@ pub(super) fn reject_unsupported_model_fit_controls( let Some(config) = config else { return Ok(()); }; - if config.cache_ram_mib.unwrap_or(0) > 0 { - bail!("skippy model_fit.cache_ram_mib is not supported by the pinned runtime"); - } if config.keep_tokens.unwrap_or(0) > 0 { bail!("skippy model_fit.keep_tokens is not supported by the pinned runtime"); } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs index 9193a99f52..2459917ade 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs @@ -2033,6 +2033,7 @@ fn staged_controls_propagate_into_stage_config_and_embedded_openai_args() { r#" [defaults.model_fit] prompt_cache = true +cache_ram_mib = 64 [defaults.model_fit.prefix_cache] enabled = true @@ -2080,6 +2081,7 @@ draft_max_tokens = 8 assert_eq!(kv_cache.shared_prefix_stride_tokens, 48); assert_eq!(kv_cache.shared_prefix_record_limit, 3); assert_eq!(kv_cache.payload, StageKvCachePayload::ResidentKv); + assert_eq!(kv_cache.l2_max_bytes, 64 * 1024 * 1024); let openai = resolved .to_embedded_openai_args(4096, true) diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs index be415bf193..8a79e5f3cd 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/translation.rs @@ -472,23 +472,25 @@ impl ResolvedSkippyConfig { &self, family_default: Option, ) -> Result> { - match &self.model_fit.prefix_cache { - ResolvedStageKvCache::FamilyDefault => Ok(family_default), - ResolvedStageKvCache::Disabled => Ok(Some(StageKvCacheConfig { + let mut resolved = match &self.model_fit.prefix_cache { + ResolvedStageKvCache::FamilyDefault => family_default, + ResolvedStageKvCache::Disabled => Some(StageKvCacheConfig { mode: StageKvCacheMode::Disabled, payload: StageKvCachePayload::Auto, max_entries: 0, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 0, shared_prefix_stride_tokens: 0, shared_prefix_record_limit: 0, - })), + }), ResolvedStageKvCache::Explicit(template) => { let mut cache = family_default.unwrap_or(StageKvCacheConfig { mode: template.mode.clone(), payload: StageKvCachePayload::Auto, max_entries: 128, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, @@ -510,9 +512,18 @@ impl ResolvedSkippyConfig { if let Some(value) = template.shared_prefix_record_limit { cache.shared_prefix_record_limit = value as u64; } - Ok(Some(cache)) + Some(cache) } + }; + if self.model_fit.l2_max_bytes > 0 && resolved.is_none() { + anyhow::bail!( + "model_fit.cache_ram_mib requires an executable prefix-cache configuration" + ); + } + if let Some(cache) = resolved.as_mut() { + cache.l2_max_bytes = self.model_fit.l2_max_bytes; } + Ok(resolved) } } diff --git a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs index f49d29e7e7..4e2217184b 100644 --- a/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs +++ b/crates/mesh-llm-host-runtime/src/inference/skippy/resolver/types.rs @@ -74,6 +74,7 @@ pub(crate) struct ResolvedModelFitConfig { pub(crate) cache_type_v: String, pub(crate) kv_cache_policy: String, pub(crate) prefix_cache: ResolvedStageKvCache, + pub(crate) l2_max_bytes: u64, pub(crate) kv_offload: String, /// Parsed `kv_offload` for the native tri-state control. `None` covers /// both "auto" and any value that did not parse to a bool. diff --git a/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs b/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs index da4127a4c0..c6408f7d11 100644 --- a/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs +++ b/crates/mesh-llm-host-runtime/src/runtime/local_model_only.rs @@ -505,6 +505,7 @@ mod tests { payload: StageKvCachePayload::Auto, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 8, shared_prefix_stride_tokens: 8, shared_prefix_record_limit: 2, diff --git a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json index ac43057d6d..cbab8a580b 100644 --- a/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json +++ b/crates/mesh-llm-host-runtime/tests/fixtures/config_schema_defaults_ui_reference.json @@ -112,6 +112,13 @@ "kind": "built_in" } }, + { + "canonical_path": "defaults.model_fit.cache_ram_mib", + "support": "supported", + "source": { + "kind": "built_in" + } + }, { "canonical_path": "defaults.model_fit.cache_type_k", "support": "supported", diff --git a/crates/skippy-cache/src/l2/mod.rs b/crates/skippy-cache/src/l2/mod.rs index 3591b589bf..5d72cc8b88 100644 --- a/crates/skippy-cache/src/l2/mod.rs +++ b/crates/skippy-cache/src/l2/mod.rs @@ -37,9 +37,8 @@ //! `CacheBytes` is a block-backed view over the shared segment storages, //! contiguous in the single-segment case. //! -//! This first slice is a standalone store with no wiring into the request -//! path; the benchmark harness drives it directly. L2 promotion/demotion -//! policy and server integration land in a later slice. +//! `skippy-server` wires this tier as a bounded mirror of repeated or +//! high-value L3 fills. An L2 hit rewarms L1. use std::{ collections::HashMap, ops::Range, @@ -49,10 +48,14 @@ use std::{ }, }; -use crate::payload::{CacheBytes, ExactStatePayloadKind}; -use crate::{HandoffManifest, segment_digest}; #[cfg(test)] -use crate::{HandoffSegmentRef, MANIFEST_VERSION, PayloadCodec, SegmentCodecIdentity}; +use crate::PayloadCodec; +use crate::payload::{CacheBytes, ExactStatePayload, ExactStatePayloadKind}; +use crate::{ + HandoffManifest, HandoffSegmentRef, MANIFEST_VERSION, SegmentCodecIdentity, segment_digest, +}; + +const DIRECT_SEGMENT_BYTES: usize = 1024 * 1024; /// Where an entry came from, for telemetry and promotion policy later. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -117,6 +120,9 @@ pub struct L2Layout { pub total_bytes: u64, pub kv_bytes: u64, pub recurrent_bytes: u64, + /// Opaque serialized runtime KV-page descriptor. The serving layer + /// validates this before importing the payload. + pub kv_desc_json: Option, /// `(segment digest, byte range within the assembled wire)` per /// manifest segment, in manifest order. Ranges concatenate to /// `0..total_bytes` exactly as the L3 manifest tiles them. @@ -150,6 +156,10 @@ impl ExactStatePayloadMirror { } } + pub fn kv_desc_json(&self) -> Option<&str> { + self.layout().kv_desc_json.as_deref() + } + /// Build a mirror from a captured L3 manifest. Callers must verify the /// payload wire against `manifest.payload_digest` — `admit` does this — /// before the mirror is stored. @@ -199,6 +209,7 @@ impl ExactStatePayloadMirror { total_bytes: manifest.total_bytes, kv_bytes: manifest.kv_bytes, recurrent_bytes: manifest.recurrent_bytes, + kv_desc_json: manifest.kv_desc_json.clone(), segments, }; Ok(match kind { @@ -574,6 +585,33 @@ fn validate_layout<'a>( } Ok(validated) } + +fn payload_wire(payload: &ExactStatePayload) -> Result<(Vec, u64, u64), L2InsertRefusal> { + let malformed = |error: anyhow::Error| L2InsertRefusal::MalformedManifest(error.to_string()); + match payload { + ExactStatePayload::FullState { bytes } => { + let wire = bytes.as_cow().map_err(malformed)?.into_owned(); + let kv_bytes = wire.len() as u64; + Ok((wire, kv_bytes, 0)) + } + ExactStatePayload::RecurrentOnly { recurrent } => { + let wire = recurrent.as_cow().map_err(malformed)?.into_owned(); + let recurrent_bytes = wire.len() as u64; + Ok((wire, 0, recurrent_bytes)) + } + ExactStatePayload::KvRecurrent { kv, recurrent } => { + let kv = kv.as_cow().map_err(malformed)?; + let recurrent = recurrent.as_cow().map_err(malformed)?; + let kv_bytes = kv.len() as u64; + let recurrent_bytes = recurrent.len() as u64; + let mut wire = Vec::with_capacity(kv.len().saturating_add(recurrent.len())); + wire.extend_from_slice(kv.as_ref()); + wire.extend_from_slice(recurrent.as_ref()); + Ok((wire, kv_bytes, recurrent_bytes)) + } + } +} + impl L2Tier { pub fn new(budget_bytes: u64) -> Self { Self { @@ -587,6 +625,64 @@ impl L2Tier { self.budget_bytes } + /// Admit an exact-state payload restored from the authoritative L3 tier. + /// + /// The payload is cut into stable one-MiB content chunks. This preserves + /// immutable sharing between related entries while keeping all hashing + /// and copying on the existing cache worker rather than the request path. + pub fn admit_payload( + &self, + cache_key: String, + token_count: u64, + expected_payload_digest: &str, + payload: &ExactStatePayload, + kv_desc_json: Option, + origin: L2Origin, + ) -> Result, L2InsertRefusal> { + let (wire, kv_bytes, recurrent_bytes) = payload_wire(payload)?; + let payload_digest = segment_digest(&wire); + if payload_digest != expected_payload_digest { + self.stats.admission_rejects.fetch_add(1, Ordering::Relaxed); + return Err(L2InsertRefusal::DigestMismatch { + expected: expected_payload_digest.to_string(), + actual: payload_digest, + }); + } + let mut manifest = HandoffManifest::new(String::new(), payload.kind().to_string()); + manifest.version = MANIFEST_VERSION; + manifest.total_bytes = wire.len() as u64; + manifest.payload_digest = expected_payload_digest.to_string(); + manifest.kv_bytes = kv_bytes; + manifest.recurrent_bytes = recurrent_bytes; + manifest.kv_desc_json = kv_desc_json; + manifest.token_count = token_count; + manifest.segments = wire + .chunks(DIRECT_SEGMENT_BYTES) + .enumerate() + .scan(0u64, |offset, (index, bytes)| { + let start = *offset; + *offset = offset.saturating_add(bytes.len() as u64); + Some(HandoffSegmentRef { + index: index as u32, + offset: start, + bytes: bytes.len() as u64, + digest: segment_digest(bytes), + codec_identity: Some(SegmentCodecIdentity::raw(bytes.len() as u64)), + meta_json: None, + }) + }) + .collect(); + let mirror = ExactStatePayloadMirror::from_manifest(&manifest)?; + self.admit( + cache_key, + token_count, + expected_payload_digest.to_string(), + &wire, + mirror, + origin, + ) + } + /// Admit an assembled entry. /// /// `wire` is the payload's concatenated L3 wire — the exact bytes whose @@ -980,6 +1076,62 @@ impl L2Tier { bytes } + /// Remove every entry that mirrors one durable payload digest. + pub fn remove_by_digest(&self, payload_digest: &str) -> Vec { + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + let keys = inner + .map + .iter() + .filter(|(_, entry)| entry.payload_digest == payload_digest) + .map(|(key, _)| key.clone()) + .collect::>(); + let mut evictions = Vec::with_capacity(keys.len()); + for key in keys { + let Some(removed) = inner.map.remove(&key) else { + continue; + }; + let mut retained = 0u64; + for digest in removed.payload.segment_digests() { + if inner + .map + .values() + .any(|other| other.payload.segment_digests().contains(&digest)) + { + retained = retained.saturating_add( + inner + .segments + .get(digest) + .map(|handle| handle.bytes.len() as u64) + .unwrap_or(0), + ); + } + } + Self::recompute_all_charges(&mut inner); + let before = inner.bytes; + self.release_entry_segments(&mut inner, &removed, &[]); + evictions.push(L2Eviction { + cache_key: key, + freed_bytes: before.saturating_sub(inner.bytes), + retained_bytes: retained, + }); + } + evictions + } + + /// Evict least-recently-used entries until physical usage is at or below + /// `target_bytes`. Targets above the configured budget are clamped. + pub fn shrink_to(&self, target_bytes: u64) -> Vec { + let mut inner = self.inner.lock().expect("L2 map lock poisoned"); + let mut journal = AdmitJournal::default(); + self.evict_to_limit( + &mut inner, + target_bytes.min(self.budget_bytes), + "", + &[], + &mut journal, + ) + } + pub fn len(&self) -> usize { self.inner.lock().expect("L2 map lock poisoned").map.len() } @@ -1214,6 +1366,7 @@ mod tests { total_bytes: len, kv_bytes: len, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(segment_digest(w), 0..len)], }, } @@ -1241,6 +1394,7 @@ mod tests { total_bytes: len, kv_bytes: len, recurrent_bytes: 0, + kv_desc_json: None, segments, }, } @@ -1280,6 +1434,94 @@ mod tests { assert_eq!(bytes.as_ref(), &w[..], "served bytes must equal the wire"); } + #[test] + fn admit_payload_round_trips_composite_state_and_descriptor() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[1, 2, 3]); + let kv = vec![1, 2, 3, 4]; + let recurrent = vec![5, 6, 7]; + let descriptor = r#"{"token_start":0,"token_count":3}"#.to_string(); + let expected_digest = segment_digest(&[kv.as_slice(), recurrent.as_slice()].concat()); + + tier.admit_payload( + k.clone(), + 3, + &expected_digest, + &ExactStatePayload::kv_recurrent(kv.clone(), recurrent.clone()), + Some(descriptor.clone()), + L2Origin::Direct, + ) + .expect("direct payload admission must fit"); + + let hit = tier.get(&k).expect("admitted payload must hit"); + assert_eq!(hit.payload.kv_desc_json(), Some(descriptor.as_str())); + let payload = hit.to_payload(); + assert_eq!( + payload + .kv_bytes() + .expect("read KV bytes") + .expect("composite payload has KV") + .as_ref(), + kv.as_slice() + ); + assert_eq!( + payload + .recurrent_state_bytes() + .expect("read recurrent bytes") + .as_ref(), + recurrent.as_slice() + ); + } + + #[test] + fn admit_payload_round_trips_full_state() { + let tier = L2Tier::new(1 << 20); + let k = key("ns", &[9]); + let bytes = vec![7; 128]; + let expected_digest = segment_digest(&bytes); + + tier.admit_payload( + k.clone(), + 1, + &expected_digest, + &ExactStatePayload::full_state(bytes.clone()), + None, + L2Origin::Direct, + ) + .expect("full-state admission must fit"); + + let payload = tier + .get(&k) + .expect("admitted payload must hit") + .to_payload(); + assert_eq!( + payload + .full_state_bytes_timed() + .expect("read full state") + .0 + .as_ref(), + bytes.as_slice() + ); + } + + #[test] + fn direct_payload_admission_requires_the_durable_digest() { + let tier = L2Tier::new(1 << 20); + let payload = ExactStatePayload::full_state(vec![7; 128]); + let error = tier + .admit_payload( + key("ns", &[9]), + 1, + &segment_digest(b"different"), + &payload, + None, + L2Origin::FromL3, + ) + .expect_err("mismatched durable identity must be refused"); + assert!(matches!(error, L2InsertRefusal::DigestMismatch { .. })); + assert!(tier.is_empty()); + } + #[test] fn admission_digest_mismatch_refuses_and_stores_nothing() { let tier = L2Tier::new(1 << 20); @@ -1856,6 +2098,65 @@ mod tests { assert_eq!(tier.stats().segments, 0); } + #[test] + fn remove_by_digest_drops_only_matching_mirrors() { + let tier = L2Tier::new(1 << 20); + let (shared, shared_digest) = wire(64, 8); + let (other, other_digest) = wire(64, 9); + for tokens in [&[1][..], &[2][..]] { + tier.admit( + key("ns", tokens), + 1, + shared_digest.clone(), + &shared, + single_segment_mirror(&shared), + L2Origin::FromL3, + ) + .expect("shared mirror fits"); + } + let other_key = key("ns", &[3]); + tier.admit( + other_key.clone(), + 1, + other_digest, + &other, + single_segment_mirror(&other), + L2Origin::FromL3, + ) + .expect("other mirror fits"); + + let removed = tier.remove_by_digest(&shared_digest); + assert_eq!(removed.len(), 2); + assert_eq!(tier.len(), 1); + assert!(tier.get(&other_key).is_some()); + } + + #[test] + fn shrink_to_evicts_lru_until_the_target_is_met() { + let tier = L2Tier::new(256); + let mut keys = Vec::new(); + for i in 0..3i32 { + let (bytes, digest) = wire(64, i as u8 + 1); + let cache_key = key("ns", &[i]); + tier.admit( + cache_key.clone(), + 1, + digest, + &bytes, + single_segment_mirror(&bytes), + L2Origin::FromL3, + ) + .expect("entry fits"); + keys.push(cache_key); + } + assert!(tier.get(&keys[0]).is_some(), "first entry becomes hottest"); + + let evicted = tier.shrink_to(128); + assert_eq!(evicted.len(), 1); + assert_eq!(evicted[0].cache_key, keys[1]); + assert_eq!(tier.stats().bytes, 128); + } + #[test] fn identical_wire_same_key_readmit_keeps_its_own_segments() { // The original failure: re-admitting an identical wire at the same @@ -1947,6 +2248,7 @@ mod tests { total_bytes: total, kv_bytes: total, recurrent_bytes: 0, + kv_desc_json: None, segments: grown_segments, }, }; @@ -2031,6 +2333,7 @@ mod tests { total_bytes: 120, kv_bytes: 120, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(segment_digest(&x), 0..60), (segment_digest(&y), 60..120)], }, }; @@ -2138,6 +2441,7 @@ mod tests { total_bytes: 32, kv_bytes: 32, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![ (segment_digest(&seg), 0..16), (segment_digest(&seg), 16..32), @@ -2287,6 +2591,7 @@ mod tests { total_bytes: 64, kv_bytes: 64, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(x_digest.clone(), 0..32), (x_digest, 32..64)], }, }; @@ -2346,6 +2651,7 @@ mod tests { total_bytes: 48, kv_bytes: 48, recurrent_bytes: 0, + kv_desc_json: None, segments: vec![(stolen, 0..24), (segment_digest(&w2[24..]), 24..48)], }, }; diff --git a/crates/skippy-prompt/src/prompt_cli/stage_config.rs b/crates/skippy-prompt/src/prompt_cli/stage_config.rs index 84964bbedc..954541104a 100644 --- a/crates/skippy-prompt/src/prompt_cli/stage_config.rs +++ b/crates/skippy-prompt/src/prompt_cli/stage_config.rs @@ -132,6 +132,7 @@ fn prompt_stage_kv_cache_config( payload: StageKvCachePayload::Auto, max_entries: 1, max_bytes: 0, + l2_max_bytes: 0, min_tokens: args.kv_page_size_tokens.max(1), shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, @@ -145,6 +146,7 @@ fn prompt_stage_kv_cache_config( payload, max_entries: 128, max_bytes, + l2_max_bytes: 0, min_tokens: args.kv_page_size_tokens.max(1), shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-protocol/src/config.rs b/crates/skippy-protocol/src/config.rs index 25afdc68d0..61148b0827 100644 --- a/crates/skippy-protocol/src/config.rs +++ b/crates/skippy-protocol/src/config.rs @@ -339,6 +339,10 @@ pub struct StageKvCacheConfig { pub max_entries: usize, #[serde(default)] pub max_bytes: u64, + /// Hard byte budget for the opt-in host-RAM L2 exact-state tier. + /// Zero keeps L2 disabled. + #[serde(default)] + pub l2_max_bytes: u64, #[serde(default = "default_kv_cache_min_tokens")] pub min_tokens: u64, #[serde(default = "default_kv_cache_shared_stride_tokens")] diff --git a/crates/skippy-server/src/binary_transport/stage_execution.rs b/crates/skippy-server/src/binary_transport/stage_execution.rs index f61a716961..66faf9da68 100644 --- a/crates/skippy-server/src/binary_transport/stage_execution.rs +++ b/crates/skippy-server/src/binary_transport/stage_execution.rs @@ -976,6 +976,7 @@ pub(in crate::binary_transport) fn prefix_cache_test_config() -> StageConfig { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/frontend/local_generation/tests.rs b/crates/skippy-server/src/frontend/local_generation/tests.rs index 72e6842ca2..3cba546ce4 100644 --- a/crates/skippy-server/src/frontend/local_generation/tests.rs +++ b/crates/skippy-server/src/frontend/local_generation/tests.rs @@ -147,6 +147,7 @@ fn recurrent_test_backend( payload: StageKvCachePayload::KvRecurrent, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 0, diff --git a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs index 41d9c2cd95..7ddd6f074e 100644 --- a/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs +++ b/crates/skippy-server/src/frontend/local_generation/token_generation/kv_restore.rs @@ -41,7 +41,7 @@ impl StageOpenAiBackend { "skippy.exact_cache.source".to_string(), json!(restored.source), ); - if restored.source == "l3" { + if restored.source != "radix" { attrs.insert( "skippy.exact_cache.fill_ms".to_string(), json!(restored.fill_ms), diff --git a/crates/skippy-server/src/frontend/prefix_cache.rs b/crates/skippy-server/src/frontend/prefix_cache.rs index 61e9ed46ca..d39695595e 100644 --- a/crates/skippy-server/src/frontend/prefix_cache.rs +++ b/crates/skippy-server/src/frontend/prefix_cache.rs @@ -1372,6 +1372,7 @@ mod tests { payload: skippy_protocol::StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, @@ -1449,6 +1450,7 @@ mod tests { payload: skippy_protocol::StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/crates/skippy-server/src/frontend/tests/support.rs b/crates/skippy-server/src/frontend/tests/support.rs index 1de05d1a68..15d99b9f48 100644 --- a/crates/skippy-server/src/frontend/tests/support.rs +++ b/crates/skippy-server/src/frontend/tests/support.rs @@ -47,6 +47,7 @@ pub(super) fn prefix_cache_test_config() -> StageConfig { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/activation.rs b/crates/skippy-server/src/kv_integration/activation.rs index 3dcf3d66d8..b7f874738e 100644 --- a/crates/skippy-server/src/kv_integration/activation.rs +++ b/crates/skippy-server/src/kv_integration/activation.rs @@ -189,6 +189,7 @@ mod tests { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 256, shared_prefix_stride_tokens: 128, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/cache_affinity.rs b/crates/skippy-server/src/kv_integration/cache_affinity.rs index aa56428c2c..612164f229 100644 --- a/crates/skippy-server/src/kv_integration/cache_affinity.rs +++ b/crates/skippy-server/src/kv_integration/cache_affinity.rs @@ -109,6 +109,7 @@ mod tests { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 64, shared_prefix_stride_tokens: 32, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/config.rs b/crates/skippy-server/src/kv_integration/config.rs index 968a4e5231..fb574935bb 100644 --- a/crates/skippy-server/src/kv_integration/config.rs +++ b/crates/skippy-server/src/kv_integration/config.rs @@ -133,7 +133,16 @@ impl KvStageIntegration { return Ok(None); } let l3_manager = manager()?; - let durable_payload = l3_manager.as_ref().map(|_| { + if cache_config.l2_max_bytes > 0 && l3_manager.is_none() { + let _ = mesh_llm_events::emit_event(OutputEvent::Warning { + message: "Skippy L2 host-RAM cache disabled for this model stage".to_string(), + context: Some(format!( + "stage_id={} model_id={} reason=L2 requires an active L3 cache", + config.stage_id, config.model_id + )), + }); + } + let durable_payload = l3_manager.is_some().then(|| { if payload == StagePrefixCachePayload::ResidentKv && dense_without_recurrent { // Resident KV stays the in-process fast path. Dense families // export KV pages with an empty recurrent snapshot for L3; @@ -147,6 +156,17 @@ impl KvStageIntegration { .zip(durable_payload) .map(|(manager, payload)| l3_tier_for_manager(config, payload, manager)) .transpose()?; + let l2 = l3 + .as_ref() + .filter(|_| cache_config.l2_max_bytes > 0) + .and(durable_payload) + .map(|payload| { + super::l2_serving::StageL2::new( + cache_config.l2_max_bytes, + numerical_model_identity_for_stage(config), + exact_state_identity_for_stage(config, l3_payload_kind(payload)), + ) + }); // FullState is architecture-neutral: the native runtime serializes the // complete session state for both dense and recurrent model families. if matches!(model_capability, ModelKvCapability::KnownRecurrent) { @@ -178,6 +198,7 @@ impl KvStageIntegration { let worker_radix = radix.clone(); let worker_exact_blobs = exact_blobs.clone(); let worker_l3 = l3.clone(); + let worker_l2 = l2.clone(); let inflight_records: Arc>> = l3.as_ref().map_or_else( || Arc::new(Mutex::new(BTreeSet::new())), |tier| tier.manager().record_claims(tier.state_identity()), @@ -222,6 +243,7 @@ impl KvStageIntegration { &worker_exact_blobs, exact_max_entries, exact_byte_limits, + worker_l2.as_ref(), worker_l3.as_deref(), pending, ) @@ -280,6 +302,7 @@ impl KvStageIntegration { split_prefill_tokens: Arc::new(Mutex::new(BTreeMap::new())), kv_lifecycle_observer: observer, exact_state_record_queue_bytes, + l2, l3, inflight_fills, dense_without_recurrent, @@ -535,6 +558,7 @@ fn store_exact_radix_record( blobs: &Mutex, max_entries: usize, limits: ExactStateByteLimits, + l2: Option<&super::l2_serving::StageL2>, l3: Option<&L3Tier>, pending: PendingExactStateRecord, ) -> Result<()> { @@ -543,7 +567,9 @@ fn store_exact_radix_record( // or failing disk must not fail the in-memory record. The refusal reason // lands in the tier's status; one warning per process keeps a full disk // from flooding the log. - if let Some(l3) = l3 { + if pending.write_through_l3 + && let Some(l3) = l3 + { let kv_desc_json = pending .extra .kv_desc @@ -573,6 +599,15 @@ fn store_exact_radix_record( } } } + if let (Some(l2), Some(payload_digest)) = (l2, pending.l2_promotion_digest.as_deref()) { + let _ = l2.promote( + &pending.namespace, + &pending.token_ids, + payload_digest, + &pending.payload, + &pending.extra, + ); + } let logical_bytes = pending.payload.byte_len(); let (payload, _) = pending.payload.dedupe_into( &mut blobs @@ -785,6 +820,7 @@ fn effective_cache_config(config: &StageConfig) -> Option { payload, max_entries, max_bytes, + l2_max_bytes: 0, min_tokens, shared_prefix_stride_tokens, shared_prefix_record_limit, @@ -833,6 +869,8 @@ mod tests { namespace: "model".to_string(), token_ids: tokens.to_vec(), l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, } } @@ -847,6 +885,7 @@ mod tests { 1, limits(0, 0), None, + None, pending("first", &[1, 2], b"aaaabbbb"), ) .unwrap(); @@ -856,6 +895,7 @@ mod tests { 1, limits(0, 0), None, + None, pending("second", &[1, 3], b"aaaacccc"), ) .unwrap(); @@ -875,6 +915,28 @@ mod tests { ); } + #[test] + fn selected_l3_fill_promotes_to_l2_on_the_record_worker_path() { + let radix = Mutex::new(UnifiedRadixCache::new()); + let blobs = Mutex::new(CacheBlobStore::new(4)); + let l2 = super::super::l2_serving::StageL2::new( + 1 << 20, + "model-identity".to_string(), + "state-identity".to_string(), + ); + let bytes = b"filled-state"; + let mut record = pending("filled", &[1, 2], bytes); + record.write_through_l3 = false; + record.l2_promotion_digest = Some(skippy_cache::segment_digest(bytes)); + + store_exact_radix_record(&radix, &blobs, 1, limits(0, 0), Some(&l2), None, record).unwrap(); + + let stats = l2.stats(); + assert_eq!(stats.entries, 1); + assert_eq!(stats.inserts, 1); + assert_eq!(stats.logical_bytes, bytes.len() as u64); + } + #[test] fn invalid_exact_radix_key_releases_deduped_payload() { let radix = Mutex::new(UnifiedRadixCache::new()); @@ -886,6 +948,7 @@ mod tests { 1, limits(0, 0), None, + None, pending("empty", &[], b"aaaabbbb"), ) .unwrap_err(); @@ -914,6 +977,7 @@ mod tests { &blobs, 1, limits(0, 0), + None, Some(&tier), pending("first", &[1, 2], b"first-exact-state"), ) @@ -923,6 +987,7 @@ mod tests { &blobs, 1, limits(0, 0), + None, Some(&tier), pending("second", &[1, 3], b"second-exact-state"), ) @@ -1076,6 +1141,7 @@ mod tests { 8, limits(4, 1024), None, + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -1107,6 +1173,7 @@ mod tests { 2, limits(4, 1_024), None, + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -1135,6 +1202,7 @@ mod tests { 8, limits(4, 8), None, + None, pending(page_id, &tokens, bytes), ) .unwrap(); @@ -1161,6 +1229,7 @@ mod tests { 8, limits(2, 4), None, + None, pending("checkpoint", &[1, 2], b"aaaabbbb"), ) .unwrap(); @@ -1402,6 +1471,58 @@ mod tests { assert_eq!(invalid.exact_state_payload(), None); } + #[test] + fn cache_ram_budget_enables_stage_l2_with_disk_authority() { + let mut config = enabled_auto_config("future/model"); + config + .kv_cache + .as_mut() + .expect("enabled cache config") + .l2_max_bytes = 64 * 1024 * 1024; + let root = tempfile::tempdir().unwrap(); + let manager = L3CacheManager::acquire(root.path(), StoreLimits::new(1 << 30, 0)).unwrap(); + + let kv = KvStageIntegration::from_loaded_model_with_l3_manager( + &config, + Some(ModelStateKind::Dense), + Some(manager), + None, + ) + .unwrap() + .expect("prefix cache should remain enabled"); + + assert!(kv.l2.is_some()); + assert!(kv.l3.is_some()); + let attrs = kv.attrs().into_iter().collect::>(); + assert_eq!(attrs["skippy.kv.l2.enabled"], serde_json::json!(true)); + assert_eq!( + attrs["skippy.kv.l2.budget_bytes"], + serde_json::json!(64 * 1024 * 1024u64) + ); + } + + #[test] + fn cache_ram_budget_does_not_create_an_authority_free_l2() { + let mut config = enabled_auto_config("future/model"); + config + .kv_cache + .as_mut() + .expect("enabled cache config") + .l2_max_bytes = 64 * 1024 * 1024; + + let kv = KvStageIntegration::from_loaded_model_with_l3_manager( + &config, + Some(ModelStateKind::Dense), + None, + None, + ) + .unwrap() + .expect("L1 prefix cache should remain enabled"); + + assert!(kv.l2.is_none()); + assert!(kv.l3.is_none()); + } + #[test] fn parses_cache_mode_and_payload_aliases() { assert_eq!( @@ -1483,6 +1604,7 @@ mod tests { payload: StageKvCachePayload::Auto, max_entries: 512, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/crates/skippy-server/src/kv_integration/exact_state.rs b/crates/skippy-server/src/kv_integration/exact_state.rs index 1d5047cef6..41d137755c 100644 --- a/crates/skippy-server/src/kv_integration/exact_state.rs +++ b/crates/skippy-server/src/kv_integration/exact_state.rs @@ -88,7 +88,7 @@ impl KvStageIntegration { (lookup, entries) }; let Some(lookup) = lookup else { - // Radix miss: the durable tier may still hold this prefix. + // The durable tiers may still hold this prefix. // Runs inside the restore transaction, so a failed import // rolls the lane back exactly as a radix restore would. if let Some(restored) = @@ -395,6 +395,8 @@ impl KvStageIntegration { namespace: identity.namespace.clone(), token_ids: identity.token_ids.clone(), l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, }) { ExactStateRecordAdmission::Queued => { // Recording owns the radix/blob locks while it hashes a potentially @@ -462,6 +464,13 @@ impl KvStageIntegration { // recorded the reason for the status surface. Ok(None) | Err(_) => return Ok(None), }; + // L3 remains the authority for L2: locate and validate the current + // manifest identity before a host-RAM mirror may serve the request. + if let Some(restored) = + self.restore_from_l2(runtime, session_id, identity, &location, lookup_started)? + { + return Ok(Some(restored)); + } // Segment and manifest digests intentionally deduplicate bytes across // numerical states. A fill claim must not: one state's fill cannot // warm another state's radix namespace, even when their payload bytes @@ -511,8 +520,14 @@ impl KvStageIntegration { // A load failure (corrupt segment, now quarantined) is a miss, not a // request failure. Import failures below do propagate: the transaction // rolls the lane back and the caller falls back to cold prefill. - let Ok(fill) = l3.load(location) else { - return Ok(None); + let fill = match l3.load(location) { + Ok(fill) => fill, + Err(_) => { + if let Some(l2) = &self.l2 { + l2.invalidate_digest(&location.manifest_key); + } + return Ok(None); + } }; if fill.payload.byte_len() == 0 { return Ok(None); @@ -607,6 +622,10 @@ impl KvStageIntegration { // Re-warm the RAM tier off the request path. A drop is fine: the // disk copy stays authoritative. The fill claim rides along so the // worker releases it only once the entry is radix-resident. + let l2_promotion_digest = self.l2.as_ref().and_then(|l2| { + l2.consider_l3_fill(&location.manifest_key, token_count, fill.payload.byte_len()) + .then(|| location.manifest_key.clone()) + }); let admission = self.enqueue_exact_state_record(PendingExactStateRecord { page_id: identity.page_id.clone(), payload: fill.payload, @@ -614,6 +633,8 @@ impl KvStageIntegration { namespace: identity.namespace.clone(), token_ids: identity.token_ids[..token_count as usize].to_vec(), l3_fill_claim: Some(l3_fill_claim_key(l3, location)), + write_through_l3: false, + l2_promotion_digest, }); let rewarm_enqueued = matches!(admission, ExactStateRecordAdmission::Queued); Ok(Some(ExactStateRestore { diff --git a/crates/skippy-server/src/kv_integration/identity.rs b/crates/skippy-server/src/kv_integration/identity.rs index 71fa2bb0f9..16f5e045f2 100644 --- a/crates/skippy-server/src/kv_integration/identity.rs +++ b/crates/skippy-server/src/kv_integration/identity.rs @@ -225,6 +225,7 @@ mod tests { payload: StageKvCachePayload::ResidentKv, max_entries: 8, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 64, shared_prefix_stride_tokens: 32, shared_prefix_record_limit: 2, diff --git a/crates/skippy-server/src/kv_integration/l2_serving.rs b/crates/skippy-server/src/kv_integration/l2_serving.rs new file mode 100644 index 0000000000..1ae1e84796 --- /dev/null +++ b/crates/skippy-server/src/kv_integration/l2_serving.rs @@ -0,0 +1,437 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use anyhow::{Context, Result}; +use skippy_cache::{L2Hit, L2Origin, L2Stats, L2Tier, l2_cache_key}; +use skippy_runtime::RuntimeKvPageDesc; + +use crate::runtime_state::RuntimeState; + +use super::{ + ExactStateExtra, ExactStateRecordAdmission, ExactStateRestore, KvStageIntegration, + PendingExactStateRecord, PrefillKvIdentity, StagePrefixCachePayload, + records::add_reconstruct_stats, +}; + +const PROMOTION_MAX_BYTES: u64 = 64 * 1024 * 1024; +const RESTORE_WORTH_TOKENS: u64 = 4_096; +const SECOND_HIT_WINDOW: Duration = Duration::from_secs(10 * 60); +const MAX_PROMOTION_OBSERVATIONS: usize = 4_096; + +#[derive(Clone, Copy)] +struct PromotionObservation { + first_seen: Instant, + hits: u8, + terminal: bool, +} + +/// Stage-scoped view of the host-RAM tier. The identities are fixed when the +/// model loads, so every lookup and promotion uses the same numerical boundary +/// as L3 without exposing model or prompt fingerprints in telemetry. +#[derive(Clone)] +pub(crate) struct StageL2 { + tier: Arc, + model_identity: String, + state_identity: String, + promotions: Arc>>, +} + +impl StageL2 { + pub(crate) fn new(budget_bytes: u64, model_identity: String, state_identity: String) -> Self { + Self { + tier: Arc::new(L2Tier::new(budget_bytes)), + model_identity, + state_identity, + promotions: Arc::new(Mutex::new(HashMap::new())), + } + } + + pub(crate) fn get( + &self, + identity: &PrefillKvIdentity, + location: &skippy_cache::L3Location, + ) -> Option<(String, L2Hit)> { + let token_count = usize::try_from(location.token_count).ok()?; + let tokens = identity.token_ids.get(..token_count)?; + let key = self.cache_key(&identity.namespace, tokens); + let hit = self.tier.get(&key)?; + if hit.token_count != location.token_count || hit.payload_digest != location.manifest_key { + self.remove(&key); + return None; + } + Some((key, hit)) + } + + pub(crate) fn remove(&self, key: &str) { + let _ = self.tier.remove(key); + } + + pub(crate) fn invalidate_digest(&self, payload_digest: &str) { + let _ = self.tier.remove_by_digest(payload_digest); + } + + pub(crate) fn consider_l3_fill( + &self, + manifest_key: &str, + token_count: u64, + payload_bytes: u64, + ) -> bool { + self.should_promote(manifest_key, token_count, payload_bytes) + } + + pub(crate) fn promote( + &self, + namespace: &str, + tokens: &[i32], + expected_payload_digest: &str, + payload: &skippy_cache::ExactStatePayload, + extra: &ExactStateExtra, + ) -> bool { + let key = self.cache_key(namespace, tokens); + let kv_desc_json = extra + .kv_desc + .as_ref() + .and_then(|desc| serde_json::to_string(desc).ok()); + let admitted = self + .tier + .admit_payload( + key, + tokens.len() as u64, + expected_payload_digest, + payload, + kv_desc_json, + L2Origin::FromL3, + ) + .is_ok(); + if let Some(observation) = self + .promotions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get_mut(expected_payload_digest) + { + observation.terminal = true; + } + admitted + } + + pub(crate) fn stats(&self) -> L2Stats { + self.tier.stats() + } + + fn cache_key(&self, namespace: &str, tokens: &[i32]) -> String { + l2_cache_key( + &self.model_identity, + &self.state_identity, + namespace, + tokens, + ) + } + + fn should_promote(&self, manifest_key: &str, token_count: u64, payload_bytes: u64) -> bool { + let now = Instant::now(); + let mut observations = self + .promotions + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + observations.retain(|_, observation| { + now.saturating_duration_since(observation.first_seen) <= SECOND_HIT_WINDOW + }); + if observations.len() >= MAX_PROMOTION_OBSERVATIONS + && !observations.contains_key(manifest_key) + && let Some(oldest) = observations + .iter() + .min_by(|left, right| { + left.1 + .first_seen + .cmp(&right.1.first_seen) + .then_with(|| left.0.cmp(right.0)) + }) + .map(|(key, _)| key.clone()) + { + observations.remove(&oldest); + } + let observation = + observations + .entry(manifest_key.to_string()) + .or_insert(PromotionObservation { + first_seen: now, + hits: 0, + terminal: false, + }); + if observation.terminal { + return false; + } + observation.hits = observation.hits.saturating_add(1); + if payload_bytes > PROMOTION_MAX_BYTES { + observation.terminal = true; + return false; + } + // Keep the observation eligible until the worker has attempted the + // promotion. If the bounded worker queue drops this fill, the next L3 + // restore may retry instead of suppressing the entry for ten minutes. + token_count >= RESTORE_WORTH_TOKENS || observation.hits >= 2 + } +} + +impl KvStageIntegration { + pub(super) fn restore_from_l2( + &self, + runtime: &mut RuntimeState, + session_id: &str, + identity: &PrefillKvIdentity, + location: &skippy_cache::L3Location, + lookup_started: Instant, + ) -> Result> { + let Some(l2) = &self.l2 else { + return Ok(None); + }; + let Some((key, hit)) = l2.get(identity, location) else { + return Ok(None); + }; + let token_count = hit.token_count; + if token_count == 0 || token_count != location.token_count { + l2.remove(&key); + return Ok(None); + } + let payload = hit.to_payload(); + let kv_desc = match hit.payload.kv_desc_json() { + Some(json) => match serde_json::from_str::(json) { + Ok(desc) => Some(desc), + Err(_) => { + l2.remove(&key); + return Ok(None); + } + }, + None => None, + }; + let fill_ms = lookup_started.elapsed().as_secs_f64() * 1000.0; + let mut reconstruct_ms = 0.0; + let mut reconstruct_bytes = 0u64; + let mut reconstruct_blocks = 0usize; + let mut kv_import_ms = 0.0; + let mut recurrent_import_ms = 0.0; + let mut deterministic_failure = false; + let restore = (|| -> Result { + match payload.kind().into() { + StagePrefixCachePayload::FullState => { + let (bytes, stats) = payload + .full_state_bytes_timed() + .context("reconstruct L2 full-state payload") + .map_err(|error| mark_failure(&mut deterministic_failure, error))?; + if bytes.is_empty() { + deterministic_failure = true; + anyhow::bail!("L2 full-state payload is empty"); + } + add_reconstruct_stats( + &mut reconstruct_ms, + &mut reconstruct_bytes, + &mut reconstruct_blocks, + stats, + ); + let started = Instant::now(); + runtime.import_full_state_for_token_count( + session_id, + bytes.as_ref(), + token_count, + )?; + kv_import_ms = started.elapsed().as_secs_f64() * 1000.0; + } + StagePrefixCachePayload::KvRecurrent => { + if let Some((kv, stats)) = payload + .kv_bytes_timed() + .context("reconstruct L2 KV payload") + .map_err(|error| mark_failure(&mut deterministic_failure, error))? + { + add_reconstruct_stats( + &mut reconstruct_ms, + &mut reconstruct_bytes, + &mut reconstruct_blocks, + stats, + ); + match kv_desc.as_ref() { + Some(desc) => { + desc.validate_payload(kv.len()).map_err(|error| { + mark_failure(&mut deterministic_failure, error) + })?; + if desc.token_start != 0 || desc.token_count != token_count { + deterministic_failure = true; + anyhow::bail!("L2 KV page token range mismatch"); + } + let started = Instant::now(); + runtime.import_kv_page(session_id, desc, kv.as_ref())?; + kv_import_ms = started.elapsed().as_secs_f64() * 1000.0; + } + None if !kv.is_empty() => { + deterministic_failure = true; + anyhow::bail!("L2 KV payload is missing its descriptor"); + } + None => {} + } + } + let (recurrent, stats) = payload + .recurrent_state_bytes_timed() + .context("reconstruct L2 recurrent payload") + .map_err(|error| mark_failure(&mut deterministic_failure, error))?; + if recurrent.is_empty() && !self.dense_without_recurrent { + deterministic_failure = true; + anyhow::bail!("L2 recurrent-state payload is empty"); + } + add_reconstruct_stats( + &mut reconstruct_ms, + &mut reconstruct_bytes, + &mut reconstruct_blocks, + stats, + ); + let started = Instant::now(); + if recurrent.is_empty() { + runtime.set_session_position(session_id, token_count)?; + } else { + runtime.import_recurrent_state_for_token_count( + session_id, + recurrent.as_ref(), + token_count, + )?; + } + recurrent_import_ms = started.elapsed().as_secs_f64() * 1000.0; + } + StagePrefixCachePayload::Disabled | StagePrefixCachePayload::ResidentKv => { + return Ok(false); + } + } + Ok(true) + })(); + let restored = match restore { + Ok(restored) => restored, + Err(error) => { + if deterministic_failure { + l2.remove(&key); + } + return Err(error); + } + }; + if !restored { + return Ok(None); + } + let logical_bytes = payload.byte_len(); + let payload_kind = payload.kind(); + let rewarm_enqueued = self.try_begin_record(&identity.page_id) + && matches!( + self.enqueue_exact_state_record(PendingExactStateRecord { + page_id: identity.page_id.clone(), + payload, + extra: ExactStateExtra { kv_desc }, + namespace: identity.namespace.clone(), + token_ids: identity.token_ids[..token_count as usize].to_vec(), + l3_fill_claim: None, + write_through_l3: false, + l2_promotion_digest: None, + }), + ExactStateRecordAdmission::Queued + ); + Ok(Some(ExactStateRestore { + page_id: identity.page_id.clone(), + token_count: token_count as usize, + payload_kind, + logical_bytes, + entries: l2.stats().entries as usize, + reconstruct_ms, + reconstruct_bytes, + reconstruct_blocks, + lookup_ms: fill_ms, + kv_import_ms, + recurrent_import_ms, + source: "l2", + fill_ms, + rewarm_enqueued, + })) + } +} + +fn mark_failure(deterministic: &mut bool, error: anyhow::Error) -> anyhow::Error { + *deterministic = true; + error +} + +#[cfg(test)] +mod tests { + use super::*; + + fn identity(tokens: Vec) -> PrefillKvIdentity { + PrefillKvIdentity { + identity: crate::kv_proto::PageIdentity::default(), + page_id: "page".to_string(), + namespace: "namespace".to_string(), + token_ids: tokens, + } + } + + fn location(token_count: u64, manifest_key: String) -> skippy_cache::L3Location { + skippy_cache::L3Location { + namespace_key: "namespace-key".to_string(), + prefix_key: "prefix-key".to_string(), + token_count, + manifest_key, + kv_desc_json: None, + kv_bytes: 0, + native_kv_passthrough: false, + } + } + + #[test] + fn second_l3_fill_promotes_and_serves_the_located_prefix() { + let l2 = StageL2::new(1 << 20, "model".to_string(), "state".to_string()); + let bytes = vec![7; 128]; + let digest = skippy_cache::segment_digest(&bytes); + let payload = skippy_cache::ExactStatePayload::full_state(bytes); + let identity = identity(vec![1, 2, 3, 4]); + let location = location(3, digest.clone()); + + assert!(!l2.consider_l3_fill(&digest, 3, payload.byte_len())); + assert!(l2.consider_l3_fill(&digest, 3, payload.byte_len())); + assert!(l2.promote( + &identity.namespace, + &identity.token_ids[..3], + &digest, + &payload, + &ExactStateExtra { kv_desc: None }, + )); + assert!(!l2.consider_l3_fill(&digest, 3, payload.byte_len())); + + let (_, hit) = l2 + .get(&identity, &location) + .expect("longer query should hit the located L3 prefix mirror"); + assert_eq!(hit.token_count, 3); + assert_eq!(hit.payload_digest, digest); + } + + #[test] + fn restart_class_prefix_promotes_on_first_fill_but_oversized_payload_does_not() { + let l2 = StageL2::new(1 << 20, "model".to_string(), "state".to_string()); + assert!(l2.consider_l3_fill("restart", RESTORE_WORTH_TOKENS, 1 << 20)); + assert!(!l2.consider_l3_fill("oversized", RESTORE_WORTH_TOKENS, PROMOTION_MAX_BYTES + 1,)); + assert!(!l2.consider_l3_fill("oversized", RESTORE_WORTH_TOKENS, PROMOTION_MAX_BYTES,)); + } + + #[test] + fn a_changed_durable_digest_invalidates_the_mirror() { + let l2 = StageL2::new(1 << 20, "model".to_string(), "state".to_string()); + let bytes = vec![7; 128]; + let digest = skippy_cache::segment_digest(&bytes); + let payload = skippy_cache::ExactStatePayload::full_state(bytes); + let identity = identity(vec![1, 2, 3]); + assert!(l2.promote( + &identity.namespace, + &identity.token_ids, + &digest, + &payload, + &ExactStateExtra { kv_desc: None }, + )); + + let changed = location(3, skippy_cache::segment_digest(b"changed")); + assert!(l2.get(&identity, &changed).is_none()); + assert_eq!(l2.stats().entries, 0); + } +} diff --git a/crates/skippy-server/src/kv_integration/mod.rs b/crates/skippy-server/src/kv_integration/mod.rs index 93db8aa401..26123ad372 100644 --- a/crates/skippy-server/src/kv_integration/mod.rs +++ b/crates/skippy-server/src/kv_integration/mod.rs @@ -27,6 +27,7 @@ mod cache_affinity; mod config; mod exact_state; mod identity; +mod l2_serving; pub mod lifecycle; mod model_capability; mod output_tokens; @@ -185,6 +186,9 @@ pub struct KvStageIntegration { /// lets one multi-GiB export sit next to another and doubles the RAM the /// cache can pin behind a request. pub(crate) exact_state_record_queue_bytes: Arc, + /// Optional bounded host-RAM tier. Qualified L3 fills enter L2; an L2 hit + /// promotes back into L1 through the existing worker. + pub(crate) l2: Option, /// Durable L3 floor under the radix cache: exact-state records write /// through to it on the worker, and radix misses fill back from it. pub(crate) l3: Option>, @@ -229,6 +233,12 @@ pub(crate) struct PendingExactStateRecord { /// so requests arriving during the asynchronous re-warm prefill normally /// instead of duplicating the disk read. pub(crate) l3_fill_claim: Option, + /// Only freshly exported request state writes through. Tier fills that + /// merely re-warm L1 must not rewrite their existing durable entry. + pub(crate) write_through_l3: bool, + /// Durable manifest digest to mirror into L2 on this worker job. `None` + /// leaves the payload out of L2. + pub(crate) l2_promotion_digest: Option, } #[derive(Debug)] @@ -793,6 +803,11 @@ impl KvStageIntegration { Err(std::sync::TryLockError::WouldBlock) => None, }; let radix = radix_stats.unwrap_or_default(); + let l2 = self + .l2 + .as_ref() + .map(|tier| tier.stats()) + .unwrap_or_default(); let activations = self .activations .lock() @@ -951,6 +966,21 @@ impl KvStageIntegration { "skippy.exact_cache.max_entries", json!(self.exact_max_entries), ), + ("skippy.kv.l2.enabled", json!(self.l2.is_some())), + ("skippy.kv.l2.budget_bytes", json!(l2.budget_bytes)), + ("skippy.kv.l2.bytes", json!(l2.bytes)), + ("skippy.kv.l2.logical_bytes", json!(l2.logical_bytes)), + ("skippy.kv.l2.entries", json!(l2.entries)), + ("skippy.kv.l2.segments", json!(l2.segments)), + ("skippy.kv.l2.hits", json!(l2.hits)), + ("skippy.kv.l2.misses", json!(l2.misses)), + ("skippy.kv.l2.inserts", json!(l2.inserts)), + ("skippy.kv.l2.evictions", json!(l2.evictions)), + ( + "skippy.kv.l2.admission_rejects", + json!(l2.admission_rejects), + ), + ("skippy.kv.l2.refused_bytes", json!(l2.refused_bytes)), ( "skippy.kv.output_token_entries", json!(output_token_entries), @@ -1140,6 +1170,8 @@ mod exact_state_record_queue_tests { namespace: "test".to_string(), token_ids: vec![1], l3_fill_claim: None, + write_through_l3: true, + l2_promotion_digest: None, } } diff --git a/crates/skippy-server/src/kv_integration/resident_prefix.rs b/crates/skippy-server/src/kv_integration/resident_prefix.rs index cf35b4c2d0..05a4dd75cc 100644 --- a/crates/skippy-server/src/kv_integration/resident_prefix.rs +++ b/crates/skippy-server/src/kv_integration/resident_prefix.rs @@ -824,6 +824,7 @@ mod proactive_eviction_tests { payload: StageKvCachePayload::ResidentKv, max_entries: 4, max_bytes: 0, + l2_max_bytes: 0, min_tokens: 1, shared_prefix_stride_tokens: 1, shared_prefix_record_limit: 1, diff --git a/docs/USAGE.md b/docs/USAGE.md index 23da197d01..e13044376c 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -468,7 +468,7 @@ kv_cache_policy = "balanced" # macro preset: auto quality balanced saver # explicit cache_type_k/v always wins over preset kv_offload = "auto" # bool or "auto" — KV residency / offload policy kv_unified = "auto" # bool or "auto" — unified KV layout (schema-reserved) -cache_ram_mib = 0 # byte cap for KV cache in MiB; 0 = no cap (schema-reserved) +cache_ram_mib = 0 # host-RAM L2 budget in MiB; 0 = disabled; requires L3 cache_idle_slots = 0 # idle slot retention count (schema-reserved) prompt_cache = "auto" # bool or "auto" — reuse previous prompt KV swa_full = false # sliding-window attention (model-family specific) diff --git a/docs/skippy/CONFIGURATION.md b/docs/skippy/CONFIGURATION.md index a210406aa2..cd3f28c851 100644 --- a/docs/skippy/CONFIGURATION.md +++ b/docs/skippy/CONFIGURATION.md @@ -91,7 +91,7 @@ website configuration reference, with the same `Wiring status`. | 5.1 | KV cache policy preset | `model_fit.kv_cache_policy` | P0 | `plugin/config.rs` | policy expander into cache_type_k, cache_type_v, kv_offload, cache_ram_mib | single-stage, staged | restart/reload only | mesh policy default | enum auto, quality, balanced, saver; explicit cache_type_k/v wins over preset expansion | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | quality=f16/f16 with no forced RAM cap; balanced=preserve runtime defaults; saver=prefer lower-memory dtypes plus offload warning if unsupported; auto=family or topology policy decides | wired | | 5.1 | KV offload | `model_fit.kv_offload` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, native `skippy_runtime_config.kv_offload` tri-state | single-stage, staged | restart/reload only | backend runtime default or kv_cache_policy expansion | boolean or auto; auto preserves llama.cpp's derived `offload_kqv` default | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; capture tests in `crates/skippy-runtime/src/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | controls KV residency and offload policy | wired | | 5.1 | Unified KV cache | `model_fit.kv_unified` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, native `skippy_runtime_config.kv_unified` tri-state | single-stage, staged | restart/reload only | backend runtime default | boolean or auto; recurrent/hybrid architectures still force this true natively regardless of the requested value | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; capture tests in `crates/skippy-runtime/src/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | advanced cache layout flag | wired | -| 5.1 | Cache RAM budget | `model_fit.cache_ram_mib` | P1 | `plugin/config.rs` | Schema-reserved only; resolver fail-closed today | single-stage, staged | restart/reload only | unset by default | integer >= 0 MiB; zero or unset means no forced cap; not executable today | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | caps cache memory when runtime path supports it; schema-reserved until a real executable surface exists | unwired | +| 5.1 | Cache RAM budget | `model_fit.cache_ram_mib` | P1 | `plugin/config.rs`, `skippy-server`, `skippy-cache` | `StageKvCacheConfig.l2_max_bytes` and bounded host-RAM L2 exact-state tier | single-stage, staged | restart/reload only | disabled by default | integer >= 0 MiB; zero or unset disables L2; positive values require prefix caching and active L3 | `#model-fit-context-and-kv-cache` | resolver propagation tests in `crates/mesh-llm-host-runtime/src/inference/skippy/resolver/tests.rs`; promotion and admission tests in `skippy-server` and `skippy-cache`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | repeated or high-value L3 fills promote to L2; L2 hits restore before disk and asynchronously rewarm L1 | wired | | 5.1 | Cache idle slots | `model_fit.cache_idle_slots` | P1 | `plugin/config.rs` | SkippyModelLoadOptions, StageConfig, `skippy-server::RuntimeState::max_idle_sessions` idle-pool bound | single-stage, staged | restart/reload only | runtime default (unbounded, capped only by `lane_count`) | integer >= 0; bounds the idle session pool size | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; `crates/skippy-server/src/runtime_state.rs` and `crates/skippy-server/src/runtime_state/lane_lifecycle.rs` unit tests; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | bounds retained idle sessions so `drop_session_timed` discards lanes past the configured cap instead of growing the pool unbounded | wired | | 5.1 | Prompt cache / cache prompt | `model_fit.prompt_cache` | P1 | `plugin/config.rs` | cache config or request adapter defaults when executable | single-stage, staged | restart/reload only | disabled unless operator enables reuse | boolean only; reject when selected runtime does not expose prompt-cache behavior | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | request prompt reuse toggle | wired | | 5.1 | Prefix cache enable | `model_fit.prefix_cache.enabled` | P1 | `plugin/config.rs` | cache config or SKIPPY_PREFIX_CACHE bridge | single-stage, staged | restart/reload only | cache env default or disabled | boolean only; when false, ignore remaining prefix_cache fields | `#model-fit-context-and-kv-cache` | `plugin::config` tests in `crates/mesh-llm-host-runtime/src/plugin/config.rs`; parity: `website_docs_parity` in `crates/mesh-llm-config/src/website_docs_parity.rs` | top-level exact-prefix cache toggle | wired | diff --git a/docs/skippy/PROMPT_CACHE.md b/docs/skippy/PROMPT_CACHE.md index b17f9a204c..b5356722c9 100644 --- a/docs/skippy/PROMPT_CACHE.md +++ b/docs/skippy/PROMPT_CACHE.md @@ -68,6 +68,27 @@ limit. Both limits are reported on `stage.openai_generation_summary` as `skippy.exact_cache.max_bytes` and `skippy.exact_cache.hard_max_bytes`. +## Host-RAM L2 Cache + +Set `model_fit.cache_ram_mib` to a positive MiB value to enable the bounded +host-RAM exact-state tier for that model. The default value, `0`, leaves L2 +disabled. Prefix caching and the node-local L3 cache must also be enabled. + +Exact-state lookup proceeds from the in-process radix cache (L1), to host RAM +(L2), then to the node-local disk cache (L3). L3 remains authoritative: the +server locates the current durable manifest before serving an L2 mirror and +requires the mirror digest to match it. The cache worker promotes an L3 fill +on its second hit within ten minutes, or on the first hit for prefixes of at +least 4,096 tokens. Payloads larger than 64 MiB stay out of L2. A verified L2 +hit restores the request immediately and queues the same payload to rewarm L1. +Rewarm records do not rewrite an existing L3 entry. Unloading the stage drops +its L2 tier. + +`stage.openai_generation_summary` reports `skippy.kv.l2.enabled`, budget and +resident byte counts, logical bytes, entries, segments, hits, misses, inserts, +evictions, and admission refusals. Exact-hit telemetry identifies the restore +source as `l2` and includes fill time and whether the L1 rewarm was queued. + ## mesh-llm Defaults mesh-llm wires Skippy prefix cache through family policy. For supported model diff --git a/tools/xtask/data/console_print_allowlist.json b/tools/xtask/data/console_print_allowlist.json index 94c3d891da..6f9e188845 100644 --- a/tools/xtask/data/console_print_allowlist.json +++ b/tools/xtask/data/console_print_allowlist.json @@ -4251,19 +4251,19 @@ ], "crates/skippy-prompt/src/prompt_cli/stage_config.rs": [ { - "line": 358, + "line": 360, "macro_name": "eprintln!" }, { - "line": 398, + "line": 400, "macro_name": "eprintln!" }, { - "line": 464, + "line": 466, "macro_name": "eprintln!" }, { - "line": 478, + "line": 480, "macro_name": "eprint!" } ], diff --git a/website/src/docs/pages/config-reference.md b/website/src/docs/pages/config-reference.md index cd5cdd9c56..91630c91c5 100644 --- a/website/src/docs/pages/config-reference.md +++ b/website/src/docs/pages/config-reference.md @@ -143,7 +143,7 @@ for the activity policy and privacy boundary. | `model_fit.kv_cache_policy` | enum | `balanced` (default), `auto`, `quality`, `saver`; expands into cache dtypes | both | model reload | wired | none | | `model_fit.kv_offload` | bool-or-`auto` | `auto` | both | model reload | wired | none | | `model_fit.kv_unified` | bool-or-`auto` | `auto` | both | model reload | wired (recurrent/hybrid architectures still force this true natively) | none | -| `model_fit.cache_ram_mib` | integer | unset (no cap) | both | model reload | unwired (any positive value fails at model load) | none | +| `model_fit.cache_ram_mib` | integer | `0`/unset = host-RAM L2 disabled | both | model reload | wired; requires prefix caching and active L3 | none | | `model_fit.cache_idle_slots` | integer | unset uses the runtime lane count; `0` drops every reset lane, positive values cap retained idle sessions | both | model reload | wired | none | | `model_fit.prompt_cache` | bool-or-`auto` | `auto` | both | model reload | wired | none | | `model_fit.prefix_cache.enabled` | boolean | unset uses family defaults; `false` disables | both | model reload | wired | none |