PagedAttention：突破大语言模型内存瓶颈的分页式注意力革命

PagedAttention：突破大语言模型内存瓶颈的分页式注意力革命

在当今大语言模型服务的浪潮中，一个看似简单却极具挑战性的问题困扰着开发者：如何高效管理数十亿参数模型推理时产生的键值缓存？传统方法面临的内存碎片和低利用率问题，正成为制约LLM服务吞吐量的关键瓶颈。

1 内存管理的困境与突破

在大语言模型推理过程中，自注意力机制需要维护一个不断增长的键值（KV）缓存，用于存储先前所有令牌的键和值向量。传统方法为每个序列预分配连续内存空间，这种做法导致了严重的内存碎片问题：由于序列长度可变且不可预测，内存分配和释放会产生大量无法利用的碎片空间。

更糟糕的是，在并行采样和束搜索等场景中，多个候选序列共享大部分前缀令牌，但传统方法无法有效利用这种共享特性，导致内存冗余和利用率低下。研究表明，这种低效的内存管理可使实际内存利用率降至不足50%，严重限制了系统吞吐量和可扩展性。

PagedAttention技术的出现彻底改变了这一局面。这一创新方法借鉴操作系统虚拟内存分页的思想，通过分块机制将KV缓存划分为固定大小的逻辑块，实现了非连续内存存储，从而显著提升内存利用率并减少碎片。

2 核心架构与工作原理

2.1 分块存储与映射机制

PagedAttention的核心创新在于将每个序列的KV缓存分割为可配置大小的逻辑块。每个块包含固定数量令牌（如16或32个）的键和值向量，这些逻辑块通过块表（block table）动态映射到GPU内存中的物理块。

这种设计类似于操作系统的页表机制：块表维护每个序列的逻辑块到物理块的映射关系，使得连续的键值数据可以存储在非连续的物理内存空间中。内存浪费因此被限制在单个块内，大幅提高了内存利用率。

# 伪代码示例：块表映射机制
class BlockTable:def __init__(self):self.logical_to_physical = {}  # 逻辑块到物理块的映射self.free_blocks = []          # 空闲物理块队列def allocate_block(self, logical_block_id):if self.free_blocks:physical_block = self.free_blocks.pop()else:physical_block = allocate_physical_memory()self.logical_to_physical[logical_block_id] = physical_blockreturn physical_block

2.2 物理块共享与写时复制

PagedAttention支持跨序列的物理块共享，这在并行采样和束搜索场景中特别有价值。当多个候选序列共享相同前缀时，它们可以映射到相同的物理KV块，仅对最后逻辑块采用写时复制（copy-on-write）机制。

实验数据表明，这种共享机制可节省37.6%-66.3%的内存使用。每个物理块维护引用计数，确保共享块的安全管理。当多个逻辑块映射到同一物理块时，仅在写入新生成令牌时才执行复制操作。