Vattention：一种无需PagedAttention的LLM服务动态内存管理方案

本文向您介绍Ramya Prabhu、Ajay Nayak（印度科学学院）、Jayashree Mohan、Ramchandran Ramjee与Ashish Panwar（均来自微软研究院印度班加罗尔团队）的研究成果。这篇论文于2025年3月30日至4月3日在荷兰鹿特丹举行的第三十届ACM编程语言和操作系统体系结构支持国际会议（ASPLOS ‘25）第一卷上发表。

本研究属于计算机科学领域，特别是大型语言模型（LLM）推理服务系统的软件与系统优化方向。当前，LLM被广泛应用于聊天机器人、搜索引擎和代码助手等场景，其推理效率的优化至关重要。其中，批处理（Batching）是提升服务吞吐量的关键技术，但这要求对存储于GPU内存中的关键值缓存（KV Cache）进行高效管理。KV Cache随着请求的处理逐令牌（token）增长，但其最终大小在请求开始时未知，这给内存分配带来了挑战。早期系统如Orca和FasterTransformer采用静态预分配方式（根据模型支持的最大上下文长度预留空间），导致了严重的内部内存碎片化，从而限制了批处理规模与吞吐量。为了解决此问题，vLLM框架引入了PagedAttention，它借鉴操作系统中的按需分页（Demand Paging）思想，将KV Cache划分为固定大小的块，并在需要时动态分配GPU内存，极大地缓解了碎片问题，并成为LLM服务系统中动态内存管理的实际标准。

然而，论文作者们发现，PagedAttention在尝试实现物理内存的动态分配时，付出了显著的代价：它改变了KV Cache在虚拟内存中的布局，使其从连续变为非连续。这种设计带来了一系列非平凡的问题：首先，它需要重写注意力计算内核（Kernel），以便能够访问分布在非连续内存块中的KV Cache数据。这给程序员带来了巨大挑战，并且使得系统难以跟上快速迭代的注意力内核优化研究（例如FlashAttention系列），导致性能落后。其次，它迫使服务框架在用户空间实现一套内存管理器，以拼接动态分配的虚拟内存块，这实质上重复了操作系统本应提供的虚拟到物理地址转换功能，增加了冗余。第三，它在CPU和GPU的执行关键路径上都引入了运行时开销。例如，PagedAttention内核需要额外的代码来查找块表（Block-table）和执行分支，导致其性能通常低于非分页（Non-paged）内核（论文数据显示预填充阶段可慢37%-42%）；同时，CPU端准备块表也可能带来高达10%的延迟开销。表格1通过vLLM、FlashAttention-2、TensorRT-LLM和FlashInfer等系统的实例，具体说明了PagedAttention方法在复杂性、性能和可维护性方面面临的挑战。

基于此，本研究提出了一种新颖的方案——Vattention。其核心目标是：在保留KV Cache虚拟内存连续性的同时，实现物理内存的动态分配，从而避免PagedAttention的上述缺陷。Vattention的基本原理是利用CUDA虚拟内存管理（VMM）API，将虚拟内存与物理内存的分配解耦。具体设计流程如下：

1. 虚拟内存预分配：在系统初始化时，Vattention为每个模型层（Layer）的K Cache和V Cache分别预分配一个大型、连续的虚拟内存缓冲区（Virtual Tensor）。该缓冲区的大小足以容纳配置的最大批处理量（Batch Size）下，每个请求达到模型最大上下文长度（Context Length）时所需的KV Cache。由于64位系统虚拟地址空间极其充裕（例如每个进程用户空间可达128TB），这种预分配不会造成问题。这保证了整个服务生命周期内，KV Cache在虚拟地址层面始终是连续的。

2. 请求级索引与物理内存按需映射：每个请求被分配一个唯一的请求标识符（reqid），该标识符用于在预分配的虚拟缓冲区中定位属于该请求的、连续的子张量（Sub-tensor）。在每次模型前向传播（Forward Pass）之前，服务框架调用Vattention的step API，传入当前所有活跃请求的上下文长度。Vattention内部会检查每个reqid对应的KV Cache子张量区域是否有足够的物理内存页面（Page）与之映射。如果没有，则通过CUDA VMM API动态地将所需数量的物理页面映射到该虚拟内存区域。物理内存的分配单位是“页组”（Page-group），其大小可配置。

3. 优化措施以应对挑战：直接使用CUDA VMM API面临两个效率挑战，Vattention通过针对LLM服务负载特性的优化予以解决：

   • 隐藏内存分配延迟：
     • 计算与分配重叠（解码阶段）：利用解码阶段每个请求每迭代只生成一个令牌、内存需求可预测的特点，在上一轮迭代的计算过程中，通过后台线程预先为下一轮迭代可能需要的物理页面执行映射操作，从而将CUDA API调用（如cumemmap）的延迟（可能达数毫秒）隐藏起来。
     • 延迟回收与预分配（预填充阶段）：当一个请求完成时，不立即回收其物理页面。当新请求到来时，直接复用已完成的请求ID及其已映射的物理内存。如果新请求需要更多内存，再额外分配。同时，Vattention还会主动为某个空闲的reqid预分配少量页面，以备新请求使用。这些策略极大地减少了在预填充关键路径上进行内存分配的需要。
   • 减少内部碎片：CUDA VMM API原生仅支持分配2MB的大页面（Large Page），这可能导致严重的内部碎片（Internal Fragmentation）。为了支持更细粒度的分配（如与PagedAttention建议的16-32令牌块大小相当），作者修改了开源的NVIDIA统一虚拟内存驱动程序，新增了一套支持64KB、128KB、256KB页组大小的API（如表3中前缀为‘v’的API所示）。评估表明，使用更小的64KB页面并未导致TLB颠簸，对注意力内核性能无负面影响，同时显著降低了碎片。

4. 无缝集成与可移植性：Vattention被实现为一个Python库，通过C++扩展与CUDA驱动交互。它向服务框架（论文中以vLLM为例）暴露一组简单的API（如表4所示），用于初始化、请求ID管理以及确保内存映射的step调用。由于KV Cache在虚拟内存中保持连续，因此任何现有的、假设连续内存布局的非分页注意力内核（如FlashAttention-2、FlashAttention-3、FlashInfer的非分页版本）都可以无需修改，直接与Vattention集成，由Vattention透明地提供动态内存管理支持。

研究的主要结果通过在三款模型（Yi-6B, LLaMA-3-8B, Yi-34B）和A100/H100 GPU上的实验得到验证：

1. 预填充（Prefill）性能：在长上下文场景下，Vattention支持的非分页内核显著优于对应的PagedAttention内核。例如，对于Yi-34B模型，在192K上下文长度时，使用FlashAttention-2非分页内核的Vattention（fa2_vattention）比使用其PagedAttention版本（fa2_paged）的预填充吞吐量高1.24倍。性能提升主要源于非分页内核更快的计算速度，避免了PagedAttention内核中额外的块表查找和分支开销。对于FlashInfer，Vattention带来的优势更明显，部分原因是还避免了其分页版本中构建压缩块表等CPU开销。

2. 解码（Decode）性能：在解码阶段，Vattention（fa2_vattention）与性能最佳的PagedAttention方案（fa2_paged）持平，但两者均显著优于vLLM原生的解码内核（最高达1.99倍）和FlashInfer的PagedAttention解码内核（最高达1.23倍）。这突显了vLLm内核由于复杂性和维护滞后导致的性能差距，也说明了Vattention在利用最新优化内核方面的优势。

3. 端到端服务吞吐量：

   • 离线场景：在arXiv摘要生成的长上下文工作负载上，Vattention相比fa2_paged和fi_paged，整体吞吐量提升最高分别达1.18倍和1.23倍。
   • 在线场景：在模拟在线服务的动态负载下，Vattention能够显著降低请求的端到端执行延迟（中位数降低最高达42%），主要得益于其更快的预填充处理能力减少了排队延迟。

4. 可移植性优势的例证：新发布的、针对Hopper架构优化的FlashAttention-3内核在发布时尚未支持PagedAttention。Vattention能够开箱即用地支持FA3，无需对其内核做任何修改。实验表明，在H100 GPU上，FA3配合Vattention相比FA2配合Vattention还能带来额外的最高1.35倍的吞吐量提升，充分展示了Vattention在加速新内核部署方面的价值。

5. 消融实验（Ablation Studies）：

   • 验证了“计算与分配重叠”优化能有效消除解码迭代中的内存分配延迟尖峰。
   • 验证了“延迟回收”策略几乎完全消除了预填充阶段同步调用CUDA API的开销。
   • 验证了使用64KB小页面不会对注意力内核性能产生负面影响，并且在追求最大批处理规模的吞吐量导向场景中，能比2MB页面支持更大的批处理量（如图15所示）。

本研究的结论是，Vattention提出了一种更简单、可移植性更强且性能更优的动态KV Cache内存管理替代方案。它通过解耦虚拟与物理内存分配，在保留KV Cache虚拟连续性的同时实现了高效的物理内存按需使用，从而避免了PagedAttention方法带来的内核重写负担、系统冗余和性能开销。Vattention使得服务系统能够无缝、即时地集成任何现有的或未来的非分页注意力内核优化，保持了与前沿研究的同步。

本研究的亮点在于：第一，提出了一个根本性的设计范式转变，从在用户空间模拟分页（PagedAttention）回归到利用系统原生支持的按需分页（Vattention），更加符合操作系统内存管理的设计哲学。第二，实现了卓越的可移植性，这是其最突出的实践价值之一，极大地降低了集成新内核的工程复杂度。第三，针对LLM服务负载特性进行了一系列精妙优化，如重叠计算与分配、延迟回收等，有效隐藏了系统调用的开销。第四，通过修改开源驱动支持小页面，在保证性能的同时进一步减少了内存碎片，展示了解决实际问题时的工程深度。第五，提供了详实的实验证据，不仅对比了吞吐量，还深入分析了延迟开销、碎片影响和可移植性案例，全面论证了Vattention的优势。这项工作为LLM推理服务系统的内存管理提供了新的、更具原则性的解决方案，对工业界和学术界均有重要的参考价值。