PagedAttention

PagedAttention 是一种注意力算法和内存管理器，用于大型语言模型（LLM）服务：它会将每个请求的 KV 缓存存储在大小固定的小区块中，这些区块可放置在 GPU 内存的任意位置，并通过每个请求各自的区块表进行追踪，类似于操作系统映射虚拟内存页。UC Berkeley 及合作大学的 Woosuk Kwon、Zhuohan Li 和同事在《通过 PagedAttention 高效管理大型语言模型服务的内存》中与 vLLM 服务引擎一同提出了该方法；这篇论文发表于 SOSP 2023。消除连续缓存分配产生的碎片后，单个 GPU 就能同时容纳更多序列；论文报告称，在相同延迟下，其吞吐量比 FasterTransformer 和 Orca 高 2 至 4 倍。

它解决的内存问题#

在文本生成期间，Transformer会保留每个先前词元的键和值向量，避免重复计算。这一缓存会在每层中为每个词元增加一个条目；单个包含 130 亿参数的 LLaMA 模型序列最多可占用 1.7 GB。请求到达时，输出长度尚不可知，因此早期的服务系统会为每个请求预留一块连续内存，并按可能的最大长度来分配。大多数区域都处于空闲状态，区域之间的空隙也太小而无法重新利用。vLLM 团队测量发现，这类系统浪费了 60% 至 80% 的 KV 缓存内存，因此批次大小早在算力用尽之前就已受限。

分页的工作方式#

PagedAttention 会将缓存拆分成多个区块，每个区块保存固定数量词元对应的键和值；在 vLLM 中，默认每个区块保存 16 个词元。只有当请求的最后一个区块填满时，系统才会从任意空闲位置为其分配新区块，并在区块表中记录顺序。注意力机制内核会按照该表逐块收集键和值，因此序列缓存无需连续存放。浪费仅限于每个序列最后一个区块未填满的尾部；vLLM 的测量显示浪费率低于 4%。

下面这段 PyTorch 示例代码将一个请求的键存储在共享池中分散的区块里，再通过该请求的区块表读取这些键：

import torch

block_size, head_dim = 16, 64
key_pool = torch.zeros(64, block_size, head_dim)  # physical blocks shared by every request
free_blocks = [41, 7, 19, 2, 33]  # free blocks are scattered across the pool

# A 40-token request needs ceil(40 / 16) = 3 blocks, wherever they happen to be free
keys = torch.randn(40, head_dim)
num_blocks = -(-len(keys) // block_size)
block_table = [free_blocks.pop(0) for _ in range(num_blocks)]
for pos, key in enumerate(keys):
    key_pool[block_table[pos // block_size], pos % block_size] = key

# The attention kernel reads the request's keys back through its block table
gathered = key_pool[block_table].reshape(-1, head_dim)[: len(keys)]
print(f"block table {block_table}, keys intact: {torch.equal(gathered, keys)}")
print(f"unused slots: {num_blocks * block_size - len(keys)} paged vs {2048 - len(keys)} with a 2,048-token reservation")

它会输出区块表 [41, 7, 19] 和 keys intact: True，然后显示：分页请求有 8 个未使用槽位，而为 2,048 个词元预留的区域有 2,008 个未使用槽位。

序列之间共享区块#

区块表还能让多个序列指向相同的物理区块。并行采样、束搜索以及以相同系统提示开头的请求都能复用同一份共享前缀副本，并通过每个区块的引用计数进行管理。当某个序列需要写入其他序列仍在使用的区块时，引擎会先复制该区块；这一步称为写时复制，借鉴自操作系统。

vLLM 作者报告称，这种共享方式能将并行采样和束搜索的内存使用量最多降低 55%，使吞吐量最多提高 2.2 倍。对完整区块进行哈希可将这一做法扩展到不同请求；这也是 vLLM 实现提示缓存的方式。

相关技术#

PagedAttention 决定缓存的键和值存放在哪里，并可与改变存储量或读取速度的方法结合使用：

• 闪存注意力：对注意力计算进行分块，避免将完整的分数矩阵写入 GPU 主内存。当前引擎会运行从 FlashAttention 衍生的内核，从分页区块中读取数据。
• 连续批处理：每当另一个请求完成时，就接收一个新请求；分页机制会立即将已完成序列的区块归还空闲池，为新请求腾出空间。
• 分组查询注意力：让多个查询头共享键头和值头，从而缩减每个区块的内容。
• FP8缓存量化：以 8 位精度存储键和值，与 16 位格式相比，缓存内存用量大约减半。

Ultralytics LLM 接口会原样转发请求参数，因此通过其 base_url 将带有 n=3 的 Chat Completions 调用发送到本地 vLLM 服务器后，会返回三个共享提示缓存区块的样本，只有彼此不同的后续内容才会占用新内存。