Ollama的内存管理策略：从KV Cache Block分配到Prefix-aware内存复用的演进

Ollama的内存管理策略从KV Cache Block分配到Prefix-aware内存复用的演进一、当13B模型的推理服务在并发8时OOM显存管理的精细化需求部署LLaMA-13B模型时单请求占用8GB显存模型权重5.2GB KV Cache 2.8GB。理论上16GB显存的GPU可以支持2个并发请求。但实际在并发数达到6时就开始OOM。分析发现KV Cache的静态分配策略按最大sequence length2048 tokens预分配而80%的请求实际只生成200-500个token。剩余的KV Cache空间闲置但不能复用。进一步分析多个请求共享相同的系统提示词system prompt每个请求都独立计算并缓存该前缀的KV值——重复计算和重复存储。Ollama基于llama.cpp通过KV Cache Block管理和Prefix-aware内存复用解决了这些问题。二、KV Cache内存管理的架构graph TB subgraph 显存布局 A[模型权重 5.2GB] B[KV Cache池] B1[Block 0] B2[Block 1] B3[Block ...] B4[Block N] end subgraph Block分配策略 C[请求到达] -- D{是否有共享前缀?} D --|是| E[引用共享前缀Block] D --|否| F[分配新Block] E -- G[仅分配增量Block] F -- G end subgraph Block回收 H[请求完成] -- I[标记Block为free] I -- J{是否被其他请求引用?} J --|是| K[仅减引用计数] J --|否| L[回收到free pool] end subgraph Prefix-aware复用 M[System Prompt] -- N[共享Block组] N -- O[请求1引用] N -- P[请求2引用] N -- Q[请求3引用] end三、KV Cache内存管理的工程实现use std::collections::{HashMap, VecDeque, HashSet}; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; /// KV Cache Block固定大小的缓存单元 /// Block大小 num_layers × num_heads × head_dim × dtype_size const BLOCK_SIZE_TOKENS: usize 64; // 每个Block覆盖64个token const BYTES_PER_TOKEN: usize 2 * 4096 * 32 * 128 * 2; // FP16 struct KVCacheBlock { /// Block在池中的索引 index: usize, /// 引用计数多个请求共享前缀时1 ref_count: AtomicUsize, /// 该Block覆盖的token范围在sequence中的位置 token_start: usize, /// 实际存储的token数≤BLOCK_SIZE_TOKENS num_tokens: usize, /// 内容的哈希值用于Prefix匹配 /// 使用token id序列的哈希而非KV值的哈希 content_hash: u64, /// 所属的sequence组用于批量回收 sequence_group: Optionu64, } /// 显存池管理器 struct KVCachePool { /// 所有Block的数组 blocks: VecKVCacheBlock, /// 空闲Block索引队列 free_blocks: VecDequeusize, /// 总显存容量字节 total_memory_bytes: usize, /// 已用显存字节 used_memory_bytes: AtomicUsize, /// 共享前缀映射hash → block indices /// 用于快速查找已有前缀的缓存 prefix_index: HashMapu64, Vecusize, } impl KVCachePool { /// 初始化根据显存容量计算Block数量 fn new(available_memory_bytes: usize) - Self { let block_bytes BLOCK_SIZE_TOKENS * BYTES_PER_TOKEN; let num_blocks available_memory_bytes / block_bytes; let blocks: VecKVCacheBlock (0..num_blocks) .map(|i| KVCacheBlock { index: i, ref_count: AtomicUsize::new(0), token_start: 0, num_tokens: 0, content_hash: 0, sequence_group: None, }) .collect(); let free_blocks: VecDequeusize (0..num_blocks).collect(); Self { blocks, free_blocks, total_memory_bytes: available_memory_bytes, used_memory_bytes: AtomicUsize::new(0), prefix_index: HashMap::new(), } } /// 为新sequence分配Block /// 检查是否有可复用的前缀Block fn allocate_sequence( mut self, tokens: [u32], max_tokens: usize, ) - ResultSequenceAllocation, CacheError { let mut allocated Vec::new(); let mut token_pos 0; let mut reused_blocks 0usize; let mut new_blocks 0usize; // 第一步查找共享前缀 while token_pos tokens.len() { let hash Self::hash_slice(tokens[token_pos..]); if let Some(cached) self.prefix_index.get(hash) { // 找到共享前缀Block for block_idx in cached { let block self.blocks[block_idx]; // 验证hash碰撞比较实际token if self.verify_block_match(block_idx, tokens[token_pos..]) { // 增加引用计数 block.ref_count.fetch_add(1, Ordering::Relaxed); allocated.push(block_idx); token_pos block.num_tokens; reused_blocks 1; continue; } } break; // hash匹配但内容不匹配中止前缀搜索 } else { break; // 无共享前缀 } } // 第二步为新token分配Block let remaining max_tokens - token_pos; let needed (remaining BLOCK_SIZE_TOKENS - 1) / BLOCK_SIZE_TOKENS; if self.free_blocks.len() needed { // 尝试驱逐优先驱逐非共享、低优先级的Block let evicted self.evict_blocks(needed - self.free_blocks.len())?; if evicted needed - self.free_blocks.len() self.free_blocks.len() { return Err(CacheError::OutOfMemory); } } for _ in 0..needed { if let Some(block_idx) self.free_blocks.pop_front() { let block mut self.blocks[block_idx]; block.ref_count.store(1, Ordering::Relaxed); block.token_start token_pos; block.num_tokens BLOCK_SIZE_TOKENS.min(remaining); allocated.push(block_idx); token_pos block.num_tokens; new_blocks 1; } } // 第三步将新分配的Block注册到前缀索引 let mut pos 0; for block_idx in allocated { if new_blocks 0 pos tokens.len() { let hash Self::hash_slice(tokens[pos..]); self.prefix_index.entry(hash) .or_insert_with(Vec::new) .push(block_idx); } pos self.blocks[block_idx].num_tokens; } let block_bytes BLOCK_SIZE_TOKENS * BYTES_PER_TOKEN * allocated.len(); self.used_memory_bytes.fetch_add(block_bytes, Ordering::Relaxed); Ok(SequenceAllocation { block_indices: allocated, total_blocks: reused_blocks new_blocks, reused_blocks, }) } /// 释放sequence的Block fn release_sequence(mut self, allocation: SequenceAllocation) { for block_idx in allocation.block_indices { let block self.blocks[block_idx]; // 减引用计数 let prev block.ref_count.fetch_sub(1, Ordering::Relaxed); // 最后一个引用者回收Block if prev 1 { // 从前缀索引中移除 self.remove_from_prefix_index(block_idx); // 回收到空闲池 self.free_blocks.push_back(block_idx); block.num_tokens 0; block.content_hash 0; } } let block_bytes BLOCK_SIZE_TOKENS * BYTES_PER_TOKEN * allocation.block_indices.len(); self.used_memory_bytes.fetch_sub(block_bytes, Ordering::Relaxed); } /// 驱逐BlockOOM时的紧急回收 fn evict_blocks(mut self, needed: usize) - Resultusize, CacheError { let mut evicted 0usize; // 驱逐策略优先选择ref_count1非共享的Block // 按LRU排序最早分配的Block优先驱逐 // 收集可驱逐的Block let mut candidates: Vec(usize, usize) self.blocks.iter() .filter(|b| b.ref_count.load(Ordering::Relaxed) 1 b.num_tokens 0) .map(|b| (b.index, b.token_start)) // token_start作为简单的时间戳代理 .collect(); // 按token_start排序越早分配的越先驱逐LRU近似 candidates.sort_by_key(|(_, ts)| *ts); for (block_idx, _) in candidates.iter().take(needed) { self.remove_from_prefix_index(*block_idx); self.free_blocks.push_back(*block_idx); evicted 1; } if evicted needed { return Err(CacheError::CannotEvictEnough); } Ok(evicted) } /// 预处理阶段预热常用前缀 /// 将System Prompt或Few-shot examples预加载到KV Cache fn warmup_prefix(mut self, tokens: [u32]) - Result(), CacheError { // 分配Block但不绑定到特定sequence // 设置ref_count为0永久保留 let _ self.allocate_sequence(tokens, tokens.len())?; // 标记为预热Block不被驱逐策略选中 // 在allocate_sequence中分配的Block已在前缀索引中注册 Ok(()) } fn hash_slice(tokens: [u32]) - u64 { // 使用wyhash或类似快速哈希 // 哈希窗口BLOCK_SIZE_TOKENS长度的滑动窗口 let mut hash: u64 0xcbf29ce484222325; let window tokens.len().min(BLOCK_SIZE_TOKENS); for token in tokens[..window] { hash ^ token as u64; hash hash.wrapping_mul(0x100000001b3); } hash } fn verify_block_match(self, block_idx: usize, tokens: [u32]) - bool { let block self.blocks[block_idx]; // 实际实现需比较token id序列 // 这里简化信任hash true } fn remove_from_prefix_index(mut self, block_idx: usize) { let hash self.blocks[block_idx].content_hash; if let Some(indices) self.prefix_index.get_mut(hash) { indices.retain(|i| i ! block_idx); if indices.is_empty() { self.prefix_index.remove(hash); } } } /// 获取池状态用于监控 fn stats(self) - CacheStats { let total_blocks self.blocks.len(); let used_blocks total_blocks - self.free_blocks.len(); let shared_blocks self.blocks.iter() .filter(|b| b.ref_count.load(Ordering::Relaxed) 1) .count(); CacheStats { total_blocks, used_blocks, shared_blocks, free_blocks: self.free_blocks.len(), used_memory_bytes: self.used_memory_bytes.load(Ordering::Relaxed), total_memory_bytes: self.total_memory_bytes, } } } /// Sequence的Block分配结果 struct SequenceAllocation { block_indices: Vecusize, total_blocks: usize, reused_blocks: usize, } /// 缓存统计 #[derive(Debug)] struct CacheStats { total_blocks: usize, used_blocks: usize, shared_blocks: usize, free_blocks: usize, used_memory_bytes: usize, total_memory_bytes: usize, } #[derive(Debug, thiserror::Error)] enum CacheError { #[error(Out of memory: cannot allocate more blocks)] OutOfMemory, #[error(Cannot evict enough blocks)] CannotEvictEnough, }关键设计固定大小Block64 tokens在内存碎片化和分配开销间取得平衡引用计数共享前缀System prompt只需缓存一次哈希前缀索引O(1)查找已有前缀LRU驱逐策略非共享Block优先驱逐Prefix-aware复用的序列化开销与跨请求一致性。前缀匹配的核心瓶颈不在于哈希查找而在于Token序列化——每个新请求到达时需要将其prompt tokens序列化为哈希链以BLOCK_SIZE_TOKENS为窗口滑动计算64-bit哈希对于2048 tokens的system prompt这需要32次哈希计算。虽然FNV-1a哈希极快~1ns/hash但关键路径上的32ns累积在高并发下不可忽视。优化的方式是将前缀哈希预计算并缓存预热阶段一次性计算所有前缀BLOCK的哈希存储在prefix_index中后续请求的匹配只需对incremental tokens不包含已缓存前缀的部分计算哈希。更复杂的问题是前缀共享的一致性语义如果请求A和请求B共享前缀但后续生成不同请求B修改了共享Block中的KV值怎么办正确的做法是CoWCopy-on-Write——共享Block通过引用计数标记为只读当请求需要修改该Block的内容时例如生成了新token需要更新该位置的KV值先复制Block再修改原Block保留给其他共享者。这种Copy-on-Write的开销就是Prefix-aware复用的代价只有当共享前缀足够长system prompt 512 tokensCoW复制的开销才小于重新计算的收益。实践中对于256 tokens的短前缀直接计算比复用更快——避免了引用计数的原子操作和潜在的CoW分配。四、内存管理策略的边界Block大小的权衡小块16 tokens碎片少但分配/回收频率高索引开销大大块256 tokens分配快但内碎片多未使用的token空间64 token的经验值Block开销约5%碎片率10%Prefix-aware复用的局限仅对完全相同的前缀有效——近义前缀的复用需要语义级别的匹配Hash碰撞风险需要token-level验证兜底预热策略不常用前缀的缓存反而浪费显存适用场景多轮对话共享system prompt显存节省30-50%批量推理相同prompt模板前缀复用率80%高并发服务Block池化消除碎片化五、总结KV Cache的Block管理是解决显存碎片和利用不足的核心机制引用计数支持前缀共享——system prompt的显存占用从N倍降为1倍前缀哈希索引实现O(1)的缓存命中判断——匹配延迟影响prefill总延迟Block大小的选择需要在碎片率和分配效率间权衡——64 token是工程经验值驱逐策略应优先回收非共享Block共享Block的驱逐直接影响多个请求