大模型推理内存优化：从 KV Cache 分页到连续批处理的工程实践
在 LLM 推理的生产部署中，真正的瓶颈往往不在 GPU 算力，而在显存带宽和容量。以 LLaMA-2 70B 为例，FP16 权重就占了 140GB，单张 A100-80GB 根本装不下。即便用张量并行把模型分摊到 4 张卡上，每个请求的 KV Cache 依然吃紧——seq_len=4096 时，单请求就要 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_size=16 时 = 21.4 GB（仅 KV Cache，不含模型权重）传统预分配方式下，请求 1（50 tokens）预分配 4096 slots，浪费 98.8%；请求 2（8000 tokens）预分配 4096 slots，直接溢出；请求 3（200 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: VecDeque<BlockId>, /// 总物理块数量 total_blocks: usize, /// 已使用物理块数量 used_blocks: usize, } impl BlockPool { fn new(num_blocks: usize) -> Self { let free_blocks: VecDeque<BlockId> = (0..num_blocks as u32).collect(); BlockPool { free_blocks, total_blocks: num_blocks, used_blocks: 0, } } /// 分配一个物理块 fn allocate(&mut self) -> Option<BlockId> { 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: Vec<BlockId>, /// 请求状态 state: RequestState, } /// 连续批处理调度器 /// 核心职责：决定哪些请求参与当前迭代，管理 KV Cache 内存 struct ContinuousBatchingScheduler { /// 物理块池 block_pool: BlockPool, /// 等待 prefill 的请求队列 waiting_queue: VecDeque<InferenceRequest>, /// 正在 decode 的请求集合 active_requests: HashMap<u64, 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) -> Vec<u64> { 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: Vec<u64> = 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 A100），offloading 会显著增加 decode 延迟。此方案仅适用于吞吐优先、延迟容忍的离线推理场景。
总结
大模型推理的内存优化是提升服务吞吐与降低延迟的核心工程路径。KV Cache 的显存占用是推理吞吐的主要瓶颈，分页管理可将显存利用率从 20-40% 提升至 90% 以上；连续批处理通过动态调度消除了静态批处理的 padding 浪费，但需要处理 prefill/decode 的资源竞争；PagedAttention 的间接寻址开销在短序列场景下需要权衡；量化、前缀缓存和 CPU Offloading 是重要的补充优化手段，但各有适用边界。推理优化的本质是在显存容量、计算带宽和延迟约束之间寻找最优操作点。
