vAttention: 让 LLM 推理更快速、更便捷

大型语言模型 (LLM) 正在改变着各个领域，但其高昂的计算资源消耗和内存占用成为了制约其应用的关键瓶颈。为了解决 LLM 推理中的内存问题，vLLM 提出了 PagedAttention，它通过动态分配内存来避免预先分配过多的内存，从而减少内存浪费。然而，PagedAttention 也存在一些缺陷，例如需要修改注意力内核代码，增加了开发难度，并且性能也受到一定影响。

为了克服 PagedAttention 的局限性，本文介绍了一种全新的内存管理技术——vAttention。vAttention 巧妙地利用了系统本身的虚拟内存机制，将预分配虚拟内存和实际分配物理内存分离，并拓展了分配小内存的 PyTorch 算子，从而实现了更灵活、更高效的内存管理。

PagedAttention 的局限性

PagedAttention 通过动态分配内存，将 KV 缓存分割成多个固定大小的块，并在需要时分配内存。然而，这种方法存在以下几个问题：

• 需要修改注意力内核代码： PagedAttention 导致 KV 缓存存储在非连续的虚拟内存中，需要修改注意力内核代码以支持分页，增加了开发难度。
• 增加了框架复杂度： PagedAttention 需要在服务框架中实现一个内存管理器，负责分配和释放 KV 缓存，增加了框架的复杂度。
• 性能下降： PagedAttention 会增加 GPU 和 CPU 的运行时开销，导致性能下降。

vAttention：突破传统内存管理的限制

vAttention 则采用了一种更巧妙的方案，它保留了 KV 缓存的虚拟内存连续性，并利用系统本身的虚拟内存机制来实现按需物理内存分配。

vAttention 的核心思想：

1. 预留虚拟内存： vAttention 预先分配足够大的虚拟内存空间来存储 KV 缓存，但并不分配物理内存。
2. 按需物理内存分配： vAttention 仅在需要时才分配物理内存页面，并将这些页面映射到虚拟内存中。
3. 利用低级 CUDA 支持： vAttention 利用 CUDA 低级 API，可以对虚拟和物理内存进行细粒度控制，实现更灵活的内存管理。
4. 拓展 PyTorch 接口： vAttention 扩展了 PyTorch 缓存分配器，允许应用程序为张量保留虚拟内存缓冲区，而无需提前提交物理内存。
5. 请求级 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 的论文，并加入了个人理解和分析。