大模型推理内存优化从 KV Cache 分页到连续批处理的工程实践在 LLM 推理的生产部署中真正的瓶颈往往不在 GPU 算力而在显存带宽和容量。以 LLaMA-2 70B 为例FP16 权重就占了 140GB单张 A100-80GB 根本装不下。即便用张量并行把模型分摊到 4 张卡上每个请求的 KV Cache 依然吃紧——seq_len4096 时单请求就要 5GB单卡最多同时跑 12 个请求。更麻烦的是请求长度差异极大短问答只要 50 token长文档摘要能到 8000 token。传统框架给每个请求预分配最大长度的 KV Cache结果短请求浪费了 90% 以上的显存。vLLM 提出的分页注意力PagedAttention就是为了解决这个问题——像操作系统管理虚拟内存一样把 KV Cache 分页按需分配彻底消除预分配浪费。KV Cache 的内存开销与分页机制Transformer 生成第 t 个 token 时需要访问前 t-1 个 token 的 Key 和 Value 向量。为了避免重复计算推理引擎把这些 KV 向量缓存在显存里这就是 KV Cache。KV Cache 的显存占用公式KV_Cache_Size 2 × num_layers × seq_len × num_kv_heads × head_dim × dtype_size × batch_size以 LLaMA-2 70B 为例80 层、8 KV 头、128 维、FP16单请求 KV Cache 2 × 80 × 4096 × 8 × 128 × 2 bytes 1.34 GB batch_size16 时 21.4 GB仅 KV Cache不含模型权重传统预分配方式下请求 150 tokens预分配 4096 slots浪费 98.8%请求 28000 tokens预分配 4096 slots直接溢出请求 3200 tokens预分配 4096 slots浪费 95.1%。分页分配方式下请求 1 分配 2 页零浪费请求 2 分配 32 页按需扩展请求 3 分配 1 页零浪费。vLLM 的 PagedAttention 借鉴了操作系统的虚拟内存分页机制物理块Block固定大小的 KV Cache 存储单元通常 16 个 token 的 KV 向量页表Block Table逻辑块到物理块的映射表每个请求维护独立的页表按需分配生成新 token 时若当前物理块已满分配新的物理块并更新页表块级共享并行采样beam search、n-best时多个输出序列共享前缀的物理块调度器选择可调度的请求查询页表获取逻辑块到物理块的映射返回物理块列表。如果需要分配新物理块并更新页表然后执行 PagedAttention Kernel。GPU 根据 Block Table 间接寻址读取 KV Cache最后返回生成的 token。PagedAttention 的 CUDA Kernel 实现是性能关键。与传统 Attention 不同PagedAttention 需要通过 Block Table 间接寻址 KV Cache先从 Block Table 查找逻辑块对应的物理块 ID再从物理块中读取 KV 向量。这引入了一次额外的 GPU 全局内存访问但通过以下优化缓解Block Table 缓存到共享内存每个 warp 加载自己负责的 Block Table 条目到共享内存KV 向量的合并访问同一 block 内的 KV 向量在内存中连续GPU 可合并coalesce访存请求Prefetch 下一 block在处理当前 block 时异步预取下一个 block 的 KV 向量生产级推理引擎的内存管理与调度以下代码展示了一个简化版的分页 KV Cache 管理器与连续批处理调度器用 Rust 实现涵盖内存池管理、请求调度和块分配逻辑。use std::collections::{HashMap, VecDeque}; /// 物理块 ID type BlockId u32; /// 逻辑块索引 type LogicalBlockIdx u32; /// 每个物理块存储的 token 数量 const BLOCK_SIZE: usize 16; /// KV Cache 物理块池 /// 管理所有物理块的分配与回收 struct BlockPool { /// 空闲物理块列表 free_blocks: VecDequeBlockId, /// 总物理块数量 total_blocks: usize, /// 已使用物理块数量 used_blocks: usize, } impl BlockPool { fn new(num_blocks: usize) - Self { let free_blocks: VecDequeBlockId (0..num_blocks as u32).collect(); BlockPool { free_blocks, total_blocks: num_blocks, used_blocks: 0, } } /// 分配一个物理块 fn allocate(mut self) - OptionBlockId { self.free_blocks.pop_front().map(|id| { self.used_blocks 1; id }) } /// 回收一个物理块 fn free(mut self, block_id: BlockId) { self.free_blocks.push_back(block_id); self.used_blocks - 1; } /// 可用物理块数量 fn available(self) - usize { self.free_blocks.len() } } /// 请求状态 #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum RequestState { /// 等待 prefill处理 prompt WaitingPrefill, /// 正在 decode逐 token 生成 Decoding, /// 已完成 Finished, } /// 推理请求 struct InferenceRequest { /// 请求唯一 ID id: u64, /// prompt token 数量 prompt_len: usize, /// 已生成的 token 数量 generated_len: usize, /// 最大生成长度 max_tokens: usize, /// 逻辑块 → 物理块映射表页表 block_table: VecBlockId, /// 请求状态 state: RequestState, } /// 连续批处理调度器 /// 核心职责决定哪些请求参与当前迭代管理 KV Cache 内存 struct ContinuousBatchingScheduler { /// 物理块池 block_pool: BlockPool, /// 等待 prefill 的请求队列 waiting_queue: VecDequeInferenceRequest, /// 正在 decode 的请求集合 active_requests: HashMapu64, InferenceRequest, /// 每次迭代的最大 batch size max_batch_size: usize, } impl ContinuousBatchingScheduler { fn new( total_gpu_memory_blocks: usize, max_batch_size: usize, ) - Self { ContinuousBatchingScheduler { block_pool: BlockPool::new(total_gpu_memory_blocks), waiting_queue: VecDeque::new(), active_requests: HashMap::new(), max_batch_size, } } /// 添加新请求到等待队列 fn add_request(mut self, req: InferenceRequest) { self.waiting_queue.push_back(req); } /// 调度一次迭代返回参与本次迭代的请求列表 fn schedule(mut self) - Vecu64 { let mut scheduled Vec::new(); // 第一步保留正在 decode 的请求 for (id, req) in self.active_requests { if req.state RequestState::Decoding { // decode 阶段每个请求需要 1 个新物理块 // 当当前 block 已满时 let total_tokens req.prompt_len req.generated_len; let blocks_needed (total_tokens BLOCK_SIZE - 1) / BLOCK_SIZE; let blocks_owned req.block_table.len(); if blocks_needed blocks_owned { // 需要分配新物理块 if self.block_pool.available() 0 { if let Some(block_id) self.block_pool.allocate() { self.active_requests.get_mut(id) .unwrap() .block_table .push(block_id); scheduled.push(id); } // 无可用块此请求本轮被抢占preempted } } else { scheduled.push(id); } } } // 第二步从等待队列中 prefill 新请求 let active_count scheduled.len(); while scheduled.len() self.max_batch_size !self.waiting_queue.is_empty() { let mut req self.waiting_queue.pop_front().unwrap(); // 计算 prefill 所需的物理块数量 let blocks_needed (req.prompt_len BLOCK_SIZE - 1) / BLOCK_SIZE; // 检查是否有足够的物理块 if self.block_pool.available() blocks_needed { // 分配物理块并构建页表 for _ in 0..blocks_needed { if let Some(block_id) self.block_pool.allocate() { req.block_table.push(block_id); } } req.state RequestState::Decoding; let id req.id; self.active_requests.insert(id, req); scheduled.push(id); } else { // 显存不足将请求放回队列头部等待下一轮 self.waiting_queue.push_front(req); break; } } // 第三步检查已完成的请求释放物理块 let finished_ids: Vecu64 self.active_requests.iter() .filter(|(_, req)| { req.generated_len req.max_tokens }) .map(|(id, _)| id) .collect(); for id in finished_ids { if let Some(req) self.active_requests.remove(id) { for block_id in req.block_table { self.block_pool.free(block_id); } } } scheduled } /// 更新请求的生成进度 fn step(mut self, request_id: u64) { if let Some(req) self.active_requests.get_mut(request_id) { req.generated_len 1; if req.generated_len req.max_tokens { req.state RequestState::Finished; } } } }这个实现的核心设计考量调度器采用连续批处理策略——每轮迭代动态决定哪些请求参与完成的请求立即释放资源新请求无缝加入避免传统静态批处理中的填充padding浪费物理块按需分配短请求只占用必要的块数长请求可动态扩展当显存不足时采用抢占策略——暂停低优先级请求释放其物理块给高优先级请求。工程边界与架构取舍PagedAttention 的间接寻址开销Block Table 的间接寻址引入了约 5-8% 的额外延迟。在短序列场景seq_len 256下这一开销占比更显著。对于延迟敏感的在线服务需要权衡内存效率与计算效率——短序列可能更适合传统的连续 KV Cache 分配。Prefill 与 Decode 的计算特征差异Prefill 阶段是计算密集型矩阵乘法Decode 阶段是访存密集型每次只生成 1 token但需读取全部 KV Cache。混合调度时prefill 请求会与 decode 请求竞争 GPU 计算资源导致 decode 延迟抖动。生产中通常采用chunked prefill策略——将长 prompt 分块处理每块插入少量 decode 步平滑延迟。前缀缓存Prefix Caching的一致性挑战多个请求共享相同 system prompt 时可复用前缀的物理块显著降低显存占用。但前缀缓存需要处理版本一致性问题——当 system prompt 更新时需使旧缓存失效。vLLM 通过哈希值标记缓存版本但哈希冲突可能导致错误的缓存命中。量化与 KV Cache 精度的权衡将 KV Cache 从 FP16 量化为 FP8 或 INT4 可将显存占用减半或降至 1/4但量化误差在长序列上累积可能导致生成质量下降。生产中通常对 KV Cache 采用比权重更保守的量化策略如 FP8 而非 INT4并在关键任务场景保留 FP16。CPU Offloading 的延迟陷阱当 GPU 显存不足时将部分 KV Cache 卸载到 CPU 内存可支持更长的序列和更大的 batch。但 PCIe 带宽约 32GB/s for PCIe 4.0 x16远低于 HBM 带宽约 2TB/s for A100offloading 会显著增加 decode 延迟。此方案仅适用于吞吐优先、延迟容忍的离线推理场景。总结大模型推理的内存优化是提升服务吞吐与降低延迟的核心工程路径。KV Cache 的显存占用是推理吞吐的主要瓶颈分页管理可将显存利用率从 20-40% 提升至 90% 以上连续批处理通过动态调度消除了静态批处理的 padding 浪费但需要处理 prefill/decode 的资源竞争PagedAttention 的间接寻址开销在短序列场景下需要权衡量化、前缀缓存和 CPU Offloading 是重要的补充优化手段但各有适用边界。推理优化的本质是在显存容量、计算带宽和延迟约束之间寻找最优操作点。