大模型推理引擎显存优化：从 KV Cache 管理到连续批处理的工程实践 大模型推理引擎显存优化从 KV Cache 管理到连续批处理的工程实践一、显存墙推理服务的真实瓶颈7B 参数模型在推理时单个请求的 KV Cache 就能占到 2-4GB取决于序列长度。并发请求一旦达到 20 个仅 KV Cache 就要消耗 40-80GB 显存直接顶满 A100 的 80GB 容量。更麻烦的是 KV Cache 的分配是动态的——每个请求的序列长度不一样如果预分配固定大小会产生大量碎片。我们在生产环境实测过用 PyTorch 默认的动态显存分配器KV Cache 的碎片率能到 35%。也就是说 80GB 显存里有 28GB 被碎片占着实际可用只有 52GB。并发能力被卡住GPU 利用率也才 45%。推理引擎的显存优化主要解决两个问题KV Cache 怎么高效管理以及多请求怎么批处理调度。前者决定单卡能扛多少并发后者决定吞吐量上限。二、KV Cache 管理与连续批处理核心机制2.1 KV Cache 的显存模型Transformer 推理时KV Cache 存的是每层注意力机制的 Key 和 Value 向量。显存占用公式如下KV_Cache_Size 2 × num_layers × seq_len × num_kv_heads × head_dim × dtype_size × batch_size以 Llama-2-7B 为例32 层32 KV 头128 维FP16序列长度 2048 时单个请求的 KV Cache 大约 1GB。这个数字随序列长度线性增长是推理服务扩展性的最大障碍。graph TD subgraph 传统PagedAttention之前 A1[请求1: 预分配2048 tokens] -- W1[浪费: 实际用512] A2[请求2: 预分配2048 tokens] -- W2[浪费: 实际用1024] A3[请求3: 预分配2048 tokens] -- W3[浪费: 实际用256] end subgraph PagedAttention B1[Block Table 请求1] -- P1[物理Block池] B2[Block Table 请求2] -- P1 B3[Block Table 请求3] -- P1 P1 -- PB[按需分配br/每Block 16 tokensbr/无碎片] end subgraph 连续批处理Continuous Batching CB1[迭代级调度] -- CB2[每步检查完成请求] CB2 -- CB3[移除已完成请求] CB3 -- CB4[插入新请求] CB4 -- CB1 end2.2 PagedAttentionKV Cache 的虚拟内存管理vLLM 提出的 PagedAttention 借鉴了操作系统的虚拟内存思想。把 KV Cache 切成固定大小的 Block通常 16 个 token每个请求维护一个 Block Table映射逻辑 Block 到物理 Block。物理 Block 从全局池里按需分配请求结束就归还。这种设计的核心优势消除碎片物理 Block 固定大小没有外部碎片按需分配只在生成新 token 时分配新 Block不用预分配最大序列长度共享前缀不同请求的相同前缀比如 System Prompt可以共享物理 Block2.3 连续批处理Continuous Batching传统 Static Batching 要等所有请求跑完才释放资源短请求会被长请求拖慢。Continuous Batching 在每个解码步检查请求状态马上移除完成的请求并插入新请求让 GPU 始终满载。三、基于 Rust 的 KV Cache 管理器生产级实现下面这段代码实现了一个简化的 PagedAttention Block Manager用 Rust 管理物理 Block 池use std::collections::{HashMap, VecDeque}; /// 物理Block大小每个Block存储16个token的KV向量 const BLOCK_SIZE: usize 16; /// 物理Block ID type BlockId usize; /// 请求ID type RequestId u64; /// KV Cache Block 管理器 /// /// 核心设计 /// 1. 物理Block池预分配固定数量的Block运行时按需分配 /// 2. Block Table每个请求维护逻辑Block到物理Block的映射 /// 3. 引用计数支持前缀共享物理Block可被多个请求引用 pub struct BlockManager { /// 物理Block池空闲Block队列 free_blocks: VecDequeBlockId, /// 总Block数量 total_blocks: usize, /// 每个请求的Block Table逻辑索引 → 物理Block ID block_tables: HashMapRequestId, VecBlockId, /// 物理Block引用计数用于前缀共享 ref_counts: HashMapBlockId, usize, /// 每个Block存储的token数量最后一个Block可能不满 block_token_counts: HashMapBlockId, usize, } impl BlockManager { /// 创建BlockManager预分配指定数量的物理Block pub fn new(num_blocks: usize) - Self { let free_blocks: VecDequeBlockId (0..num_blocks).collect(); Self { free_blocks, total_blocks: num_blocks, block_tables: HashMap::new(), ref_counts: HashMap::new(), block_token_counts: HashMap::new(), } } /// 为新请求分配初始Block pub fn allocate_request( mut self, request_id: RequestId, num_tokens: usize, ) - Result(), String { if self.block_tables.contains_key(request_id) { return Err(format!(请求 {} 已存在, request_id)); } let num_blocks_needed (num_tokens BLOCK_SIZE - 1) / BLOCK_SIZE; if self.free_blocks.len() num_blocks_needed { return Err(format!( Block不足: 需要 {}, 可用 {}, num_blocks_needed, self.free_blocks.len() )); } let mut table Vec::with_capacity(num_blocks_needed); for i in 0..num_blocks_needed { let block_id self.free_blocks.pop_front().unwrap(); table.push(block_id); *self.ref_counts.entry(block_id).or_insert(0) 1; // 最后一个Block可能不满 let tokens_in_block if i num_blocks_needed - 1 num_tokens % BLOCK_SIZE ! 0 { num_tokens % BLOCK_SIZE } else { BLOCK_SIZE }; self.block_token_counts.insert(block_id, tokens_in_block); } self.block_tables.insert(request_id, table); Ok(()) } /// 为请求追加一个token可能需要分配新Block pub fn append_token(mut self, request_id: RequestId) - Result(), String { let table self.block_tables.get_mut(request_id) .ok_or_else(|| format!(请求 {} 不存在, request_id))?; let last_block *table.last().unwrap(); let current_count self.block_token_counts.get(last_block).copied().unwrap_or(0); if current_count BLOCK_SIZE { // 当前Block还有空间直接追加 self.block_token_counts.insert(last_block, current_count 1); } else { // 当前Block已满分配新Block let new_block self.free_blocks.pop_front() .ok_or_else(|| Block池已耗尽无法分配新Block.to_string())?; table.push(new_block); *self.ref_counts.entry(new_block).or_insert(0) 1; self.block_token_counts.insert(new_block, 1); } Ok(()) } /// 释放请求的所有Block考虑引用计数 pub fn free_request(mut self, request_id: RequestId) - Result(), String { let table self.block_tables.remove(request_id) .ok_or_else(|| format!(请求 {} 不存在, request_id))?; for block_id in table { let count self.ref_counts.get_mut(block_id).unwrap(); *count - 1; if *count 0 { // 引用计数归零归还Block池 self.ref_counts.remove(block_id); self.block_token_counts.remove(block_id); self.free_blocks.push_back(block_id); } } Ok(()) } /// 共享前缀将源请求的前n个Block共享给目标请求 pub fn share_prefix( mut self, src_request: RequestId, num_blocks: usize, dst_request: RequestId, ) - Result(), String { let src_table self.block_tables.get(src_request) .ok_or_else(|| format!(源请求 {} 不存在, src_request))?; if num_blocks src_table.len() { return Err(共享Block数量超过源请求的Block数.into()); } let shared_blocks: VecBlockId src_table[..num_blocks].to_vec(); // 增加引用计数 for block_id in shared_blocks { *self.ref_counts.entry(block_id).or_insert(0) 1; } // 为目标请求分配额外的Block用于非共享部分 self.block_tables.insert(dst_request, shared_blocks); Ok(()) } /// 返回可用Block数量 pub fn available_blocks(self) - usize { self.free_blocks.len() } /// 返回显存利用率 pub fn utilization(self) - f64 { let used self.total_blocks - self.free_blocks.len(); used as f64 / self.total_blocks as f64 } } #[cfg(test)] mod tests { use super::*; #[test] fn test_allocate_and_free() { let mut mgr BlockManager::new(100); // 分配请求33个token需要3个Block16161 mgr.allocate_request(1, 33).unwrap(); assert_eq!(mgr.available_blocks(), 97); // 释放请求 mgr.free_request(1).unwrap(); assert_eq!(mgr.available_blocks(), 100); } #[test] fn test_append_token() { let mut mgr BlockManager::new(100); mgr.allocate_request(1, 16).unwrap(); // 恰好1个Block // 追加17个token第1个填满当前Block后续需要新Block for _ in 0..17 { mgr.append_token(1).unwrap(); } assert_eq!(mgr.available_blocks(), 97); // 用了3个Block } #[test] fn test_prefix_sharing() { let mut mgr BlockManager::new(100); mgr.allocate_request(1, 64).unwrap(); // 4个Block // 请求2共享请求1的前2个Block mgr.share_prefix(1, 2, 2).unwrap(); assert_eq!(mgr.available_blocks(), 96); // 只多用了0个物理Block // 释放请求1共享Block不应被回收 mgr.free_request(1).unwrap(); assert_eq!(mgr.available_blocks(), 96); // 共享Block仍被请求2引用 // 释放请求2共享Block才被回收 mgr.free_request(2).unwrap(); assert_eq!(mgr.available_blocks(), 100); } #[test] fn test_out_of_blocks() { let mut mgr BlockManager::new(3); let result mgr.allocate_request(1, 65); // 需要5个Block assert!(result.is_err()); } }四、推理引擎显存优化的架构权衡Block 大小的选择Block 太小比如 1 tokenBlock Table 会过大索引开销增加Block 太大比如 256 tokens最后一个 Block 的内部碎片会增加。vLLM 默认 16 是在 A100 上的经验值。在显存更紧张的 GPU比如 4090 24GB上适当增大 Block Size 可以减少 Block Table 的显存开销。前缀共享的缓存一致性共享 Block 的请求在生成阶段会写入新 Block但不修改共享 Block。如果 System Prompt 变了需要让所有共享该前缀的 Block 失效。这要求维护一个Block → 请求的反向索引增加了管理复杂度。连续批处理的调度延迟每步检查请求状态并调度新请求引入了 CPU 端的调度开销。在 A100 上这个开销大约 50-100μs/步对于 20ms 的解码步长影响不大。但在 CPU 推理场景下解码步长可能只有 1-2ms调度开销占比会显著上升。量化与 KV Cache 精度的权衡FP8 KV Cache 可以把显存占用减半但部分模型在 FP8 下会出现精度退化尤其是长序列场景。KV Cache 量化需要逐模型验证精度不能一刀切。建议在 CI 中加入困惑度Perplexity回归测试确保量化后精度下降在可接受范围内。五、总结推理引擎的显存优化核心是 PagedAttention 的虚拟内存管理和 Continuous Batching 的迭代级调度。前者通过固定大小 Block 消除碎片、按需分配提升利用率、引用计数支持前缀共享后者通过每步调度保持 GPU 满载。两者结合能把 A100 上的并发能力提升 2-3 倍。落地建议先实现基础的 PagedAttention Block Manager验证显存利用率提升再接入 Continuous Batching 调度器验证吞吐量提升前缀共享作为第三阶段优化需要配套反向索引和缓存失效机制KV Cache 量化需要逐模型验证精度。推理优化是一个从显存到计算到调度的系统工程每个环节都需要数据驱动决策。