diff --git a/mistralrs-core/src/lib.rs b/mistralrs-core/src/lib.rs index 35ca3c69a2..c1d361b8ff 100644 --- a/mistralrs-core/src/lib.rs +++ b/mistralrs-core/src/lib.rs @@ -155,7 +155,11 @@ pub use mistralrs_mcp::{ }; pub use mistralrs_quant::{IsqBits, IsqType}; pub use mistralrs_sandbox::{NetworkMode, SandboxPolicy}; -pub use paged_attention::{MemoryGpuConfig, PagedAttentionConfig, PagedCacheType}; +pub use paged_attention::{ + compute_block_hashes, BlockHash, BlockHashWithGroupId, InMemoryKvCacheConnector, + KVCacheManager, KvCacheConnector, MemoryGpuConfig, NoopKvCacheConnector, PagedAttentionConfig, + PagedCacheType, +}; pub use pipeline::hf::{ get_model_file, hf_home_dir, hf_hub_cache_dir, hf_token_path, is_hf_hub_offline, list_model_files, probe_hf_repo_files, read_model_file_range, try_get_model_file, @@ -2689,7 +2693,7 @@ impl MistralRs { loader_config.silent, loader_config.device_map_setting.clone(), loader_config.isq, - loader_config.paged_attn_config, + loader_config.paged_attn_config.clone(), ) .map_err(|e| MistralRsError::ReloadFailed(format!("Failed to load model: {e}")))?; diff --git a/mistralrs-core/src/paged_attention/block_pool.rs b/mistralrs-core/src/paged_attention/block_pool.rs index 475ce38ecd..ebdbb92947 100644 --- a/mistralrs-core/src/paged_attention/block_pool.rs +++ b/mistralrs-core/src/paged_attention/block_pool.rs @@ -12,7 +12,10 @@ use std::collections::HashMap; +use std::sync::Arc; + use super::block_hash::{BlockHash, BlockHashWithGroupId}; +use super::kv_cache_connector::{KvCacheConnector, NoopKvCacheConnector}; /// Sentinel value for "no link" in the doubly-linked free list. const NO_LINK: usize = usize::MAX; @@ -279,15 +282,32 @@ pub struct BlockPool { null_block_id: usize, /// The block size (number of tokens per block) for hash computation. hash_block_size: usize, + /// Optional external KV connector (no-op by default). + connector: Arc, } impl BlockPool { - /// Create a new block pool. + /// Create a new block pool with the default no-op connector. /// /// `num_gpu_blocks`: Number of physical GPU blocks. /// `enable_caching`: Whether to enable prefix caching. /// `hash_block_size`: Block size used for hash computation. pub fn new(num_gpu_blocks: usize, enable_caching: bool, hash_block_size: usize) -> Self { + Self::with_connector( + num_gpu_blocks, + enable_caching, + hash_block_size, + Arc::new(NoopKvCacheConnector), + ) + } + + /// Create a new block pool with a custom KV cache connector. + pub fn with_connector( + num_gpu_blocks: usize, + enable_caching: bool, + hash_block_size: usize, + connector: Arc, + ) -> Self { assert!(num_gpu_blocks > 0, "Must have at least 1 GPU block"); // Allocate blocks: [0..num_gpu_blocks) are real blocks, @@ -310,6 +330,7 @@ impl BlockPool { num_gpu_blocks, null_block_id: 0, // Will be set below hash_block_size, + connector, }; // Pop the first block as the null block (placeholder, never freed) @@ -496,6 +517,9 @@ impl BlockPool { /// Evict a cached block's hash from the cache map and reset its hash. fn maybe_evict_cached_block(&mut self, block_id: usize) { let block_hashes = std::mem::take(&mut self.blocks[block_id].block_hashes); + if !block_hashes.is_empty() { + self.connector.observe_evict(&block_hashes, block_id); + } for hash in &block_hashes { self.cached_block_hash_to_block.pop(hash, block_id); } diff --git a/mistralrs-core/src/paged_attention/cache_engine.rs b/mistralrs-core/src/paged_attention/cache_engine.rs index f6158e6a1b..c897b2a698 100644 --- a/mistralrs-core/src/paged_attention/cache_engine.rs +++ b/mistralrs-core/src/paged_attention/cache_engine.rs @@ -38,12 +38,25 @@ impl FromStr for PagedCacheType { } } -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct CacheConfig { pub block_size: usize, pub num_gpu_blocks: usize, pub cache_type: PagedCacheType, pub kv_cache_group_ids: Vec, + pub kv_cache_connector: std::sync::Arc, +} + +impl std::fmt::Debug for CacheConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CacheConfig") + .field("block_size", &self.block_size) + .field("num_gpu_blocks", &self.num_gpu_blocks) + .field("cache_type", &self.cache_type) + .field("kv_cache_group_ids", &self.kv_cache_group_ids) + .field("kv_cache_connector", &"Arc") + .finish() + } } pub type KVCache = (Tensor, Tensor); diff --git a/mistralrs-core/src/paged_attention/in_memory_kv_cache_connector.rs b/mistralrs-core/src/paged_attention/in_memory_kv_cache_connector.rs new file mode 100644 index 0000000000..3cf610d642 --- /dev/null +++ b/mistralrs-core/src/paged_attention/in_memory_kv_cache_connector.rs @@ -0,0 +1,107 @@ +//! In-memory reference connector for the KV cache seam. +//! +//! This is an example "external" tier: it keeps a hash -> local block-id index +//! outside the block pool's own prefix map. It does not offload tensor bytes +//! (disk/S3 would still need a hydrate path before returning block IDs). + +use std::collections::HashMap; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Mutex; + +use super::block_hash::{BlockHash, BlockHashWithGroupId}; +use super::kv_cache_connector::KvCacheConnector; + +/// Process-local external KV index used as a reference `KvCacheConnector`. +#[derive(Debug, Default)] +pub struct InMemoryKvCacheConnector { + blocks: Mutex>, + lookups: AtomicUsize, + hits: AtomicUsize, + stores: AtomicUsize, + evicts: AtomicUsize, +} + +impl InMemoryKvCacheConnector { + pub fn new() -> Self { + Self::default() + } + + pub fn len(&self) -> usize { + self.blocks.lock().expect("connector lock").len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } + + pub fn lookup_count(&self) -> usize { + self.lookups.load(Ordering::SeqCst) + } + + pub fn hit_count(&self) -> usize { + self.hits.load(Ordering::SeqCst) + } + + pub fn store_count(&self) -> usize { + self.stores.load(Ordering::SeqCst) + } + + pub fn evict_count(&self) -> usize { + self.evicts.load(Ordering::SeqCst) + } +} + +impl KvCacheConnector for InMemoryKvCacheConnector { + fn lookup_blocks( + &self, + block_hashes: &[BlockHash], + _group_ids: &[u32], + max_blocks: usize, + ) -> Option> { + self.lookups.fetch_add(1, Ordering::SeqCst); + let guard = self.blocks.lock().expect("connector lock"); + let mut ids = Vec::new(); + for hash in block_hashes.iter().take(max_blocks) { + match guard.get(hash) { + Some(&id) => ids.push(id), + None => break, + } + } + if ids.is_empty() { + None + } else { + Some(ids) + } + } + + fn observe_hit( + &self, + _block_hashes: &[BlockHash], + _block_ids: &[usize], + _num_computed_tokens: usize, + ) { + self.hits.fetch_add(1, Ordering::SeqCst); + } + + fn observe_store( + &self, + block_hashes: &[BlockHash], + block_ids: &[usize], + _num_full_blocks: usize, + ) { + self.stores.fetch_add(1, Ordering::SeqCst); + let mut guard = self.blocks.lock().expect("connector lock"); + for (hash, &block_id) in block_hashes.iter().zip(block_ids.iter()) { + guard.insert(*hash, block_id); + } + } + + fn observe_evict(&self, block_hashes: &[BlockHashWithGroupId], block_id: usize) { + self.evicts.fetch_add(1, Ordering::SeqCst); + let mut guard = self.blocks.lock().expect("connector lock"); + for entry in block_hashes { + guard.remove(&entry.block_hash); + } + guard.retain(|_, id| *id != block_id); + } +} diff --git a/mistralrs-core/src/paged_attention/kv_cache_connector.rs b/mistralrs-core/src/paged_attention/kv_cache_connector.rs new file mode 100644 index 0000000000..4b750d866f --- /dev/null +++ b/mistralrs-core/src/paged_attention/kv_cache_connector.rs @@ -0,0 +1,47 @@ +//! Optional KV cache connector seam for external cache tiers. +//! +//! Default behavior is local-only prefix caching. A connector may observe +//! lookup/store/evict traffic or supply already-materialized local block IDs. + +use super::block_hash::{BlockHash, BlockHashWithGroupId}; + +/// Hook into paged-attention prefix-cache block flow. +/// +/// Implementations must keep default local behavior intact when they miss: +/// `lookup_blocks` returning `None` means "fall back to the in-process pool." +/// Returned block IDs, when `Some`, must already be valid IDs in the local +/// `BlockPool` (e.g. after the connector has hydrated external KV into them). +pub trait KvCacheConnector: Send + Sync { + fn lookup_blocks( + &self, + _block_hashes: &[BlockHash], + _group_ids: &[u32], + _max_blocks: usize, + ) -> Option> { + None + } + + fn observe_hit( + &self, + _block_hashes: &[BlockHash], + _block_ids: &[usize], + _num_computed_tokens: usize, + ) { + } + + fn observe_store( + &self, + _block_hashes: &[BlockHash], + _block_ids: &[usize], + _num_full_blocks: usize, + ) { + } + + fn observe_evict(&self, _block_hashes: &[BlockHashWithGroupId], _block_id: usize) {} +} + +/// Default connector: always miss, never observe. +#[derive(Debug, Default, Clone, Copy)] +pub struct NoopKvCacheConnector; + +impl KvCacheConnector for NoopKvCacheConnector {} diff --git a/mistralrs-core/src/paged_attention/kv_cache_manager.rs b/mistralrs-core/src/paged_attention/kv_cache_manager.rs index 7c2ad67bc9..7502d80132 100644 --- a/mistralrs-core/src/paged_attention/kv_cache_manager.rs +++ b/mistralrs-core/src/paged_attention/kv_cache_manager.rs @@ -10,9 +10,11 @@ //! - `cache_blocks`: Cache newly-full blocks after computation. use std::collections::HashMap; +use std::sync::Arc; use super::block_hash::BlockHash; use super::block_pool::BlockPool; +use super::kv_cache_connector::{KvCacheConnector, NoopKvCacheConnector}; /// Result of `get_computed_blocks`: cached block IDs and how many tokens they cover. #[derive(Debug)] @@ -50,10 +52,12 @@ pub struct KVCacheManager { kv_cache_group_ids: Vec, /// Per-request block tracking. req_to_blocks: HashMap, + /// Optional external KV connector (no-op by default). + connector: Arc, } impl KVCacheManager { - /// Create a new KV cache manager. + /// Create a new KV cache manager with the default no-op connector. /// /// - `num_gpu_blocks`: Total number of physical GPU blocks. /// - `block_size`: Tokens per block. @@ -64,13 +68,36 @@ impl KVCacheManager { block_size: usize, enable_caching: bool, kv_cache_group_ids: Vec, + ) -> Self { + Self::with_connector( + num_gpu_blocks, + block_size, + enable_caching, + kv_cache_group_ids, + Arc::new(NoopKvCacheConnector), + ) + } + + /// Create a new KV cache manager with a custom KV cache connector. + pub fn with_connector( + num_gpu_blocks: usize, + block_size: usize, + enable_caching: bool, + kv_cache_group_ids: Vec, + connector: Arc, ) -> Self { Self { - block_pool: BlockPool::new(num_gpu_blocks, enable_caching, block_size), + block_pool: BlockPool::with_connector( + num_gpu_blocks, + enable_caching, + block_size, + Arc::clone(&connector), + ), block_size, enable_caching, kv_cache_group_ids, req_to_blocks: HashMap::new(), + connector, } } @@ -139,28 +166,45 @@ impl KVCacheManager { let mut cached_block_ids = Vec::new(); - for (i, &block_hash) in block_hashes.iter().enumerate() { - if i >= max_num_blocks { - break; + if let Some(ids) = + self.connector + .lookup_blocks(block_hashes, &self.kv_cache_group_ids, max_num_blocks) + { + if !ids.is_empty() { + let take = ids.len().min(max_num_blocks); + cached_block_ids.extend_from_slice(&ids[..take]); } + } - if let Some(ids) = self - .block_pool - .get_cached_block(block_hash, &self.kv_cache_group_ids) - { - let Some(first) = ids.first().copied() else { + if cached_block_ids.is_empty() { + for (i, &block_hash) in block_hashes.iter().enumerate() { + if i >= max_num_blocks { break; - }; - if ids.iter().any(|&id| id != first) { + } + + if let Some(ids) = self + .block_pool + .get_cached_block(block_hash, &self.kv_cache_group_ids) + { + let Some(first) = ids.first().copied() else { + break; + }; + if ids.iter().any(|&id| id != first) { + break; + } + cached_block_ids.push(first); + } else { break; } - cached_block_ids.push(first); - } else { - break; } } let num_computed_tokens = cached_block_ids.len() * self.block_size; + if !cached_block_ids.is_empty() { + let hit_hashes = &block_hashes[..cached_block_ids.len().min(block_hashes.len())]; + self.connector + .observe_hit(hit_hashes, &cached_block_ids, num_computed_tokens); + } ComputedBlocks { block_ids: cached_block_ids, @@ -318,9 +362,8 @@ impl KVCacheManager { return; } - let req = match self.req_to_blocks.get_mut(&request_id) { - Some(r) => r, - None => return, + let Some(req) = self.req_to_blocks.get(&request_id) else { + return; }; // Clamp to allocated blocks: callers may pass token counts that run ahead of @@ -330,18 +373,29 @@ impl KVCacheManager { return; } + let previously_cached = req.num_cached_blocks; + let block_ids = req.block_ids.clone(); + // Cache each full block for each group ID for &group_id in &self.kv_cache_group_ids { self.block_pool.cache_full_blocks( - &req.block_ids, + &block_ids, block_hashes, - req.num_cached_blocks, + previously_cached, num_full_blocks, group_id, ); } - req.num_cached_blocks = num_full_blocks; + self.req_to_blocks + .get_mut(&request_id) + .expect("request must still exist") + .num_cached_blocks = num_full_blocks; + + let new_ids = &block_ids[previously_cached..num_full_blocks]; + let new_hashes = &block_hashes[previously_cached..num_full_blocks.min(block_hashes.len())]; + self.connector + .observe_store(new_hashes, new_ids, num_full_blocks); } /// Get the block IDs allocated for a request. @@ -672,4 +726,144 @@ mod tests { let computed = mgr.get_computed_blocks(&hashes, 8); assert_eq!(computed.num_computed_tokens, 0); } + + struct RecordingConnector { + lookups: std::sync::atomic::AtomicUsize, + hits: std::sync::atomic::AtomicUsize, + stores: std::sync::atomic::AtomicUsize, + evicts: std::sync::atomic::AtomicUsize, + } + + impl RecordingConnector { + fn new() -> Self { + Self { + lookups: std::sync::atomic::AtomicUsize::new(0), + hits: std::sync::atomic::AtomicUsize::new(0), + stores: std::sync::atomic::AtomicUsize::new(0), + evicts: std::sync::atomic::AtomicUsize::new(0), + } + } + + fn count(&self, field: &std::sync::atomic::AtomicUsize) -> usize { + field.load(std::sync::atomic::Ordering::SeqCst) + } + } + + impl KvCacheConnector for RecordingConnector { + fn lookup_blocks( + &self, + _block_hashes: &[BlockHash], + _group_ids: &[u32], + _max_blocks: usize, + ) -> Option> { + self.lookups + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + None + } + + fn observe_hit( + &self, + _block_hashes: &[BlockHash], + _block_ids: &[usize], + _num_computed_tokens: usize, + ) { + self.hits.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + + fn observe_store( + &self, + _block_hashes: &[BlockHash], + _block_ids: &[usize], + _num_full_blocks: usize, + ) { + self.stores + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + + fn observe_evict( + &self, + _block_hashes: &[crate::paged_attention::BlockHashWithGroupId], + _block_id: usize, + ) { + self.evicts + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + } + } + + #[test] + fn test_connector_observes_store_hit_and_keeps_local_behavior() { + let connector = Arc::new(RecordingConnector::new()); + let mut mgr = KVCacheManager::with_connector( + 16, + 4, + true, + vec![0], + Arc::clone(&connector) as Arc, + ); + + let tokens: Vec = (1..=8).collect(); + let hashes = compute_block_hashes(&tokens, 4, &[], &[]); + + mgr.allocate_slots(1, 8, &[]).unwrap(); + mgr.cache_blocks(1, &hashes, 8); + assert_eq!(connector.count(&connector.stores), 1); + + mgr.free(1); + + let computed = mgr.get_computed_blocks(&hashes, 12); + assert_eq!(computed.num_computed_tokens, 8); + assert_eq!(computed.block_ids.len(), 2); + assert_eq!(connector.count(&connector.lookups), 1); + assert_eq!(connector.count(&connector.hits), 1); + } + + #[test] + fn test_in_memory_connector_serves_prefix_hit() { + let external = Arc::new(crate::paged_attention::InMemoryKvCacheConnector::new()); + let mut mgr = KVCacheManager::with_connector( + 16, + 4, + true, + vec![0], + Arc::clone(&external) as Arc, + ); + + let tokens: Vec = (1..=8).collect(); + let hashes = compute_block_hashes(&tokens, 4, &[], &[]); + mgr.allocate_slots(1, 8, &[]).unwrap(); + mgr.cache_blocks(1, &hashes, 8); + mgr.free(1); + + // Drop the local prefix map so only the external connector can hit. + assert!(mgr.reset_prefix_cache()); + + let computed = mgr.get_computed_blocks(&hashes, 12); + assert_eq!(computed.num_computed_tokens, 8); + assert_eq!(computed.block_ids.len(), 2); + assert!(external.lookup_count() >= 1); + assert!(external.hit_count() >= 1); + assert_eq!(external.len(), 2); + } + + #[test] + fn test_connector_observes_evict_on_reallocation() { + use crate::paged_attention::block_hash::hash_block_tokens; + use crate::paged_attention::block_pool::BlockPool; + + let connector = Arc::new(RecordingConnector::new()); + let mut pool = BlockPool::with_connector( + 4, + true, + 4, + Arc::clone(&connector) as Arc, + ); + + let block_ids = pool.get_new_blocks(3).unwrap(); + let h0 = hash_block_tokens(None, &[1, 2, 3, 4], None); + pool.cache_full_blocks(&block_ids, &[h0, h0, h0], 0, 1, 0); + pool.free_blocks(&block_ids); + + let _new_ids = pool.get_new_blocks(3).unwrap(); + assert!(connector.count(&connector.evicts) >= 1); + } } diff --git a/mistralrs-core/src/paged_attention/mod.rs b/mistralrs-core/src/paged_attention/mod.rs index 220cead19c..ffad3cb376 100644 --- a/mistralrs-core/src/paged_attention/mod.rs +++ b/mistralrs-core/src/paged_attention/mod.rs @@ -10,6 +10,10 @@ mod cache_engine; mod config; /// Encoder output cache for multimodal models (vision/audio encoder outputs). pub mod encoder_cache; +/// In-memory reference external KV cache connector. +pub mod in_memory_kv_cache_connector; +/// Optional external KV cache connector seam. +pub mod kv_cache_connector; /// KV Cache Manager: high-level block allocation, prefix cache lookups, per-request tracking. pub mod kv_cache_manager; mod layers; @@ -20,9 +24,12 @@ mod scheduler; pub const _PAD_SLOT_ID: i64 = -1; pub use attention_backend::AttentionBackendKind; +pub use block_hash::{compute_block_hashes, BlockHash, BlockHashWithGroupId}; pub use cache_engine::{CacheConfig, CacheEngine, PagedCacheType}; use candle_core::{DType, Device}; pub use config::{KvCacheLayout, KvCacheTopology, ModelConfigLike, ModelConfigMetadata}; +pub use in_memory_kv_cache_connector::InMemoryKvCacheConnector; +pub use kv_cache_connector::{KvCacheConnector, NoopKvCacheConnector}; pub use kv_cache_manager::KVCacheManager; pub use layers::PagedAttention; pub use scheduler::{ @@ -72,11 +79,23 @@ mod tests { } /// All memory counts in MB. Default for block size is 32. -#[derive(Clone, Copy, Debug)] +#[derive(Clone)] pub struct PagedAttentionConfig { pub(crate) block_size: Option, pub(crate) mem_gpu: MemoryGpuConfig, pub(crate) cache_type: PagedCacheType, + pub(crate) kv_cache_connector: std::sync::Arc, +} + +impl std::fmt::Debug for PagedAttentionConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PagedAttentionConfig") + .field("block_size", &self.block_size) + .field("mem_gpu", &self.mem_gpu) + .field("cache_type", &self.cache_type) + .field("kv_cache_connector", &"Arc") + .finish() + } } impl PagedAttentionConfig { @@ -89,8 +108,17 @@ impl PagedAttentionConfig { block_size, mem_gpu, cache_type, + kv_cache_connector: std::sync::Arc::new(NoopKvCacheConnector), }) } + + pub fn with_kv_cache_connector( + mut self, + connector: std::sync::Arc, + ) -> Self { + self.kv_cache_connector = connector; + self + } } #[derive(Debug, Clone, Copy, PartialEq)] @@ -270,5 +298,6 @@ pub fn calculate_cache_config( num_gpu_blocks, cache_type, kv_cache_group_ids: config.kv_cache_group_ids(), + kv_cache_connector: std::sync::Arc::new(NoopKvCacheConnector), }) } diff --git a/mistralrs-core/src/paged_attention/scheduler.rs b/mistralrs-core/src/paged_attention/scheduler.rs index 3719e02b9f..e10858f531 100644 --- a/mistralrs-core/src/paged_attention/scheduler.rs +++ b/mistralrs-core/src/paged_attention/scheduler.rs @@ -87,11 +87,12 @@ impl PagedAttentionScheduler { Self { waiting: VecDeque::new(), running: VecDeque::new(), - kv_cache_manager: Arc::new(tokio::sync::Mutex::new(KVCacheManager::new( + kv_cache_manager: Arc::new(tokio::sync::Mutex::new(KVCacheManager::with_connector( cache_config.num_gpu_blocks, cache_config.block_size, true, cache_config.kv_cache_group_ids.clone(), + cache_config.kv_cache_connector, ))), block_size: cache_config.block_size, config, @@ -799,6 +800,9 @@ mod tests { num_gpu_blocks: 128, cache_type: PagedCacheType::Auto, kv_cache_group_ids: vec![0], + kv_cache_connector: std::sync::Arc::new( + crate::paged_attention::NoopKvCacheConnector, + ), }, ) } diff --git a/mistralrs-core/src/pipeline/gguf.rs b/mistralrs-core/src/pipeline/gguf.rs index bed3361e3f..c93b7da16a 100644 --- a/mistralrs-core/src/pipeline/gguf.rs +++ b/mistralrs-core/src/pipeline/gguf.rs @@ -495,7 +495,7 @@ impl Loader for GGUFLoader { let (cache_config, cache_engine) = if let Some(paged_attn_config) = paged_attn_config { let model_config: &dyn ModelConfigLike = &model_config_metadata; - let cache_config = calculate_cache_config( + let mut cache_config = calculate_cache_config( paged_attn_config.mem_gpu, paged_attn_config.block_size, internal_dtype, @@ -507,6 +507,7 @@ impl Loader for GGUFLoader { None, max_kv_tokens, )?; + cache_config.kv_cache_connector = paged_attn_config.kv_cache_connector; let cache_engine = CacheEngine::new( model_config, &cache_config, diff --git a/mistralrs-core/src/pipeline/multimodal.rs b/mistralrs-core/src/pipeline/multimodal.rs index c7558b4913..b94db283a1 100644 --- a/mistralrs-core/src/pipeline/multimodal.rs +++ b/mistralrs-core/src/pipeline/multimodal.rs @@ -880,7 +880,7 @@ impl Loader for MultimodalLoader { } let model_metadata = model.model_config(); let (cache_config, cache_engine) = if let Some(paged_attn_config) = paged_attn_config { - let cache_config = calculate_cache_config( + let mut cache_config = calculate_cache_config( paged_attn_config.mem_gpu, paged_attn_config.block_size, dtype, @@ -892,6 +892,7 @@ impl Loader for MultimodalLoader { None, max_kv_tokens, )?; + cache_config.kv_cache_connector = paged_attn_config.kv_cache_connector; let cache_engine = CacheEngine::new( model_metadata.as_ref(), &cache_config, diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 5a8e612245..8e1b2ed043 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -888,7 +888,7 @@ impl Loader for NormalLoader { } let model_metadata = model.model_config(); let (cache_config, cache_engine) = if let Some(paged_attn_config) = paged_attn_config { - let cache_config = calculate_cache_config( + let mut cache_config = calculate_cache_config( paged_attn_config.mem_gpu, paged_attn_config.block_size, dtype, @@ -904,6 +904,7 @@ impl Loader for NormalLoader { None, max_kv_tokens, )?; + cache_config.kv_cache_connector = paged_attn_config.kv_cache_connector; let mut layer_devices = Vec::new(); for layer in 0..self.inner.num_layers(&config)? { diff --git a/mistralrs-core/src/resource_plan.rs b/mistralrs-core/src/resource_plan.rs index ef2eb2790c..aaa9292440 100644 --- a/mistralrs-core/src/resource_plan.rs +++ b/mistralrs-core/src/resource_plan.rs @@ -8,7 +8,7 @@ pub enum PagedKvPolicy { FairContext, } -#[derive(Clone, Copy, Debug)] +#[derive(Clone, Debug)] pub struct PagedKvModelRequest { pub paged_attn: Option, pub max_num_seqs: usize, @@ -54,6 +54,7 @@ fn plan_fair_context_paged_kv(models: &[PagedKvModelRequest]) -> anyhow::Result< .map(|model| { model .paged_attn + .clone() .map(|config| split_paged_config(config, model.max_num_seqs.max(1), active_weight)) .transpose() }) @@ -96,5 +97,8 @@ fn split_paged_config( MemoryGpuConfig::ContextSize(tokens) => MemoryGpuConfig::ContextSize(share(tokens)), }; - PagedAttentionConfig::new(config.block_size, mem_gpu, config.cache_type) + Ok( + PagedAttentionConfig::new(config.block_size, mem_gpu, config.cache_type)? + .with_kv_cache_connector(config.kv_cache_connector), + ) } diff --git a/mistralrs/Cargo.toml b/mistralrs/Cargo.toml index 57c723e04f..186dff91f8 100644 --- a/mistralrs/Cargo.toml +++ b/mistralrs/Cargo.toml @@ -191,6 +191,10 @@ path = "examples/advanced/batching_embeddings/main.rs" name = "paged_attn" path = "examples/advanced/paged_attn/main.rs" +[[example]] +name = "kv_cache_connector" +path = "examples/advanced/kv_cache_connector/main.rs" + [[example]] name = "auto_device_map" path = "examples/advanced/auto_device_map/main.rs" diff --git a/mistralrs/examples/advanced/kv_cache_connector/main.rs b/mistralrs/examples/advanced/kv_cache_connector/main.rs new file mode 100644 index 0000000000..2a7a06fdb3 --- /dev/null +++ b/mistralrs/examples/advanced/kv_cache_connector/main.rs @@ -0,0 +1,91 @@ +//! Realistic usage of an external KV cache connector with paged attention. +//! +//! Installs `InMemoryKvCacheConnector` through the normal model builder path +//! (same place you'd later plug disk/S3), then runs two chats that share a +//! prefix so the connector sees store/lookup traffic. +//! +//! Run with: +//! `cargo run --release -p mistralrs --example kv_cache_connector` + +use std::sync::Arc; + +use anyhow::Result; +use mistralrs::{ + InMemoryKvCacheConnector, IsqBits, KvCacheConnector, MemoryGpuConfig, ModelBuilder, + PagedAttentionMetaBuilder, TextMessageRole, TextMessages, +}; + +#[tokio::main] +async fn main() -> Result<()> { + let connector = Arc::new(InMemoryKvCacheConnector::new()); + + let model = ModelBuilder::new("Qwen/Qwen3-4B") + .with_auto_isq(IsqBits::Eight) + .with_logging() + .with_paged_attn( + PagedAttentionMetaBuilder::default() + .with_block_size(32) + .with_gpu_memory(MemoryGpuConfig::ContextSize(1024)) + .with_kv_cache_connector(Arc::clone(&connector) as Arc) + .build()?, + ) + .build() + .await?; + + let shared = "You are a concise assistant."; + let first = TextMessages::new() + .add_message(TextMessageRole::System, shared) + .add_message(TextMessageRole::User, "Say hello in one short sentence."); + let second = TextMessages::new() + .add_message(TextMessageRole::System, shared) + .add_message( + TextMessageRole::User, + "Say hello in one short sentence. Then add a second short sentence.", + ); + + let response_a = model.send_chat_request(first).await?; + println!( + "turn A: {}", + response_a.choices[0] + .message + .content + .as_deref() + .unwrap_or("") + ); + println!( + "after turn A: stores={}, lookups={}, hits={}, external_entries={}", + connector.store_count(), + connector.lookup_count(), + connector.hit_count(), + connector.len() + ); + + let response_b = model.send_chat_request(second).await?; + println!( + "turn B: {}", + response_b.choices[0] + .message + .content + .as_deref() + .unwrap_or("") + ); + println!( + "after turn B: stores={}, lookups={}, hits={}, external_entries={}", + connector.store_count(), + connector.lookup_count(), + connector.hit_count(), + connector.len() + ); + + anyhow::ensure!( + connector.store_count() > 0, + "expected the live paged-attention path to observe_store via the connector" + ); + anyhow::ensure!( + connector.lookup_count() > 0, + "expected the live paged-attention path to call lookup_blocks via the connector" + ); + + println!("ok: connector received traffic from the real paged-attention scheduler"); + Ok(()) +} diff --git a/mistralrs/src/lib.rs b/mistralrs/src/lib.rs index 9bc58b515f..d16c40bbb5 100644 --- a/mistralrs/src/lib.rs +++ b/mistralrs/src/lib.rs @@ -304,6 +304,10 @@ pub use mistralrs_core::{ }; // ========== Config Types ========== +pub use mistralrs_core::{ + compute_block_hashes, BlockHash, BlockHashWithGroupId, InMemoryKvCacheConnector, + KVCacheManager, KvCacheConnector, NoopKvCacheConnector, +}; pub use mistralrs_core::{ DefaultSchedulerMethod, IsqType, MemoryGpuConfig, MistralRsConfig, ModelDType, PagedAttentionConfig, PagedCacheType, SchedulerConfig, WebSearchOptions, diff --git a/mistralrs/src/model_builder_trait.rs b/mistralrs/src/model_builder_trait.rs index 98f7fa29fb..6b8de116ab 100644 --- a/mistralrs/src/model_builder_trait.rs +++ b/mistralrs/src/model_builder_trait.rs @@ -60,10 +60,10 @@ impl AnyModelBuilder { fn paged_attn_cfg(&self) -> Option { match self { - AnyModelBuilder::Text(b) => b.paged_attn_cfg, - AnyModelBuilder::Multimodal(b) => b.paged_attn_cfg, - AnyModelBuilder::Auto(b) => b.paged_attn_cfg, - AnyModelBuilder::Gguf(b) => b.paged_attn_cfg, + AnyModelBuilder::Text(b) => b.paged_attn_cfg.clone(), + AnyModelBuilder::Multimodal(b) => b.paged_attn_cfg.clone(), + AnyModelBuilder::Auto(b) => b.paged_attn_cfg.clone(), + AnyModelBuilder::Gguf(b) => b.paged_attn_cfg.clone(), AnyModelBuilder::Diffusion(_) | AnyModelBuilder::Speech(_) | AnyModelBuilder::Embedding(_) => None, @@ -445,7 +445,7 @@ pub(crate) async fn build_pipeline_from_text_loader( !builder.with_logging, device_map_setting, isq_type, - builder.paged_attn_cfg, + builder.paged_attn_cfg.clone(), )?; if let Some(mtp_config) = builder.mtp_config.clone() { pipeline @@ -498,7 +498,7 @@ pub(crate) async fn build_pipeline_from_gguf_loader( !builder.with_logging, device_map_setting, None, - builder.paged_attn_cfg, + builder.paged_attn_cfg.clone(), )?; let scheduler_config = @@ -607,7 +607,7 @@ pub async fn build_text_pipeline( .clone() .unwrap_or(DeviceMapSetting::Auto(AutoDeviceMapParams::default_text())), isq_type, - builder.paged_attn_cfg, + builder.paged_attn_cfg.clone(), )?; if let Some(mtp_config) = builder.mtp_config.clone() { pipeline @@ -665,7 +665,7 @@ pub async fn build_text_pipeline( device, device_map_setting, isq: isq_type, - paged_attn_config: builder.paged_attn_cfg, + paged_attn_config: builder.paged_attn_cfg.clone(), silent: !builder.with_logging, chat_template: builder.chat_template.clone(), jinja_explicit: builder.jinja_explicit.clone(), @@ -732,7 +732,7 @@ pub async fn build_multimodal_pipeline( AutoDeviceMapParams::default_multimodal(), )), isq_type, - builder.paged_attn_cfg, + builder.paged_attn_cfg.clone(), )?; if let Some(mtp_config) = builder.mtp_config.clone() { pipeline @@ -795,7 +795,7 @@ pub async fn build_multimodal_pipeline( device, device_map_setting, isq: isq_type, - paged_attn_config: builder.paged_attn_cfg, + paged_attn_config: builder.paged_attn_cfg.clone(), silent: !builder.with_logging, chat_template: builder.chat_template.clone(), jinja_explicit: builder.jinja_explicit.clone(), @@ -850,7 +850,7 @@ pub async fn build_gguf_pipeline( .clone() .unwrap_or(DeviceMapSetting::Auto(AutoDeviceMapParams::default_text())), None, - builder.paged_attn_cfg, + builder.paged_attn_cfg.clone(), )?; let scheduler_config = scheduler_config_from_pipeline( @@ -891,7 +891,7 @@ pub async fn build_gguf_pipeline( device, device_map_setting, isq: None, - paged_attn_config: builder.paged_attn_cfg, + paged_attn_config: builder.paged_attn_cfg.clone(), silent: !builder.with_logging, chat_template: builder.chat_template.clone(), jinja_explicit: builder.jinja_explicit.clone(), @@ -1208,7 +1208,7 @@ pub async fn build_auto_pipeline( .clone() .unwrap_or(DeviceMapSetting::Auto(AutoDeviceMapParams::default_text())), isq_type, - builder.paged_attn_cfg, + builder.paged_attn_cfg.clone(), )?; if let Some(mtp_config) = builder.mtp_config.clone() { pipeline @@ -1268,7 +1268,7 @@ pub async fn build_auto_pipeline( device, device_map_setting, isq: isq_type, - paged_attn_config: builder.paged_attn_cfg, + paged_attn_config: builder.paged_attn_cfg.clone(), silent: !builder.with_logging, chat_template: builder.chat_template.clone(), jinja_explicit: builder.jinja_explicit.clone(), diff --git a/mistralrs/src/text_model.rs b/mistralrs/src/text_model.rs index b23452d26b..77883907d6 100644 --- a/mistralrs/src/text_model.rs +++ b/mistralrs/src/text_model.rs @@ -74,6 +74,7 @@ pub struct PagedAttentionMetaBuilder { block_size: Option, mem_gpu: MemoryGpuConfig, cache_type: PagedCacheType, + kv_cache_connector: Option>, } impl Default for PagedAttentionMetaBuilder { @@ -82,6 +83,7 @@ impl Default for PagedAttentionMetaBuilder { block_size: None, mem_gpu: MemoryGpuConfig::ContextSize(4096), cache_type: PagedCacheType::Auto, + kv_cache_connector: None, } } } @@ -105,9 +107,19 @@ impl PagedAttentionMetaBuilder { self } + /// Install an external KV cache connector used by the paged-attention scheduler. + pub fn with_kv_cache_connector(mut self, connector: Arc) -> Self { + self.kv_cache_connector = Some(connector); + self + } + /// Build the [`PagedAttentionConfig`]. Returns an error if the configuration is invalid. pub fn build(self) -> anyhow::Result { - PagedAttentionConfig::new(self.block_size, self.mem_gpu, self.cache_type) + let mut cfg = PagedAttentionConfig::new(self.block_size, self.mem_gpu, self.cache_type)?; + if let Some(connector) = self.kv_cache_connector { + cfg = cfg.with_kv_cache_connector(connector); + } + Ok(cfg) } }