为什么KV Cache如此重要? 在自回归生成中,每生成一个新token都需要计算它与之前所有token的注意力。如果不缓存,生成第N个token需要重新计算前N-1个token的K和V,复杂度为O(N²)。KV Cache将已计算的K/V存储起来,使每步生成只需O(N)的计算。
但缓存本身带来了巨大的显存压力。以Llama-3-70B为例,FP16精度下单个token的KV Cache占用约320KB。在32K上下文、batch_size=32的场景下,KV Cache总量达到约320GB——远超模型本身的140GB。
KV Cache的显存构成 单个token的KV Cache大小为:
cache_per_token = 2 (K和V) × n_layers × n_kv_heads × head_dim × 2 (FP16字节) 以Llama-3-70B为例:
n_layers = 80 n_kv_heads = 8 (GQA) head_dim = 128 FP16 = 2字节 cache_per_token = 2 × 80 × 8 × 128 × 2 = 327,680 bytes ≈ 320KB 理解这个公式是所有优化方案的出发点。
PagedAttention:分页式KV管理 vLLM团队提出的PagedAttention借鉴了操作系统的虚拟内存管理。核心思想:
块式管理 将KV Cache分为固定大小的块(Block),每个块存储若干token的KV。块不需要在物理上连续,通过块表(Block Table)映射逻辑位置到物理位置。
class PagedKVCache: def __init__(self, n_blocks, block_size, n_layers, n_kv_heads, head_dim): self.block_size = block_size # 每块的token数 self.n_blocks = n_blocks # 预分配所有块的存储 self.k_blocks = torch.zeros(n_blocks, n_layers, block_size, n_kv_heads, head_dim) self.v_blocks = torch.zeros(n_blocks, n_layers, block_size, n_kv_heads, head_dim) self.free_blocks = list(range(n_blocks)) self.block_tables = {} # seq_id -> [block_indices] def allocate(self, seq_id, n_tokens): n_blocks_needed = (n_tokens + self.block_size - 1) // self.block_size blocks = [] for _ in range(n_blocks_needed): if not self.free_blocks: raise OOMError("No free KV cache blocks") blocks.append(self.free_blocks.pop()) self.block_tables[seq_id] = blocks def append(self, seq_id, token_idx, k, v): block_idx = token_idx // self.block_size offset = token_idx % self.block_size physical_block = self.block_tables[seq_id][block_idx] self.k_blocks[physical_block, :, offset] = k self.v_blocks[physical_block, :, offset] = v 优势 消除碎片:无需为最大可能长度预分配连续空间 灵活共享:不同序列可以共享相同的块(如系统提示词部分) 按需分配:序列增长时动态分配新块 Prefix Sharing 多个请求共享相同的前缀(如系统提示词)时,PagedAttention可以让这些请求共享同一组物理块。这在对话AI场景中可节省30-50%的显存。
...