PagedAttention：KV Cache的虚拟内存（15分钟阅读）

# PagedAttention：面向KV缓存的虚拟内存管理 来源：https://thegustafson.com/blog/paged-attention 系列文章第六篇：推理引擎与服务系统 (https://thegustafson.com/series#arc-6) 我常思考一个现象：计算机科学的某个子领域偶尔会诞生出一个想法，它悄然跨越学科边界，与另一个子领域碰撞，最终彻底改变游戏规则。PagedAttention正是这样一个时刻。其核心思想不过是将操作系统自1960年代起就使用的虚拟内存技术，移植到了语言模型的键值缓存管理中。仅凭这一创见，现代推理引擎就能在单张GPU上承载比传统方案多两到四倍的用户。 本文正是关于这次技术融合的探索。我将为你构建直觉：为什么键值缓存默认会浪费如此多的内存，为何虚拟内存是完美的类比，以及PagedAttention如何以保持注意力内核兼容性的方式实现这一思想。 在深入探讨前，假设你未读过关于键值缓存的文章（https://thegustafson.com/blog/kv-cache-from-first-principles），先做个简要回顾：键值缓存是模型为每个请求存储的注意力键值对，它使模型避免在每个解码步骤重复计算。缓存大小随序列长度线性增长，在长上下文场景下消耗的GPU显存甚至超过模型权重本身。因此它成了稀缺资源——本系列所有服务系统创新，本质上都是围绕如何更高效地管理这一资源。 --- ## 内存碎片化与GPU浪费 当你为大语言模型提供服务时，无法预知每个请求的长度。用户A可能只生成25个词元，用户B可能持续输出2000个。系统承诺支持每个请求最大2048个词元的上下文，但你无法预判哪个用户会耗尽配额。因此朴素的做法显而易见：为每个活跃请求预留`max_context`大小的键值缓存槽位。若模型参数量达700亿，上下文窗口为32k词元，每个请求就会预先占用数GB的GPU显存，且大部分闲置。 这将导致两个叠加问题： **内部碎片化**：仅生成47个词元的请求仍占用2048个槽位，其中2001个始终空闲。若同时有30个并发请求，为实际仅使用1-2k个槽位的总量，却预留了6万个槽位。这些闲置内存无法分配给其他请求，因为它们已被占用。 **外部碎片化**：请求的到达与离开时间不同。当请求A结束，其连续的2048槽位块被释放。但新请求同样需要2048个连续槽位，而空闲区域仅由A留下的1500槽位空洞和另一个请求的800槽位空洞组成——两者无法合并。总空闲内存足够，但形状不匹配。要么进行内存整理（复制数据，代价高昂且可能中断服务），要么坐视内存浪费。 三个请求的示例（长度分别为25、47、12个词元，上下文预算为2048词元）：连续分配方式会预留6144个槽位仅存储84个实际词元；分页分配方式通过边界对齐仅使用96个槽位。在这个小案例中差异达64倍，且这并非特例。 Kwon等人在vLLM论文中量化了这种浪费：在生产负载下，他们测得键值缓存利用率仅20-40%，意味着集群中最昂贵GPU上60-80%的显存处于闲置状态。这正是PagedAttention要解决的问题。 --- ## 本质上仍是虚拟内存 vLLM论文让我豁然开朗的时刻，是意识到它解决的问题与操作系统数十年前解决的问题*完全相同*。 在多道程序设计的早期，每个进程都请求连续的内存块。操作系统需决定如何分配。不同大小的进程动态到来离开，空闲内存逐渐变成无法利用的碎片。即使总空闲内存充足，程序也无法装入。解决方案是**虚拟内存**——这一技术由曼彻斯特大学Atlas计算机开创，到1970年代趋于成熟。 虚拟内存的核心机制简述如下：每个进程拥有连续的虚拟地址空间，但操作系统将其分割为固定大小的页面（通常4KB），并任意存储在物理内存中。每个进程的**页表**将虚拟页号映射到物理帧号。当进程访问虚拟地址`0x1000`时，CPU（借助内存管理单元）查询页表转换关系，从分散的物理帧中获取数据。进程始终认为页面是连续的，而操作系统能以任意方式组织物理内存。 重读上文，将“进程”替换为“请求”，“页”替换为“KV块”，“物理帧”替换为“共享池中的物理块”，“页表”替换为“块表”——这便是PagedAttention。整个体系存在严格的一一对应关系： ``` 操作系统虚拟内存 → PagedAttention --------------------------- ---------------------------- 进程 → 请求（一次提示+生成） 页（4KB） → KV块（例如16个词元） 页表 → 块表（逻辑→物理块索引映射） 物理帧 → 共享池中的物理KV块 缺页中断/分配 → 按需分配新块 写时复制 → 共享前缀的写时复制 ``` 我认为这是理解本文最有效的心智模型。只要记住“面向键值缓存的虚拟内存”，你就能在纸上推导出大部分机制细节。 --- ## 一图解析PagedAttention 这是具体的实现机制：键值缓存不再为每个请求维护一个连续的大数组，而是： - GPU显存维护一个**共享池**，内含固定大小的**物理块**。典型块大小为16个词元，每个块可存储一个请求中连续16个位置的K和V向量。块按需从池中分配，使用完毕后释放回池。 - 每个请求维护一个**逻辑块序列**：L0、L1、L2……按顺序排列。前16个词元存入L0，接下来16个存入L1，依此类推。 - 每个请求有一个小型**块表**，记录逻辑块索引→物理块索引的映射。`table[0] = 7`表示“该请求的逻辑块0存储在池中物理块7的位置”。 三个请求映射到共享物理块池的示例：每个请求的块表非常小（仅几个整数），其指向的物理块在池中分散存储。不相关的请求共享分散的物理内存，如同不相关的进程共享分散的RAM。 注意力计算时，内核不再假设键值是连续的。它遍历块表：“逻辑块0对应物理块7，逻辑块1对应物理块3……”并从实际位置收集K和V向量。每个块有少量间接寻址开销，但由于块规模较大（16词元×隐藏维度），查找的固定成本被大量运算摊薄，最终开销仅为个位数百分比。 其核心优势是**仅分配实际使用的内存**，并向上取整到块大小。47个词元的请求只需3个块（48槽位），每个请求的浪费被限定在`BLOCK_SIZE - 1`个词元内，因此系统总浪费小且恒定，不会随`max_context`增长。Kwon等测量显示，生产负载下的缓存利用率从20-40%提升至约96%。 --- ## 具体实例演示 以上述对比中的三个请求为例：假设我们有64个物理块，每个存储16个词元的键值缓存。三个请求到达： - **请求A**：25个词元，需要`ceil(25/16)=2`个块。调度器分配P3和P7，块表为`[3, 7]`。前16个词元存于P3，后9个存于P7（P7中有7个填充槽位）。 - **请求B**：47个词元，需要3个块。分配P0、P5、P9，块表为`[0, 5, 9]`。47 = 16 + 16 + 15，因此P9中有1个填充槽位。 - **请求C**：12个词元，需要1个块。分配P4，块表为`[4]`。有4个填充槽位。 物理槽位总使用量：2×16 + 3×16 + 1×16 = 96个槽位，存储了84个生成词元。浪费：12个槽位（约12%）。对比连续分配方式（`max_context=2048`）：3 × 2048 = 6144个槽位已预留，实际仅使用84个，浪费率98.6%。 采用分页方式后，池中仍有64 - 6 = 58个空闲块。我们可以轻松再接纳30多个平均长度的请求。而在连续分配方式下，若无额外2048个空闲槽位，即使`max_context=2048`也无法接入新请求。分页分配不会增加内存总量，它只是停止了“每个请求都需要最坏情况预算”的假设。 若请求B再生成20个词元进入第4个块，调度器从池中分配一个新块（如P11），追加到B的块表，下一步注意力计算使用`[0, 5, 9, 11]`。这仅需一次表更新，无需复制或内存整理。 简要回顾：这是键值缓存在每个解码步骤增加一行的过程。PagedAttention不改变缓存内容，仅改变其内存布局。每个16词元块对应池中的一个物理块。 --- 此时类比再次发挥双重价值。操作系统有个精妙技巧叫**写时复制**：当进程创建副本时，内核不会复制全部内存，而是让子进程共享父进程的页面并标记为只读。仅当子进程写入某页面时，才会为其创建独立副本。只要无人修改，内存就保持共享——大多数时候确实如此。 该技巧可直接迁移。在大模型服务中，许多请求共享前缀。聊天应用可能为每条用户消息附加相同的800词元系统提示。并行采样（相同提示，多个采样续写）共享首个采样词元前的所有内容。束搜索同理，且粒度更细：不同束仅最新分叉词元不同。若每个共享前缀都维护独立键值缓存副本，相当于为相同工作重复付费。 采用PagedAttention后，调度器可让**多个块表指向同一物理块**。每个物理块维护引用计数，只要计数大于1，就不允许原地写入。当某请求需要写入（例如即将在仍与其他请求共享的块中生成新词元），引擎会克隆该块，减少原块的引用计数，并更新写入请求的块表指向克隆块，其他请求继续指向原块。 其直觉是：共享前缀在写入前保持共享。系统提示让请求A和B都指向P0和P1；当B生成新词元时，仅被写入的块被克隆（P1→P2），A的视图始终不变。只复制页面，而非整个缓存。 实际意义在于：只要共享前缀*真正被共享*（无人写入），则所有使用该前缀的请求只需为其付费一次。这也是前缀缓存的工作机制。Kwon等人展示了写时复制如何赋能束搜索和并行采样，但在生产环境中，更大的收益来自系统提示。若聊天机器人有2k词元的系统提示，同时服务百个并发用户，PagedAttention+写时复制相比朴素方案可节省约20万词元的缓存。在700亿参数模型的fp16精度下，这意味着数十GB的显存节省。 需注意写时复制是分页基础思想的优化。不实现该特性也能获得PagedAttention的大部分收益。写时复制将“良好”提升至“卓越”，特别适合提示高度复用的负载。SGLang的RadixAttention更进一步，将所有缓存块组织为按词元前缀索引的基数树，使得任意共享前缀即使未被显式声明也能被发现和复用。 --- ## 注意力内核为何能正常工作 初读论文时我有个疑惑：注意力本是密集矩阵运算，如何在非连续内存上运行而不崩溃？ 答案是：在内核层面，注意力运算本就是按块遍历键值的循环。例如FlashAttention将K和V矩阵分块为适合SRAM的块，迭代处理。PagedAttention的块大小与内核原本使用的分块尺寸对齐。每个块内部内存连续，跨块则通过块表查找下一个物理块地址继续。内核仅增加少量指针运算，并非重写注意力机制。 间接寻址确有开销。块过小则块表查找开销比例增大，内存访问更随机；块过大则内部碎片化重现。实践中最佳平衡点是每个块约16个词元，这也是vLLM、TGI和TensorRT-LLM的共识。 同一47词元请求在不同块大小的表现：块大小=4需12次查找但仅浪费1槽位；块大小=64仅需1次查找但浪费17槽位（26%）；块大小=16折中：3次查找，浪费1槽位。这是生态系统收敛的最佳平衡点。 --- ## 调度器机制简述 为展示代码层面的实现，以下是分页KV调度器的精简伪代码。虽非完整可运行引擎，但结构忠实： ```python class PagedScheduler: def __init__(self, block_size, num_blocks): self.block_size = block_size self.free_blocks = list(range(num_blocks)) # 空闲块列表 self.block_tables = {} # 请求ID → 块表 def allocate_block(self): if not self.free_blocks: return -1 # 显存不足 return self.free_blocks.pop() def handle_new_request(self, request_id, prompt_tokens): num_blocks = (len(prompt_tokens) + self.block_size - 1) // self.block_size blocks = [] for _ in range(num_blocks): block = self.allocate_block() if block == -1: raise MemoryError("无法分配足够缓存块") blocks.append(block) self.block_tables[request_id] = blocks def handle_generation(self, request_id): table = self.block_tables[request_id] last_block_tokens = len(prompt_tokens) % self.block_size if last_block_tokens == 0 or last_block_tokens == self.block_size: # 需要分配新块 new_block = self.allocate_block() table.append(new_block) def handle_write_with_cow(self, request_id, block_index): table = self.block_tables[request_id] physical_block = table[block_index] if self.ref_counts[physical_block] > 1: # 执行写时复制 new_block = self.allocate_block() copy_block_data(physical_block, new_block) table[block_index] = new_block self.ref_counts[physical_block] -= 1 self.ref_counts[new_block] = 1 def free_request(self, request_id): for block in self.block_tables[request_id]: self.ref_counts[block] -= 1 if self.ref_counts[block] == 0: self.free_blocks.append(block) del self.block_tables[request_id] ``` 实际vLLM调度器更复杂（分页策略、淘汰、抢占、GPU内核集成），但此版本呈现了核心结构。块表仅是整数列表，分配仅是从空闲列表弹出，写时复制仅需引用计数。唯一精妙之处在于注意力内核内部，且改动是局部的。 --- ## 常见误解 **“PagedAttention会改变模型输出。”** 否。它完全不改变数学运算。无论K和V张量是连续布局还是分散在池中，注意力产生的logits完全相同。唯一区别是内核获取内存的方式。输出词元、采样概率等所有内容都保持逐位一致（GPU上浮点运算的常规非确定性除外）。 **“更小的块总是更好，因为能减少浪费。”** 这更像是权衡。更小的块减少内部碎片（每个请求最后一个块的填充更少），但会增大块表规模，增加注意力内核的每块开销。它们也使内存访问更随机，可能影响缓存性能。Kwon等人发现16个词元对大多数模型是强默认值，这也是主流引擎的共识。部分以短序列为主的负载可从更小块获益；部分均匀长序列负载则偏好32或64。这个数字并非魔法，而是经过充分调优的结果。 **“写时复制主要是束搜索的优化。”** 论文强调了束搜索，但在生产环境中，共享前缀的主要来源是**系统提示**。一个提示被多个用户、多个并发请求共享，都以相同的500-2000个词元开头。写时复制使该缓存前缀在所有请求中仅存储一次——这正是服务负载中大部分节省的来源。 ## 下一步进展