PagedAttention 与 vLLM:大模型推理显存管理的工程革命
大模型推理的瓶颈往往不是算力而是显存。一个 70B 参数的模型在推理时KV Cache 的显存占用可能超过模型权重本身。vLLM 引入的 PagedAttention 机制借鉴了操作系统的虚拟内存管理思想从根本上重构了 KV Cache 的存储方式将显存利用率从 20%-40% 提升到 90% 以上。一、KV Cache 的显存困境在自回归生成中每个 token 的注意力计算需要访问之前所有 token 的 Key 和 Value 向量。为避免重复计算这些向量被缓存下来称为 KV Cache。textKV Cache 显存 2 × num_layers × num_heads × head_dim × seq_len × batch_size × dtype_sizetext以 Llama-2-70B 为例单条请求 2048 token 的 KV Cache 就需要约 1.6GB。当并发请求增多时显存碎片化问题变得极为严重——不同请求的序列长度不同分配和释放 KV Cache 时产生大量外部碎片。传统方案使用连续内存分配 KV Cache每条请求预分配最大序列长度的空间。这导致两个问题短请求浪费大量空间长请求可能超出预分配上限。## 二、PagedAttention分页式显存管理### 核心思想PagedAttention 将 KV Cache 组织为固定大小的块block每个块包含固定数量 token 的 KV 向量。逻辑上连续的序列物理上可以分散在不连续的块中。pythonclass BlockTable: 简化的块表管理类比 OS 页表 def __init__(self, num_blocks, block_size, num_kv_heads, head_dim, dtype): self.block_size block_size # 预分配整个 KV Cache 池 [num_blocks, block_size, 2, num_kv_heads, head_dim] self.kv_pool torch.zeros( num_blocks, block_size, 2, num_kv_heads, head_dim, dtypedtype ) self.free_blocks list(range(num_blocks)) self.block_tables {} # seq_id - list of block indices def allocate(self, seq_id, num_tokens): 为序列分配所需块 num_blocks_needed (num_tokens self.block_size - 1) // self.block_size blocks [] for _ in range(num_blocks_needed): if not self.free_blocks: raise OOMError(No free KV cache blocks available) blocks.append(self.free_blocks.pop(0)) self.block_tables[seq_id] blocks return blocks def append_kv(self, seq_id, token_idx, k, v): 向序列追加 KV 向量 block_idx token_idx // self.block_size offset token_idx % self.block_size physical_block self.block_tables[seq_id][block_idx] self.kv_pool[physical_block, offset, 0] k # Key self.kv_pool[physical_block, offset, 1] v # Value def free(self, seq_id): 释放序列占用的所有块 for block_idx in self.block_tables[seq_id]: self.free_blocks.append(block_idx) del self.block_tables[seq_id]text### 块大小的选择策略块大小是 PagedAttention 中影响性能的关键参数它直接决定了显存碎片率和寻址开销的平衡。vLLM 默认使用 16 个 token 作为块大小这个值是在大量实验基础上得出的经验最优值。当块大小为 16 时一个块存储的 KV 向量约为几 KB 到几十 KB取决于模型配置碎片浪费通常不超过 3%。减小块大小可以进一步降低碎片率但会增加块表的大小和注意力计算中的寻址开销。块大小为 4 时块表条目数量是块大小 16 时的 4 倍注意力计算需要更多的间接寻址操作在长序列场景下会导致明显的性能下降。增大块大小则相反——寻址开销减小但碎片率上升尤其当请求的平均输出长度远小于块大小时浪费的显存比例可能达到 20% 以上。### 注意力计算中的分页寻址在注意力计算时PagedAttention 通过块表将逻辑 token 位置映射到物理块位置pythondef paged_attention(q, block_table, kv_pool, block_size): q: [num_heads, head_dim] 当前 token 的 Query block_table: list[int] 该序列映射到的物理块索引 kv_pool: [num_blocks, block_size, 2, num_kv_heads, head_dim] scores [] for block_idx in block_table: # 获取该物理块的所有 KV block_kv kv_pool[block_idx] # [block_size, 2, num_kv_heads, head_dim] keys block_kv[:, 0] # [block_size, num_kv_heads, head_dim] values block_pool[:, 1] # 分块计算注意力分数 block_scores torch.einsum(hd,nhd-hn, q, keys) / math.sqrt(head_dim) scores.append(block_scores) # 拼接所有块的分数做 softmax all_scores torch.cat(scores, dim-1) attn_weights F.softmax(all_scores, dim-1) # 加权求和 values output torch.zeros_like(q) for i, block_idx in enumerate(block_table): values kv_pool[block_idx][:, 1] output torch.einsum(nhd,hn-hd, values, attn_weights[:, i*block_size:(i1)*block_size]) return outputtext## 三、Continuous Batching动态批处理的协同机制PagedAttention 的分页管理使得 Continuous Batching连续批处理成为可能。传统静态批处理要求所有请求同时到达、同时完成导致 GPU 空闲率高。Continuous Batching 在每个 iteration 级别动态加入新请求、移除已完成请求。pythonclass ContinuousBatchScheduler: def __init__(self, max_batch_size): self.max_batch_size max_batch_size self.running [] self.waiting [] def schedule(self): 每次解码步执行调度 # 移除已完成的序列 self.running [s for s in self.running if not s.is_finished()] # 从等待队列补充新请求 while len(self.running) self.max_batch_size and self.waiting: req self.waiting.pop(0) self.running.append(req) return self.runningtext## 四、工程效果与性能数据vLLM 的基准测试显示在 Llama-2-13B 模型上-吞吐量相比 HuggingFace Transformers 提升 14-24 倍-显存利用率从约 40% 提升到 96% 以上-P99 延迟在同等吞吐量下降低 50% 以上显存节省的核心来源是消除了三种浪费内部碎片预分配过量、外部碎片无法利用的间隙、以及重复存储多请求共享前缀时各自缓存。## 五、Prefix Caching共享前缀的进一步优化PagedAttention 的分页结构天然支持前缀共享。当多个请求使用相同的 System Prompt 或 Few-shot 示例时其 KV Cache 可以直接复用pythondef find_prefix_match(new_prompt, cached_prefixes): 查找最长的可复用前缀 best_match None best_len 0 for cached in cached_prefixes: match_len compute_common_prefix_len(new_prompt, cached.tokens) if match_len best_len: best_len match_len best_match cached if best_len MIN_REUSE_THRESHOLD: # 复用前 N 个块的 KV Cache只需计算剩余部分 return best_match.blocks[:best_len // block_size] return Nonetext这一优化在 Agent 场景中效果尤为显著——多个 Agent 共享相同的工具描述和系统指令前缀复用率可达 80% 以上。### 工程实践中的调参经验在将 PagedAttention 部署到生产环境时块大小block_size是一个需要仔细调优的参数。vLLM 默认使用 16但在不同场景下最优值不同。块太小会导致块表过大增加寻址开销块太大会增加内部碎片降低显存利用率。对于短序列密集型场景如客服对话平均序列长度 100-200 token块大小设为 8 效果更好对于长序列场景如文档摘要序列长度 4000 token块大小设为 32 可以减少块表管理开销。另一个关键实践是合理设置 GPU 内存利用率上限。vLLM 默认使用 90% 的 GPU 显存剩余 10% 留给临时计算和框架开销。如果系统中还运行了其他进程如 Embedding 模型推理需要降低这个上限避免 OOM。在生产环境中建议先用 70% 的利用率进行压测确认稳定后逐步提高至 85%-90%。对于多租户场景需要特别注意 KV Cache 的隔离问题。虽然 PagedAttention 的分页机制天然支持物理隔离但不同租户的请求在同一 GPU 上共享权重读取带宽。当某个租户的长序列请求消耗大量 KV Cache 时可能导致其他租户的短请求排队等待。解决方案是根据租户等级设置不同的 max_num_seqs 参数或者使用物理隔离的多 GPU 实例。## 总结PagedAttention 是大模型推理工程中的一次范式转变它将操作系统成熟的虚拟内存管理思想引入了 GPU 显存管理。分页式 KV Cache 不仅解决了显存碎片化问题更解锁了 Continuous Batching 和 Prefix Caching 两大优化。理解这一机制对于构建高吞吐、低延迟的 LLM 推理服务至关重要——当显存不再是瓶颈时GPU 算力才能真正被充分利用。