vLLM推理框架源码初探：PagedAttention的显存管理设计

vLLM推理框架源码初探PagedAttention的显存管理设计一、KV Cache 的内存碎片是推理吞吐的第一大隐形杀手在大模型推理中KV Cache 的显存管理比模型权重本身更棘手。原因是每生成一个新 tokenKV Cache 就需要扩展一行——这类似于一个不断 append 的动态数组。在传统的实现中如 HuggingFace Transformers系统为每个请求预分配最大生成长度 × batch_size 的连续显存空间。这导致两个问题显存浪费大多数请求的实际生成长度远小于最大长度预分配的空间大量闲置。无法灵活调度当一个请求完成后其预留的大块连续显存可能无法被其他请求使用碎片化即使总空闲显存足够。vLLM 的 PagedAttention 借鉴了操作系统的虚拟内存管理思想将 KV Cache 切分为固定大小的块block每个块可以独立分配和释放。请求的 KV Cache 不再是连续的大数组而是一个块映射表类似页表。这种设计将显存碎片率从传统方案的 40-60% 降低到 5%。flowchart TB subgraph Traditional[传统方案: 连续预分配] T1[请求1: 预分配 2048 token] T2[请求2: 预分配 2048 token] T3[请求3: 预分配 2048 token] T1 -- T1W[实际只用 500 token → 75% 浪费] T2 -- T2W[实际只用 1200 token → 40% 浪费] end subgraph Paged[PagedAttention: 分块管理] P1[Block Table 请求1] P2[Block Table 请求2] P3[Block Table 请求3] P1 -- B1[Block 0] P1 -- B2[Block 3] P1 -- B3[Block 7] P2 -- B4[Block 1] P2 -- B5[Block 5] P3 -- B6[Block 2] P3 -- B7[Block 4] P3 -- B8[Block 6] P3 -- B9[Block 8] end style Traditional fill:#ffebee style Paged fill:#e8f5e9二、PagedAttention 的核心数据结构PagedAttention 的设计围绕三个核心数据结构Logical Token → Physical Block 映射表Block Table每个序列维护一个块索引列表。第 i 个逻辑 token 属于block_idx i // block_size在 block 内的偏移为i % block_size。物理块分配器维护一个空闲物理块列表。当请求需要新块时从空闲列表中分配请求完成时回收所有块到空闲列表。Copy-on-Write 机制用于并行采样当多个序列共享相同的 prompt如 beam search 或并行采样时它们在 prompt 阶段共享相同的物理块。只有当某个序列的生成导致 KV Cache 需要修改时才触发块的复制。from typing import List, Optional, Dict, Tuple from dataclasses import dataclass, field import torch dataclass class BlockTable: 单个序列的块映射表。 类似操作系统的页表将逻辑 token 位置映射到物理显存块。 block_ids: List[int] field(default_factorylist) num_tokens: int 0 property def num_blocks(self) - int: return len(self.block_ids) def append_block(self, physical_block_id: int): self.block_ids.append(physical_block_id) def get_physical_block(self, token_idx: int, block_size: int) - int: 根据 token 索引获取对应的物理块 ID。 block_idx token_idx // block_size if block_idx len(self.block_ids): raise IndexError( fToken {token_idx} 对应的 block {block_idx} 尚未分配 ) return self.block_ids[block_idx] def get_offset(self, token_idx: int, block_size: int) - int: 获取 token 在块内的偏移位置。 return token_idx % block_size class BlockManager: 物理块分配器。 设计决策 - 使用简单的自由列表free list而非伙伴系统buddy system 因为所有块大小相同不需要合并/分裂操作。 - 引用计数用于支持 Copy-on-Write。 def __init__(self, num_blocks: int, block_size: int): self.num_blocks num_blocks self.block_size block_size # 自由块列表 self.free_blocks list(range(num_blocks)) # 块引用计数ref_count[i] 表示有多少序列引用块 i self.ref_count [0] * num_blocks def allocate(self, num_blocks: int 1) - List[int]: 分配一个或多个物理块。 if len(self.free_blocks) num_blocks: raise RuntimeError( f显存不足: 需要 {num_blocks} 个块, f仅剩 {len(self.free_blocks)} 个空闲块 ) allocated [] for _ in range(num_blocks): block_id self.free_blocks.pop() self.ref_count[block_id] 1 allocated.append(block_id) return allocated def free(self, block_ids: List[int]): 释放物理块引用计数减一归零时回收。 为什么使用引用计数而非直接释放 Copy-on-Write 导致多个序列共享同一物理块。 只有所有引用者都释放后块才能真正被回收。 for block_id in block_ids: self.ref_count[block_id] - 1 if self.ref_count[block_id] 0: self.free_blocks.append(block_id) elif self.ref_count[block_id] 0: raise RuntimeError( f块 {block_id} 的引用计数为负, 可能存在重复释放的 bug ) def copy_on_write(self, block_id: int) - int: Copy-on-Write: 复制块并返回新块的 ID。 if self.ref_count[block_id] 1: return block_id new_block_id self.allocate(1)[0] # 实际应用中需要拷贝块内的 KV Cache 数据 # kv_cache_data[new_block_id] kv_cache_data[block_id].clone() self.ref_count[block_id] - 1 return new_block_id def usage_ratio(self) - float: 当前显存块使用率。 return 1.0 - len(self.free_blocks) / self.num_blocks三、PagedAttention 的 Kernel 实现要点PagedAttention 的实际推理 kernel 需要处理一个额外复杂性每个 token 的 K/V 可能分散在不同的物理块中。标准 Attention 假设 K 和 V 是连续存储的矩阵[seq_len, head_dim]而 PagedAttention 需要先通过 Block Table 查找每个 token 的物理位置然后从分散的块中 gather K/V。这引入了一个计算效率上的挑战分散的内存访问模式对 GPU 缓存不友好。vLLM 的解决方案是在 kernel 中先按 block 遍历而非按 token 遍历充分利用一个 block 内的连续内存访问的缓存局部性。四、PagedAttention 的局限固定 block_size 的粒度问题block_size 太小 → 块表开销大、kernel 循环次数多block_size 太大 → 内部碎片块内未使用的空间被浪费。vLLM 默认 block_size16对大多数场景这是一个合理的折中。Copy-on-Write 的延迟开销并行采样场景中当多个序列的生成路径分化时触发的块复制是同步操作可能引入突发延迟。预填充阶段的大块分配Prompt 处理阶段一次性需要大量连续 token 的 KV Cache。如果 prompt 很长可能导致块分配失败即使总空闲块数足够但空闲块的分布太散。五、总结vLLM 的 PagedAttention 通过借鉴操作系统的虚拟内存管理解决了 KV Cache 的显存碎片问题将 KV Cache 切分为固定大小的块通过 Block Table 映射逻辑序列到物理块。引用计数 Copy-on-Write 支持高效的并行采样beam search、并行生成。显存利用率从传统方案的 40-60% 提升到 90%。分块内存访问模式对 GPU kernel 设计要求更高是实现的主要工程挑战。