分类： 未分类

大型语言模型 (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 的论文，并加入了个人理解和分析。

近年来，大型语言模型（LLM）在各个领域都取得了显著进展，但其高昂的计算资源消耗和内存占用成为了制约其应用的关键瓶颈。PagedAttention 作为一种动态内存分配机制，在一定程度上解决了 LLM 的内存问题，但也存在一些缺陷，例如内存不连续导致的效率低下和开发难度增加。

本文将介绍一种全新的内存管理技术——vAttention，它可以看作是 PagedAttention 的升级版本，有效地解决了这些问题，并显著提升了 LLM 的性能。

PagedAttention 的局限性

PagedAttention 通过动态分配内存来应对 LLM 自回归生成时内存需求不确定的问题。然而，这种方法存在以下几个不足：

• 内存不连续： PagedAttention 会导致内存分配不连续，需要在框架中引入额外的内存管理代码，增加了开发难度。
• 注意力计算效率降低： kv cache 在注意力计算中无法连续读取，需要手动修改注意力内核代码进行分段读取，降低了计算效率。
• 性能劣于 FlashAttention： 在小数据量时，PagedAttention 的速度明显慢于 FlashAttention，可能是由于 cache 命中率高，连续内存读取效率更高导致的。

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

vAttention 巧妙地利用了虚拟内存机制，将预分配虚拟内存和实际分配物理内存分离，并拓展了分配小内存的 PyTorch 算子，将内存管理从注意力内核中剥离出来。

vAttention 的核心思想:

1. 预留虚拟内存： 创建足够长的虚拟内存，但并不分配物理内存，利用系统本身的虚拟内存管理机制。
2. 按需物理内存分配： 每次优先分配一个物理内存页面，仅当请求已使用完其先前分配的所有物理内存页面时，再分配新的物理内存，并映射到虚拟内存中。
3. 利用低级 CUDA 支持： vAttention 利用 CUDA 低级 API，可以对虚拟和物理内存进行细粒度控制，实现更灵活的内存管理。
4. 拓展 PyTorch 接口： vAttention 扩展了 PyTorch 缓存分配器，允许应用程序为张量保留虚拟内存缓冲区，而无需提前提交物理内存。
5. 请求级 KV-cache 索引： vAttention 使用唯一的整数标识符来定位请求的子张量，并进行相应的内存分配和释放。

vAttention 的优势:

• 连续内存： vAttention 保证了 KV 缓存存储在连续的虚拟内存中，无需提前提交物理内存，避免了内存碎片化问题。
• 性能提升： vAttention 能够无缝接入 FlashAttention 等高性能注意力计算内核，并通过优化内存管理机制，显著提升了 LLM 的性能。

性能评估

实验结果表明，vAttention 在吞吐量和推理时间方面都取得了显著提升，性能比 vLLM 高出 1.97 倍。

总结

vAttention 作为 PagedAttention 的升级版本，有效地解决了传统内存管理机制的缺陷，为 LLM 的应用提供了更强大的支持。它不仅提高了 LLM 的性能，也简化了开发流程，为 LLM 的进一步发展提供了新的方向。

注： 本文内容主要参考了 akaihaoshuai 的知乎文章，并加入了个人理解和分析。