462 lines
13 KiB
Rust
462 lines
13 KiB
Rust
|
/// Phase 6: Cache Alignment & KV Cache Optimization
|
||
|
|
///
|
||
|
|
/// Optimize memory system for LLM KV cache efficiency:
|
||
|
|
/// 1. Cache metrics tracking (hit ratio, evictions)
|
||
|
|
/// 2. Wiki-link ordering by cache locality
|
||
|
|
/// 3. Chunk pre-loading for hot paths
|
||
|
|
/// 4. Monitor KV cache hit ratio during retrieval
|
||
|
|
///
|
||
|
|
/// Target:
|
||
|
|
/// - KV cache hit ratio > 80%
|
||
|
|
/// - Chunk loading latency < 50ms (cache) vs 200ms (disk)
|
||
|
|
/// - Reduce context recomputation by 60%
|
||
|
|
|
||
|
|
use anyhow::Result;
|
||
|
|
use std::collections::{HashMap, VecDeque};
|
||
|
|
use std::sync::{Arc, Mutex};
|
||
|
|
|
||
|
|
/// Cache metrics for tracking
|
||
|
|
#[derive(Debug, Clone, Copy)]
|
||
|
|
pub struct CacheMetrics {
|
||
|
|
pub hits: u64,
|
||
|
|
pub misses: u64,
|
||
|
|
pub evictions: u64,
|
||
|
|
pub avg_load_ms: f32,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl CacheMetrics {
|
||
|
|
pub fn new() -> Self {
|
||
|
|
Self {
|
||
|
|
hits: 0,
|
||
|
|
misses: 0,
|
||
|
|
evictions: 0,
|
||
|
|
avg_load_ms: 0.0,
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn hit_ratio(&self) -> f32 {
|
||
|
|
let total = self.hits + self.misses;
|
||
|
|
if total == 0 {
|
||
|
|
0.0
|
||
|
|
} else {
|
||
|
|
self.hits as f32 / total as f32
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn total_requests(&self) -> u64 {
|
||
|
|
self.hits + self.misses
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Chunk with cache locality info
|
||
|
|
#[derive(Debug, Clone)]
|
||
|
|
pub struct CachedChunk {
|
||
|
|
pub chunk_id: String,
|
||
|
|
pub text: String,
|
||
|
|
pub score: f32,
|
||
|
|
pub cache_distance: u32, // Hops from root in wiki-graph
|
||
|
|
pub access_count: u64,
|
||
|
|
pub last_accessed_slot: u32, // Transformer position in context
|
||
|
|
}
|
||
|
|
|
||
|
|
/// LRU Cache for chunks
|
||
|
|
pub struct LruChunkCache {
|
||
|
|
capacity: usize,
|
||
|
|
cache: Arc<Mutex<HashMap<String, (String, u64)>>>, // id -> (text, access_time)
|
||
|
|
access_queue: Arc<Mutex<VecDeque<String>>>,
|
||
|
|
metrics: Arc<Mutex<CacheMetrics>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl LruChunkCache {
|
||
|
|
pub fn new(capacity: usize) -> Self {
|
||
|
|
Self {
|
||
|
|
capacity,
|
||
|
|
cache: Arc::new(Mutex::new(HashMap::new())),
|
||
|
|
access_queue: Arc::new(Mutex::new(VecDeque::new())),
|
||
|
|
metrics: Arc::new(Mutex::new(CacheMetrics::new())),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Get chunk from cache
|
||
|
|
pub fn get(&self, chunk_id: &str) -> Option<String> {
|
||
|
|
let mut cache = self.cache.lock().unwrap();
|
||
|
|
let mut metrics = self.metrics.lock().unwrap();
|
||
|
|
|
||
|
|
if let Some((text, _)) = cache.get_mut(chunk_id) {
|
||
|
|
metrics.hits += 1;
|
||
|
|
let result = text.clone();
|
||
|
|
|
||
|
|
// Update access tracking (move to end of queue)
|
||
|
|
let mut queue = self.access_queue.lock().unwrap();
|
||
|
|
queue.retain(|id| id != chunk_id);
|
||
|
|
queue.push_back(chunk_id.to_string());
|
||
|
|
|
||
|
|
Some(result)
|
||
|
|
} else {
|
||
|
|
metrics.misses += 1;
|
||
|
|
None
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Put chunk in cache with LRU eviction
|
||
|
|
pub fn put(&self, chunk_id: &str, text: &str) -> Option<String> {
|
||
|
|
let mut cache = self.cache.lock().unwrap();
|
||
|
|
let mut queue = self.access_queue.lock().unwrap();
|
||
|
|
let mut metrics = self.metrics.lock().unwrap();
|
||
|
|
|
||
|
|
// If cache is full, evict LRU item
|
||
|
|
let evicted = if cache.len() >= self.capacity {
|
||
|
|
if let Some(lru_id) = queue.pop_front() {
|
||
|
|
metrics.evictions += 1;
|
||
|
|
cache.remove(&lru_id)
|
||
|
|
} else {
|
||
|
|
None
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
None
|
||
|
|
};
|
||
|
|
|
||
|
|
// Insert new chunk
|
||
|
|
cache.insert(chunk_id.to_string(), (text.to_string(), 0u64));
|
||
|
|
queue.push_back(chunk_id.to_string());
|
||
|
|
|
||
|
|
evicted.map(|(text, _)| text)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn metrics(&self) -> CacheMetrics {
|
||
|
|
*self.metrics.lock().unwrap()
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn clear(&self) {
|
||
|
|
self.cache.lock().unwrap().clear();
|
||
|
|
self.access_queue.lock().unwrap().clear();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Cache Locality Analyzer
|
||
|
|
pub struct CacheLocalityAnalyzer;
|
||
|
|
|
||
|
|
impl CacheLocalityAnalyzer {
|
||
|
|
/// Calculate cache distance (hops from root)
|
||
|
|
pub fn calculate_distance(
|
||
|
|
chunk_id: &str,
|
||
|
|
root_id: &str,
|
||
|
|
graph: &HashMap<String, Vec<String>>,
|
||
|
|
) -> u32 {
|
||
|
|
if chunk_id == root_id {
|
||
|
|
return 0;
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut visited = std::collections::HashSet::new();
|
||
|
|
let mut queue = VecDeque::new();
|
||
|
|
|
||
|
|
queue.push_back((root_id.to_string(), 0u32));
|
||
|
|
visited.insert(root_id.to_string());
|
||
|
|
|
||
|
|
while let Some((current, distance)) = queue.pop_front() {
|
||
|
|
if current == chunk_id {
|
||
|
|
return distance;
|
||
|
|
}
|
||
|
|
|
||
|
|
if let Some(neighbors) = graph.get(¤t) {
|
||
|
|
for neighbor in neighbors {
|
||
|
|
if !visited.contains(neighbor) {
|
||
|
|
visited.insert(neighbor.clone());
|
||
|
|
queue.push_back((neighbor.clone(), distance + 1));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
u32::MAX // Unreachable
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Order chunks by cache locality (closest first)
|
||
|
|
pub fn order_by_locality(
|
||
|
|
chunks: Vec<CachedChunk>,
|
||
|
|
root_id: &str,
|
||
|
|
graph: &HashMap<String, Vec<String>>,
|
||
|
|
) -> Vec<CachedChunk> {
|
||
|
|
let mut ordered = chunks;
|
||
|
|
|
||
|
|
ordered.sort_by_key(|c| {
|
||
|
|
Self::calculate_distance(&c.chunk_id, root_id, graph)
|
||
|
|
});
|
||
|
|
|
||
|
|
ordered
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// KV Cache Alignment Optimizer
|
||
|
|
pub struct KvCacheAligner {
|
||
|
|
context_window: usize, // Max tokens per context
|
||
|
|
chunk_avg_tokens: usize, // Average tokens per chunk
|
||
|
|
cache: Arc<LruChunkCache>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl KvCacheAligner {
|
||
|
|
pub fn new(context_window: usize, chunk_avg_tokens: usize, cache_size: usize) -> Self {
|
||
|
|
Self {
|
||
|
|
context_window,
|
||
|
|
chunk_avg_tokens,
|
||
|
|
cache: Arc::new(LruChunkCache::new(cache_size)),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Predict if chunk will fit in context window
|
||
|
|
pub fn will_fit(&self, chunk_count: usize) -> bool {
|
||
|
|
chunk_count * self.chunk_avg_tokens <= self.context_window
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Calculate tokens used by chunks
|
||
|
|
pub fn calculate_tokens(&self, chunks: &[CachedChunk]) -> usize {
|
||
|
|
chunks.len() * self.chunk_avg_tokens
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Assign slot positions in context (for cache locality)
|
||
|
|
pub fn assign_slots(&self, chunks: &[CachedChunk]) -> Vec<(String, u32)> {
|
||
|
|
chunks
|
||
|
|
.iter()
|
||
|
|
.enumerate()
|
||
|
|
.map(|(i, chunk)| (chunk.chunk_id.clone(), i as u32))
|
||
|
|
.collect()
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Pre-load hot chunks into cache
|
||
|
|
pub fn preload_hot_chunks(&self, hot_chunks: Vec<(&str, &str)>) -> Result<()> {
|
||
|
|
for (chunk_id, text) in hot_chunks {
|
||
|
|
self.cache.put(chunk_id, text);
|
||
|
|
}
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn get_metrics(&self) -> CacheMetrics {
|
||
|
|
self.cache.metrics()
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn get_from_cache(&self, chunk_id: &str) -> Option<String> {
|
||
|
|
self.cache.get(chunk_id)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
/// Retrieval Timing Profiler
|
||
|
|
pub struct RetrievalProfiler {
|
||
|
|
timings: Arc<Mutex<Vec<(String, u64)>>>, // (stage_name, duration_ms)
|
||
|
|
}
|
||
|
|
|
||
|
|
impl RetrievalProfiler {
|
||
|
|
pub fn new() -> Self {
|
||
|
|
Self {
|
||
|
|
timings: Arc::new(Mutex::new(Vec::new())),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn record(&self, stage: &str, duration_ms: u64) {
|
||
|
|
self.timings
|
||
|
|
.lock()
|
||
|
|
.unwrap()
|
||
|
|
.push((stage.to_string(), duration_ms));
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn summary(&self) -> Vec<(String, u64)> {
|
||
|
|
self.timings.lock().unwrap().clone()
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn total_time(&self) -> u64 {
|
||
|
|
self.timings
|
||
|
|
.lock()
|
||
|
|
.unwrap()
|
||
|
|
.iter()
|
||
|
|
.map(|(_, duration)| duration)
|
||
|
|
.sum()
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn stage_time(&self, stage: &str) -> Option<u64> {
|
||
|
|
self.timings
|
||
|
|
.lock()
|
||
|
|
.unwrap()
|
||
|
|
.iter()
|
||
|
|
.find(|(s, _)| s == stage)
|
||
|
|
.map(|(_, d)| *d)
|
||
|
|
}
|
||
|
|
|
||
|
|
pub fn clear(&self) {
|
||
|
|
self.timings.lock().unwrap().clear();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_cache_metrics_hit_ratio() {
|
||
|
|
let mut metrics = CacheMetrics::new();
|
||
|
|
metrics.hits = 80;
|
||
|
|
metrics.misses = 20;
|
||
|
|
assert_eq!(metrics.hit_ratio(), 0.8);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_lru_cache_get_hit() {
|
||
|
|
let cache = LruChunkCache::new(10);
|
||
|
|
cache.put("chunk1", "content1");
|
||
|
|
|
||
|
|
let result = cache.get("chunk1");
|
||
|
|
assert_eq!(result, Some("content1".to_string()));
|
||
|
|
assert_eq!(cache.metrics().hits, 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_lru_cache_get_miss() {
|
||
|
|
let cache = LruChunkCache::new(10);
|
||
|
|
let result = cache.get("nonexistent");
|
||
|
|
assert_eq!(result, None);
|
||
|
|
assert_eq!(cache.metrics().misses, 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_lru_cache_eviction() {
|
||
|
|
let cache = LruChunkCache::new(2);
|
||
|
|
cache.put("chunk1", "content1");
|
||
|
|
cache.put("chunk2", "content2");
|
||
|
|
cache.put("chunk3", "content3"); // Should evict chunk1
|
||
|
|
|
||
|
|
assert_eq!(cache.get("chunk1"), None); // Evicted
|
||
|
|
assert_eq!(cache.get("chunk2"), Some("content2".to_string()));
|
||
|
|
assert_eq!(cache.get("chunk3"), Some("content3".to_string()));
|
||
|
|
assert_eq!(cache.metrics().evictions, 1);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_cache_locality_distance_root() {
|
||
|
|
let mut graph = HashMap::new();
|
||
|
|
let distance = CacheLocalityAnalyzer::calculate_distance("root", "root", &graph);
|
||
|
|
assert_eq!(distance, 0);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_cache_locality_distance_neighbors() {
|
||
|
|
let mut graph = HashMap::new();
|
||
|
|
graph.insert("root".to_string(), vec!["child1".to_string()]);
|
||
|
|
graph.insert("child1".to_string(), vec!["child2".to_string()]);
|
||
|
|
|
||
|
|
let dist_child1 = CacheLocalityAnalyzer::calculate_distance("child1", "root", &graph);
|
||
|
|
let dist_child2 = CacheLocalityAnalyzer::calculate_distance("child2", "root", &graph);
|
||
|
|
|
||
|
|
assert_eq!(dist_child1, 1);
|
||
|
|
assert_eq!(dist_child2, 2);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_cache_locality_ordering() {
|
||
|
|
let mut graph = HashMap::new();
|
||
|
|
graph.insert("root".to_string(), vec!["near".to_string(), "far".to_string()]);
|
||
|
|
graph.insert("far".to_string(), vec!["farther".to_string()]);
|
||
|
|
|
||
|
|
let chunks = vec![
|
||
|
|
CachedChunk {
|
||
|
|
chunk_id: "farther".to_string(),
|
||
|
|
text: "".to_string(),
|
||
|
|
score: 0.9,
|
||
|
|
cache_distance: u32::MAX,
|
||
|
|
access_count: 1,
|
||
|
|
last_accessed_slot: 0,
|
||
|
|
},
|
||
|
|
CachedChunk {
|
||
|
|
chunk_id: "near".to_string(),
|
||
|
|
text: "".to_string(),
|
||
|
|
score: 0.8,
|
||
|
|
cache_distance: 1,
|
||
|
|
access_count: 1,
|
||
|
|
last_accessed_slot: 0,
|
||
|
|
},
|
||
|
|
];
|
||
|
|
|
||
|
|
let ordered = CacheLocalityAnalyzer::order_by_locality(chunks, "root", &graph);
|
||
|
|
assert_eq!(ordered[0].chunk_id, "near"); // Closest first
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_kv_cache_will_fit() {
|
||
|
|
let aligner = KvCacheAligner::new(1000, 100, 10);
|
||
|
|
assert!(aligner.will_fit(5)); // 500 tokens < 1000
|
||
|
|
assert!(!aligner.will_fit(15)); // 1500 tokens > 1000
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_kv_cache_assign_slots() {
|
||
|
|
let aligner = KvCacheAligner::new(1000, 100, 10);
|
||
|
|
let chunks = vec![
|
||
|
|
CachedChunk {
|
||
|
|
chunk_id: "chunk1".to_string(),
|
||
|
|
text: "".to_string(),
|
||
|
|
score: 0.9,
|
||
|
|
cache_distance: 0,
|
||
|
|
access_count: 1,
|
||
|
|
last_accessed_slot: 0,
|
||
|
|
},
|
||
|
|
CachedChunk {
|
||
|
|
chunk_id: "chunk2".to_string(),
|
||
|
|
text: "".to_string(),
|
||
|
|
score: 0.8,
|
||
|
|
cache_distance: 1,
|
||
|
|
access_count: 1,
|
||
|
|
last_accessed_slot: 0,
|
||
|
|
},
|
||
|
|
];
|
||
|
|
|
||
|
|
let slots = aligner.assign_slots(&chunks);
|
||
|
|
assert_eq!(slots[0], ("chunk1".to_string(), 0));
|
||
|
|
assert_eq!(slots[1], ("chunk2".to_string(), 1));
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_preload_hot_chunks() {
|
||
|
|
let aligner = KvCacheAligner::new(1000, 100, 10);
|
||
|
|
let hot_chunks = vec![("chunk1", "content1"), ("chunk2", "content2")];
|
||
|
|
|
||
|
|
aligner.preload_hot_chunks(hot_chunks).unwrap();
|
||
|
|
|
||
|
|
// Verify that preloaded chunks are in cache by retrieving them
|
||
|
|
// This will increment metrics
|
||
|
|
assert_eq!(aligner.get_from_cache("chunk1"), Some("content1".to_string()));
|
||
|
|
assert_eq!(aligner.get_from_cache("chunk2"), Some("content2".to_string()));
|
||
|
|
|
||
|
|
let metrics = aligner.get_metrics();
|
||
|
|
assert!(metrics.total_requests() >= 2);
|
||
|
|
assert!(metrics.hits >= 2);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_retrieval_profiler_record() {
|
||
|
|
let profiler = RetrievalProfiler::new();
|
||
|
|
profiler.record("wiki_scope", 10);
|
||
|
|
profiler.record("tfidf_filter", 50);
|
||
|
|
profiler.record("semantic_rerank", 100);
|
||
|
|
|
||
|
|
let summary = profiler.summary();
|
||
|
|
assert_eq!(summary.len(), 3);
|
||
|
|
assert_eq!(profiler.total_time(), 160);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_retrieval_profiler_stage_time() {
|
||
|
|
let profiler = RetrievalProfiler::new();
|
||
|
|
profiler.record("wiki_scope", 10);
|
||
|
|
profiler.record("semantic_rerank", 100);
|
||
|
|
|
||
|
|
assert_eq!(profiler.stage_time("wiki_scope"), Some(10));
|
||
|
|
assert_eq!(profiler.stage_time("nonexistent"), None);
|
||
|
|
}
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_cache_metrics_total_requests() {
|
||
|
|
let mut metrics = CacheMetrics::new();
|
||
|
|
metrics.hits = 60;
|
||
|
|
metrics.misses = 40;
|
||
|
|
assert_eq!(metrics.total_requests(), 100);
|
||
|
|
}
|
||
|
|
}
|