LLM 前沿技术跟踪：PagedAttention 升级版 vAttention

近年来，大型语言模型（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 的知乎文章，并加入了个人理解和分析。