Skip to content

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 高性能的核心秘诀。