LLM 推理中的 PageAttention 实现原理：虚拟内存管理在 KV Cache 中的创新应用

LLM 推理中的 PageAttention 实现原理虚拟内存管理在 KV Cache 中的创新应用一、KV Cache 的内碎片问题固定分配如何浪费 40% 显存传统 LLM 推理为每个请求预分配一块连续的 KV Cache 空间大小 max_seq_len × num_layers × 2 × head_dim × num_heads × dtype_size。以 Llama-2-70B 为例max_seq_len4096 时的单请求 KV Cache 为 2.5GB。当实际生成长度仅 128 tokens 时预分配的 96.9% 空间闲置。更严重的是连续分配导致严重的外碎片。多个请求的 KV Cache 在显存中交错分配短请求完成后释放的空间被长请求的活跃 KV Cache 块隔开无法合并成大块。最终显存利用率仅 60%但 GPU OOM 频率上升——不是因为总量不够而是无法分配连续空间。此外连续分配策略在动态批处理Dynamic Batching场景下表现极差。当 Batch 中某个请求提前结束时其释放的 KV Cache 空间若小于其他请求所需就会成为无法利用的空洞。PageAttention 的非连续分配天然解决了这个问题——释放的 Block 可以立即分配给任何其他请求无论其大小。PageAttentionvLLM 提出的 PagedAttention将虚拟内存的分页思想引入 KV Cache 管理KV Cache 被分割为固定大小的 Block类似 4KB 内存页每个 Block 存储固定数量 token 的 Key 和 Value。请求的 KV Cache 由 Block 链表组成不需要物理连续。释放时逐 Block 回收无外碎片。二、PageAttention 的虚拟内存管理架构Attention kernel 计算时通过block_table[request_id]查找每个 token 对应的物理 Block 地址然后使用 GPU 的gather指令从非连续的 Block 中收集 Key/Value。这个收集操作的额外开销约 3~5%——远小于显存利用率的提升从 60% 到 96%。gather操作的性能瓶颈在于显存带宽利用率。当 Block 分散在 VRAM 的不同区域时GPU 的 Memory Controller 需要发起多次显存事务Transaction降低了有效带宽。但 Attention 计算本身是 Bandwidth-Bound 的计算量小数据搬运量大额外的gather开销被掩盖在计算延迟中。实测显示在 A100 80GB 上PageAttention 的 Attention Kernel 吞吐量仅比连续分配低 2-3%。Block 大小是一个关键超参数。Block 越大映射表的条目越少减少查找开销但内碎片越大。vLLM 默认 Block Size 16这是基于实验的折中——A100 的 L2 Cache line 128 字节16 个 fp16 token 的 KV 约为 16 × 2 × 128 × 32 × 2 / 8 32KB适配 L1 Cache 大小。选择 Block Size 16 的另一个原因是与 GPU Warp Size32 threads对齐。每个 Block 的 Attention 计算可以由一个 Warp 独立完成避免 Warp 内分支发散。如果 Block Size 不是 Warp Size 的倍数会导致部分线程空闲降低 GPU 利用率。三、PageAttention 在 Rust 推理引擎中的实现use std::collections::{HashMap, VecDeque}; use std::sync::Arc; use tokio::sync::Mutex; /// Block 配置 const BLOCK_SIZE: usize 16; // 每个 Block 存储 16 个 token /// 物理 Block 在 GPU 显存中的表示 #[derive(Debug, Clone, Copy)] struct PhysicalBlock { /// Block 在 Block 池中的索引 block_id: usize, /// GPU 显存中的虚拟地址 gpu_addr: u64, /// 引用计数多个请求可能共享前缀 Block ref_count: u32, } /// Block 池管理器GPU 显存管理器 struct BlockPool { /// 总 Block 数量 total_blocks: usize, /// 空闲 Block 队列 free_blocks: VecDequePhysicalBlock, /// 所有 Block索引 → Block all_blocks: VecPhysicalBlock, /// 当前已分配的 Block 数量 allocated_blocks: usize, } impl BlockPool { fn new(num_gpu_blocks: usize, base_addr: u64, block_size_bytes: u64) - Self { let mut all_blocks Vec::with_capacity(num_gpu_blocks); let mut free_blocks VecDeque::with_capacity(num_gpu_blocks); for i in 0..num_gpu_blocks { let block PhysicalBlock { block_id: i, gpu_addr: base_addr (i as u64) * block_size_bytes, ref_count: 0, }; all_blocks.push(block); free_blocks.push_back(block); } BlockPool { total_blocks: num_gpu_blocks, free_blocks, all_blocks, allocated_blocks: 0, } } /// 分配一个空闲 Block /// 设计原因使用 VecDeque::pop_front 实现 O(1) 分配 /// 返回的 Block 引用计数设为 1首次分配 fn allocate(mut self) - OptionPhysicalBlock { self.free_blocks.pop_front().map(|mut block| { block.ref_count 1; self.allocated_blocks 1; block }) } /// 释放一个 Block引用计数减 1归零时回收 /// 设计原因使用引用计数支持前缀共享 /// 多个请求可以共享同一个 Block如系统 prompt 的 KV Cache fn free(mut self, block_id: usize) { let block mut self.all_blocks[block_id]; block.ref_count block.ref_count.saturating_sub(1); if block.ref_count 0 { self.free_blocks.push_back(*block); self.allocated_blocks - 1; } } /// 增加引用计数共享 Block 时使用 /// 设计原因前缀共享时多个请求需要访问同一个 Block /// 只有当所有持有者都释放后Block 才真正回收 fn add_ref(mut self, block_id: usize) { self.all_blocks[block_id].ref_count 1; } fn usage_ratio(self) - f64 { self.allocated_blocks as f64 / self.total_blocks as f64 } } /// 逻辑 Block 映射一个请求的 Block 表 #[derive(Debug, Clone)] struct BlockTable { /// token_index → physical_block_id 的映射 /// key: 逻辑 Block 序号0, 1, 2, ... /// value: 物理 Block ID blocks: Vecusize, } impl BlockTable { fn new() - Self { BlockTable { blocks: Vec::new() } } /// 获取指定 token 索引对应的物理 Block 和块内偏移 fn get_physical_location(self, token_idx: usize) - Option(usize, usize) { let logical_block token_idx / BLOCK_SIZE; let offset token_idx % BLOCK_SIZE; self.blocks.get(logical_block).map(|phys_id| (phys_id, offset)) } /// 需要多少个逻辑 Block 来存储给定数量 token fn num_required_blocks(num_tokens: usize) - usize { (num_tokens BLOCK_SIZE - 1) / BLOCK_SIZE } } /// PageAttention 管理器协调 KV Cache 的分配与释放 struct PageAttentionManager { /// GPU 显存 Block 池 block_pool: ArcMutexBlockPool, /// 每个活跃请求的 Block 表 request_blocks: ArcMutexHashMapu64, BlockTable, } impl PageAttentionManager { fn new(num_gpu_blocks: usize) - Self { let block_pool BlockPool::new( num_gpu_blocks, 0x1000_0000, // GPU 显存基地址 2 * BLOCK_SIZE as u64 * 4096 * 32 * 2, // 每 Block 字节数 ); PageAttentionManager { block_pool: Arc::new(Mutex::new(block_pool)), request_blocks: Arc::new(Mutex::new(HashMap::new())), } } /// 为新请求分配 KV Cache Blocks /// 返回分配的 Block 数量 /// 设计原因分配可能部分失败在分配过程中 Block 池耗尽 /// 必须回滚已分配的 Block避免资源泄漏 async fn allocate_for_request( self, request_id: u64, prompt_len: usize, ) - Resultusize, String { let num_blocks BlockTable::num_required_blocks(prompt_len); let mut pool self.block_pool.lock().await; // 检查是否有足够的空闲 Block // 设计原因先检查再分配避免部分分配后失败 // 但此检查与后续分配之间存在 TOCTOU 竞态 // 生产实现中应使用 try_allocate 原子操作 if pool.free_blocks.len() num_blocks { return Err(format!( Insufficient blocks: need {}, available {}, num_blocks, pool.free_blocks.len() )); } let mut table BlockTable::new(); for _ in 0..num_blocks { match pool.allocate() { Some(block) table.blocks.push(block.block_id), None { // 回滚释放已分配的 Block // 设计原因必须保证原子性全部分配或完全不分配 for block_id in table.blocks { pool.free(block_id); } return Err(Allocation failed mid-way.to_string()); } } } let mut requests self.request_blocks.lock().await; requests.insert(request_id, table); Ok(num_blocks) } /// 为生成阶段追加新的 KV Cache Block /// 每次自回归解码生成一个新 token可能触发新 Block 分配 async fn append_slot( self, request_id: u64, current_len: usize, ) - ResultOptionusize, String { // 检查当前是否正好是 Block 边界 if current_len % BLOCK_SIZE ! 0 { return Ok(None); // 无需新 Block } let mut pool self.block_pool.lock().await; let mut requests self.request_blocks.lock().await; let table requests.get_mut(request_id) .ok_or(Request not found)?; match pool.allocate() { Some(block) { table.blocks.push(block.block_id); Ok(Some(block.block_id)) } None { // 无空闲 Block → 触发抢占Preemption // 策略选择 oldest request 驱逐其 Blocks Err(Preemption required: no free blocks.to_string()) } } } /// 释放请求的所有 Blocks async fn free_request(self, request_id: u64) { let mut requests self.request_blocks.lock().await; let mut pool self.block_pool.lock().await; if let Some(table) requests.remove(request_id) { for block_id in table.blocks { pool.free(block_id); } } } /// 前缀共享新请求复用已有请求的 prompt Block /// 例如系统 prompt 被多个请求共享无需重复计算 KV Cache async fn share_prefix_blocks( self, source_request_id: u64, target_request_id: u64, shared_len: usize, ) - Result(), String { let shared_blocks BlockTable::num_required_blocks(shared_len); let requests self.request_blocks.lock().await; let mut pool self.block_pool.lock().await; let source_table requests.get(source_request_id) .ok_or(Source request not found)?; if source_table.blocks.len() shared_blocks { return Err(Source has fewer blocks than requested share length.to_string()); } // 为共享的 Blocks 增加引用计数 for i in 0..shared_blocks { pool.add_ref(source_table.blocks[i]); } let mut target_table BlockTable::new(); target_table.blocks source_table.blocks[..shared_blocks].to_vec(); drop(requests); let mut requests self.request_blocks.lock().await; requests.insert(target_request_id, target_table); Ok(()) } } #[tokio::main] async fn main() { let manager PageAttentionManager::new(1024); // 1024 个物理 Blocks // 分配请求 A 的 KV Cacheprompt50 tokens → 需要 4 个 Blocks let blocks_a manager.allocate_for_request(1, 50).await.unwrap(); println!(Request A allocated {} blocks, blocks_a); // 自回归生成第 17 个 token触发新 Block 分配 let new_block manager.append_slot(1, 17).await.unwrap(); println!(Request A new block: {:?}, new_block); // 请求 B 共享 A 的前 2 个 Block系统 prompt manager.share_prefix_blocks(1, 2, 32).await.unwrap(); println!(Request B shared prefix blocks from request A); // 释放请求 manager.free_request(1).await; manager.free_request(2).await; println!(All requests freed); }前缀共享是 PageAttention 额外带来的重要优化。在 RAG 场景中多个请求共享相同的检索文档 prompt前缀 KV Cache 的复用可以减少 50~80% 的 prefill 计算量。引用计数机制保证共享 Block 在所有使用者释放后才回收。Block 分配失败时的抢占策略需要根据请求优先级选择性驱逐。低优先级请求的 KV Cache 被卸载到 CPU 内存Swap Out待 GPU 有空闲 Block 时再换入。这需要在 GPU 和 CPU 之间做高效的数据传输。四、PageAttention 的性能边界与退化场景额外开销gather操作从 Block 表中查找物理地址每次 Attention 增加约 3% 计算开销Block 管理的元数据内存每 Block 约 64 字节1000 个 Block 仅 64KB退化场景短序列 32 tokensBlock 内碎片占比大连续分配可能更优极高并发 100 requestsBlock 表查找的 TLB 压力增大需增大 Block SizeBeam Search每个 Beam 需要独立的 Block 表fork 操作需要复制 Block 引用显存效率数据连续分配利用率 60%总 KV Cache 浪费约 40%PageAttention利用率 96%Block 内碎片占比约 4%Block Size16PageAttention 的退化场景与应对短序列 32 tokensBlock 内碎片占比大可能 50%连续分配可能更优。应对策略动态选择 Block Size——短提示用 Block Size4长生成用 Block Size32。极高并发 100 requestsBlock 表查找的 TLB 压力增大每次 Attention 需要多次显存访问。应对策略增大 Block Size 到 32 或 64减少 Block 表大小。Beam Search每个 Beam 需要独立的 Block 表fork 操作需要复制 Block 引用引用计数1。当 Beam Width5 时KV Cache 占用放大 5 倍。应对策略使用共享前缀 独立后缀仅 fork 生成部分的 Block 表。Multi-Modal 模型图像文本图像 token 的 KV Cache 远大于文本固定 Block Size 导致图像 Block 浪费显存。应对策略使用可变大小 BlockvLLM 的 v2 引擎支持。与 DeepSpeed Inference 的 Tensor Parallelism 结合当使用 TPTensor Parallelism时每个 GPU 只持有部分注意力头。PageAttention 需要跨 GPU 同步 Block 表——这会导致额外的通信开销。优化方案每个 GPU 独立管理 Block 池仅共享逻辑 Block 表的元数据通过 NCCL AllReduce。五、总结PageAttention 将虚拟内存的分页管理引入 KV Cache消除物理连续分配的外碎片将显存利用率从 60% 提升至 96%。KV Cache 被分割为固定大小的 Block默认 16 tokens通过 Block 表完成逻辑→物理地址映射Attention kernel 使用 gather 操作访问非连续 Block。gather 操作的额外开销约 3%Block 内碎片约 4%净收益显存利用率提升 36%远超额外开销。前缀共享利用引用计数实现多请求复用公共 prompt 的 KV CacheRAG 场景下减少 50~80% prefill 计算。Block 耗尽时的抢占策略将低优先级请求的 KV Cache 交换到 CPU 内存GPU 空闲后换回。