外观
PagedAttention 内存管理与 KV 缓存优化
内容整理自学习笔记,仅供面试备考参考;不构成录用、培训或考试承诺。
1. vLLM 用于大模型并行推理加速存在什么问题?
在 vLLM 中,LLM 服务的性能受到内存瓶颈的限制。在自回归解码过程中,LLM 的所有输入标记都会生成其 key 和 value 张量,并且这些张量保存在 GPU 内存中以生成下一个 token。这些缓存的 key 和 value 张量通常称为 KV 缓存。
KV 缓存的特点
| 特点 | 说明 |
|---|---|
| 占用大 | LLaMA-13B 中的单个序列最多占用 1.7GB |
| 动态变化 | 其大小取决于序列长度,序列长度变化很大且不可预测 |
现有系统由于碎片和过度预留而浪费了 60% - 80% 的内存。
2. vLLM 如何优化大模型并行推理加速?
vLLM 引入了 PagedAttention,这是一种受操作系统中虚拟内存和分页的经典思想启发的注意力算法。
3. 什么是 PagedAttention?
与传统的注意力算法不同,PagedAttention 允许在不连续的内存空间中存储连续的 key 和 value。
核心思想:将操作系统的虚拟内存和分页机制引入到 LLM 服务中。
4. PagedAttention 如何存储连续的 key 和 value?
具体来说,PagedAttention 将每个序列的 KV 缓存划分为块,每个块包含固定数量 token 的 key 和 value。在注意力计算过程中,PagedAttention 内核有效地识别并获取这些块。
类比映射关系
| 操作系统概念 | PagedAttention 对应 |
|---|---|
| 页面(Page) | 块(Block) |
| 字节(Byte) | Token |
| 进程(Process) | 序列(Sequence) |
| 虚拟内存 | 连续逻辑块 |
| 物理内存 | 非连续物理块 |
序列的连续逻辑块通过块表映射到非连续物理块。当新 token 生成时,物理块会按需分配。
5. PagedAttention 技术细节
5.1 内存效率
| 指标 | 传统系统 | PagedAttention |
|---|---|---|
| 内存浪费率 | 60% - 80% | 低于 4% |
| 内存浪费来源 | 碎片和过度预留 | 仅在序列的最后一个块中 |
内存效率的提高允许系统将更多序列一起批处理,提高 GPU 利用率,从而显著提高吞吐量。
5.2 内存共享
PagedAttention 的另一个关键优势是高效的内存共享。例如,在并行采样中,从同一提示生成多个输出序列时,提示的计算和内存可以在输出序列之间共享。
PagedAttention 通过其块表自然实现内存共享,与进程共享物理页的方式类似。
5.3 性能提升效果
| 优化场景 | 效果 |
|---|---|
| 内存共享降低开销 | 并行采样和波束搜索的内存占用降低高达 55% |
| 吞吐量提升 | 吞吐量提高高达 2.2 倍 |
6. PagedAttention 如何实现安全共享?
6.1 共享问题
PagedAttention 中的不同序列可以通过将其逻辑块映射到同一物理块来共享块,这就涉及到如何安全共享的问题。
6.2 解决方案
| 方面 | 说明 |
|---|---|
| 动机 | 不同序列需要共享物理块时,需确保数据一致性 |
| 思路 | PagedAttention 跟踪物理块的引用计数并实现 Copy-on-Write(写时复制) 机制 |
Copy-on-Write 机制确保当一个序列需要修改共享块时,系统会自动复制该块,而不影响其他序列。
7. PagedAttention 源码介绍
PagedAttention 是 vLLM 背后的核心技术,vLLM 是 LLM 推理和服务引擎,支持各种具有高性能和易于使用的界面的模型。
从 vLLM 的源码中可以看出,vLLM 是如何对 HuggingFace Models 上的模型进行推理优化的:
PagedAttention 核心组件:
├── 块表(Block Table) → 逻辑块到物理块的映射
├── 引用计数(Reference Count) → 跟踪物理块的使用情况
├── Copy-on-Write → 安全的内存共享机制
└── 按需分配(On-demand Allocation) → 新 token 生成时分配物理块PagedAttention 通过将操作系统的内存管理思想引入 LLM 推理,实现了接近最佳的内存使用效率,是 vLLM 高性能的核心秘诀。