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_len=4096 时的单请求 KV Cache 为 2.5GB。当实际生成长度仅 128 tokens 时，预分配的 96.9% 空间闲置。

更严重的是：连续分配导致严重的外碎片。多个请求的 KV Cache 在显存中交错分配，短请求完成后释放的空间被长请求的活跃 KV Cache 块隔开，无法合并成大块。最终显存利用率仅 60%，但 GPU OOM 频率上升——不是因为总量不够，而是无法分配连续空间。

此外，连续分配策略在动态批处理（Dynamic Batching）场景下表现极差。当 Batch 中某个请求提前结束时，其释放的 KV Cache 空间若小于其他请求所需，就会成为无法利用的"空洞"。PageAttention 的非连续分配天然解决了这个问题——释放的 Block 可以立即分配给任何其他请求，无论其大小。

PageAttention（vLLM 提出的 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 Size（32 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: VecDeque<PhysicalBlock>, /// 所有 Block（索引 → Block） all_blocks: Vec<PhysicalBlock>, /// 当前已分配的 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) -> Option<PhysicalBlock> { 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: Vec<usize>, } 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: Arc<Mutex<BlockPool>>, /// 每个活跃请求的 Block 表 request_blocks: Arc<Mutex<HashMap<u64, 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, ) -> Result<usize, 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, ) -> Result<Option<usize>, 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 Cache（prompt=50 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 tokens）：Block 内碎片占比大，连续分配可能更优
极高并发（> 100 requests）：Block 表查找的 TLB 压力增大，需增大 Block Size
Beam Search：每个 Beam 需要独立的 Block 表，fork 操作需要复制 Block 引用

显存效率数据：

连续分配：利用率 60%，总 KV Cache 浪费约 40%
PageAttention：利用率 96%，Block 内碎片占比约 4%（Block Size=16）

PageAttention 的退化场景与应对：

短序列（< 32 tokens）：Block 内碎片占比大（可能 > 50%），连续分配可能更优。应对策略：动态选择 Block Size——短提示用 Block Size=4，长生成用 Block Size=32。
极高并发（> 100 requests）：Block 表查找的 TLB 压力增大，每次 Attention 需要多次显存访问。应对策略：增大 Block Size 到 32 或 64，减少 Block 表大小。
Beam Search：每个 Beam 需要独立的 Block 表，fork 操作需要复制 Block 引用（引用计数+1）。当 Beam Width=5 时，KV Cache 占用放大 5 倍。应对策略：使用共享前缀 + 独立后缀，仅 fork 生成部分的 Block 表。
Multi-Modal 模型（图像+文本）：图像 token 的 KV Cache 远大于文本，固定 Block Size 导致图像 Block 浪费显存。应对策略：使用可变大小 Block（vLLM 的 v2 引擎支持）。

与 DeepSpeed Inference 的 Tensor Parallelism 结合：
当使用 TP（Tensor 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 Cache，RAG 场景下减少 50~80% prefill 计算。
Block 耗尽时的抢占策略将低优先级请求的 KV Cache 交换到 CPU 内存，GPU 空闲后换回。

企业官网建设流程全解析