企业官网建设流程全解析

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 length（2048 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: Option<u64>, } /// 显存池管理器 struct KVCachePool { /// 所有Block的数组 blocks: Vec<KVCacheBlock>, /// 空闲Block索引队列 free_blocks: VecDeque<usize>, /// 总显存容量（字节） total_memory_bytes: usize, /// 已用显存（字节） used_memory_bytes: AtomicUsize, /// 共享前缀映射：hash → block indices /// 用于快速查找已有前缀的缓存 prefix_index: HashMap<u64, Vec<usize>>, } 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: Vec<KVCacheBlock> = (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: VecDeque<usize> = (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, ) -> Result<SequenceAllocation, 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); } /// 驱逐Block：OOM时的紧急回收 fn evict_blocks(&mut self, needed: usize) -> Result<usize, CacheError> { let mut evicted = 0usize; // 驱逐策略：优先选择ref_count=1（非共享）的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: Vec<usize>, 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, }

关键设计：

• 固定大小Block（64 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值怎么办？正确的做法是CoW（Copy-on-Write）——共享Block通过引用计数标记为只读，当请求需要修改该Block的内容时（例如生成了新token需要更新该位置的KV值），先复制Block再修改，原Block保留给其他共享者。这种Copy-on-Write的开销就是Prefix-aware复用的代价：只有当共享前缀足够长（system prompt > 512 tokens），CoW复制的开销才小于重新计算的收益。实践中，对于<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池化消除碎片化

五、总结

1. KV Cache的Block管理是解决显存碎片和利用不足的核心机制
2. 引用计数支持前缀共享——system prompt的显存占用从N倍降为1倍
3. 前缀哈希索引实现O(1)的缓存命中判断——匹配延迟影响prefill总延迟
4. Block大小的选择需要在碎片率和分配效率间权衡——64 token是工程经验值
5. 驱逐策略应优先回收非共享Block：共享Block的驱逐直接影响多个请求