大型语言模型 (LLM) 正在改变着各个领域，但其高昂的计算资源消耗和内存占用成为了制约其应用的关键瓶颈。为了解决 LLM 推理中的内存问题，vLLM 提出了 PagedAttention，它通过动态分配内存来避免预先分配过多的内存，从而减少内存浪费。然而，PagedAttention 也存在一些缺陷，例如需要修改注意力内核代码，增加了开发难度，并且性能也受到一定影响。
友情链接：ACEJoy
为了克服 PagedAttention 的局限性，本文介绍了一种全新的内存管理技术——vAttention。vAttention 巧妙地利用了系统本身的虚拟内存机制，将预分配虚拟内存和实际分配物理内存分离，并拓展了分配小内存的 PyTorch 算子，从而实现了更灵活、更高效的内存管理。
PagedAttention 的局限性
PagedAttention 通过动态分配内存，将 KV 缓存分割成多个固定大小的块，并在需要时分配内存。然而，这种方法存在以下几个问题：
- 需要修改注意力内核代码： PagedAttention 导致 KV 缓存存储在非连续的虚拟内存中，需要修改注意力内核代码以支持分页，增加了开发难度。
- 增加了框架复杂度： PagedAttention 需要在服务框架中实现一个内存管理器，负责分配和释放 KV 缓存，增加了框架的复杂度。
- 性能下降： PagedAttention 会增加 GPU 和 CPU 的运行时开销，导致性能下降。
vAttention：突破传统内存管理的限制
vAttention 则采用了一种更巧妙的方案，它保留了 KV 缓存的虚拟内存连续性，并利用系统本身的虚拟内存机制来实现按需物理内存分配。
vAttention 的核心思想：
- 预留虚拟内存： vAttention 预先分配足够大的虚拟内存空间来存储 KV 缓存，但并不分配物理内存。
- 按需物理内存分配： vAttention 仅在需要时才分配物理内存页面，并将这些页面映射到虚拟内存中。
- 利用低级 CUDA 支持： vAttention 利用 CUDA 低级 API，可以对虚拟和物理内存进行细粒度控制，实现更灵活的内存管理。
- 拓展 PyTorch 接口： vAttention 扩展了 PyTorch 缓存分配器，允许应用程序为张量保留虚拟内存缓冲区，而无需提前提交物理内存。
- 请求级 KV-cache 索引： vAttention 使用唯一的标识符来定位每个请求的子张量，并进行相应的内存分配和释放。
vAttention 的优势：
- 连续内存： vAttention 保证了 KV 缓存存储在连续的虚拟内存中，避免了内存碎片化问题。
- 性能提升： vAttention 能够无缝接入 FlashAttention 等高性能注意力计算内核，并通过优化内存管理机制，显著提升了 LLM 的性能。
- 更易于移植： vAttention 能够直接使用现有的注意力内核，无需修改代码，方便移植到不同的 LLM 模型和框架中。
性能评估
实验结果表明，vAttention 在吞吐量和推理时间方面都取得了显著提升，性能比 vLLM 高出 1.97 倍。
总结
vAttention 作为 PagedAttention 的升级版本，有效地解决了传统内存管理机制的缺陷，为 LLM 的应用提供了更强大的支持。它不仅提高了 LLM 的性能，也简化了开发流程，为 LLM 的进一步发展提供了新的方向。
注： 本文内容主要参考了 vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention 的论文，并加入了个人理解和分析。