PagedAttention
PagedAttention 是一种在 LLM 推理中高效管理 KV 缓存(KV cache)的内存优化方法。它带来的主要服务收益并非来自更快的注意力核函数,而是来自服务引擎如何分配和管理 KV 缓存内存。注意力核函数实现了其中的部分机制(它们通过查找表读取 KV 块),但真正的收益在服务层,这也是为什么本页位于"推理优化"章节。关于内核级别的注意力效率,请参阅 FlashAttention。
注意力机制与 KV 缓存
注意力(attention)是让 Transformer 评估 token 之间关联强度的机制。对于每个 token,模型会计算三个向量:
- 查询(Query,Q):当前 token 在寻找什么
- 键(Key,K):每个 token 为匹配提供了什么
- 值(Value,V):每个 token 贡献的内容
模型将查询与键进行比较以产生注意力分数,用 softmax 进行归一化,并利用得到的权重对值做加权求和。
在自回归生成过程中,每个新 token 都需要所有先前 token 的键和值。服务引擎不会在每一步重新计算它们,而是将它们存储在 KV 缓存中。该缓存会随序列长度增长,并可能跨并发请求占用大量 GPU 内存。引擎如何分配这部分内存,正是 PagedAttention 要解决的问题。
更多信息请参阅注意力机制。
为什么连续分配的 KV 缓存会浪费内存
通常情况下,KV 缓存会占用一大块内存,因为它被存储为一个巨大的连续内存块。这可能导致内存碎片化或空间浪费,因为你即使不会完全填满,也必须预留一大块内存。
具体来说,早期的服务引擎常常将 KV 缓存分配为按最坏情况预留的连续张量。简化的形状公式为:
2 × num_layers × num_heads × head_dim × max_seq_len
这种分配在每个活跃请求上都会发生,而且是在计入批大小和每元素字节数之前。它虽然简单,但假设了每个请求都会用到最大序列长度。真实流量是变化的:一个请求可能生成很短的答案,另一个可能维持很长的对话,还有一个可能提前停止。如果每个请求都为 max_seq_len 预留内存,那么大部分预留的 GPU 内存可能一直处于闲置状态。
结果是有效批大小更低、内存碎片化更严重、并发请求数更少。
更多信息请参阅 PagedAttention 博客文章。
PagedAttention 是如何工作的?
PagedAttention 将这个巨大的内存块拆分成更小的块,就像一本书中的页面一样。换句话说,KV 缓存以非连续的内存块存储。然后它使用查找表来跟踪这些块。LLM 只加载它需要的块,而不是一次性加载所有内容。
这样可以节省内存,并使整个过程更高效。如果有需要,它甚至允许在不同输出之间共享相同的块。
最初的 PagedAttention 论文报告称,不使用 PagedAttention 时,只有 20.4%-38.2% 的已分配 KV 缓存内存被用于存储实际的 token 状态,其余部分都因碎片化而浪费。相比之下,PagedAttention 将 KV 缓存的内存浪费降低到接近零。
这就是 PagedAttention 的意义超越单个注意力核函数的原因。它为服务引擎提供了一个更好的 KV 缓存内存分配器,从而使得诸如连续批处理、前缀缓存和 KV 缓存卸载等技术更容易组合使用。
PagedAttention 最初由 vLLM 实现。此后,MAX 等其他项目也采纳并实现了它。它们将块大小暴露为服务标志,尽管名称可能不同:
MAX
max serve --model meta-llama/Llama-3.1-8B-Instruct \
--kv-cache-page-size 256
vLLM
vllm serve --model meta-llama/Llama-3.1-8B-Instruct \
--block-size 16