vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

发表时间: 2025-03 · arXiv:2405.04437 (ASPLOS 2025)

Ramya Prabhu (Microsoft Research Bengaluru, India), Ajay Nayak (Indian Institute of Science Bengaluru, India), Jayashree Mohan, Ramchandran Ramjee, Ashish Panwar (Microsoft Research Bengaluru, India)

速读

一句话结论
vAttention 通过解耦虚拟内存和物理内存的分配，在保持 KV cache 虚拟内存连续性的同时实现了物理内存的动态按需分配，在无需重写注意力算子的前提下，将大语言模型推理吞吐量提升了最高 1.23 倍。

要解决什么问题
在大语言模型推理中，为了避免预分配显存带来的严重内部碎片并提升批处理大小，业界广泛采用 PagedAttention 机制进行物理显存的按需动态分配。然而，PagedAttention 的原有做法卡在一个关键机制上：它在实现物理内存动态分配的同时，破坏了 KV cache 在虚拟内存上的连续性。这种非连续的虚拟内存布局带来了三个致命卡点。首先是算子重写成本极高，开发者必须修改底层的注意力算子来支持非连续内存的寻址，这导致生产系统很难跟上学术界最新算子（如 FlashAttention-3）的迭代速度。其次是 GPU 运行时代价，为了在算子内部查找块表（Block-Table）并处理额外的代码分支，指令数量增加了，寄存器溢出压力变大，导致分页算子的运行速度比非分页算子慢最高 42%。最后是 CPU 运行时代价，推理框架被迫在用户态实现一套内存管理器来拼接这些虚拟内存块，这不仅与操作系统的虚拟地址翻译功能冗余，还在调度关键路径上引入了显著的延迟。

怎么做的
vAttention 的核心思路是利用 CUDA 虚拟内存管理（VMM）API，对虚拟内存和物理内存采用不同的分配策略：在虚拟内存中提前为 KV cache 预留巨大的连续空间，但将物理内存的实际分配推迟到运行时按需进行。因为现代 64 位系统的用户态虚拟内存空间极大（例如多卡环境下可达 256TB），所以提前预留不会造成任何物理资源浪费。在这一设计下，单个请求在单层网络中的最大 KV cache 容量定义为：

$$S = L \times H \times D \times P$$

其中 $L$ 为模型支持的最大上下文长度，$H$ 为 KV 头数，$D$ 为头部维度，$P$ 为数据精度对应的字节数。系统会按最大批处理大小 $B$ 预留总大小为 $B \times S$ 的连续虚拟内存。每个请求会被分配一个唯一的标识符 $\text{reqId}$，其在连续虚拟张量中的内存偏移量即为 $\text{reqId} \times S$。
为了让这套机制高效运转，vAttention 包含了三个关键设计。第一是异步显存分配，由于调用 CUDA API 映射物理内存需要陷入操作系统内核，延迟较高，vAttention 在解码阶段利用后台线程，将下一次迭代所需的物理内存分配与当前迭代的 GPU 计算完全重叠，从而隐藏了分配延迟。第二是延迟回收与预分配，在预填充阶段，系统不会立即释放已完成请求的物理页，而是让新请求直接复用这些已映射好物理内存的虚拟张量，避免了关键路径上的分配开销。第三是细粒度页表支持，标准的 CUDA API 仅支持 2MB 的大页分配，这会引发严重的内部碎片；作者通过修改开源的 NVIDIA 统一内存驱动，增加了对 64KB、128KB 和 256KB 物理页的支持，在不引发 TLB 抖动的前提下实现了细粒度的内存管理。

效果如何
实验在 1 至 2 张 A100 GPU 以及 1 至 2 张 H100 GPU 上进行，测试了 Yi-6B、Llama-3-8B 和 Yi-34B 模型。对比基线包括 vLLM 默认的解码算子，以及代表当前最优性能路线的 FlashAttention-2 和 FlashInfer 的分页（Paged）版本。在解码吞吐量上，vAttention 配合非分页的 FlashAttention-2 算子，比 vLLM 默认实现高出最高 1.99 倍。在处理 arXiv-Summarization 长上下文数据集的端到端在线与离线推理场景下，vAttention 的吞吐量比基于 PagedAttention 的 FlashAttention-2 和 FlashInfer 分别提升了最高 1.18 倍和 1.23 倍。最能体现其架构优势的是，vAttention 能够零代码修改、开箱即用地支持专为 Hopper 架构优化且发布时不支持分页的 FlashAttention-3 算子，使其吞吐量比分页版 FlashAttention-2 进一步提升 1.26 至 1.5 倍。该方法的局限在于，若要彻底消除内部碎片，必须修改底层的 NVIDIA 驱动以支持 64KB 小页，否则只能依赖 2MB 大页（尽管可以通过张量切片技术缓解碎片问题）；此外，当物理显存真正耗尽时，它仍需要像现有系统一样抢占并挂起部分请求。

主要贡献

大语言模型（LLMs）的推理优化至关重要，其中批处理（Batching）是提升吞吐量的核心技术。然而，实现大批处理量需要仔细分配GPU内存，以存储自回归生成过程中不断增长的KV cache。早期的静态内存分配会导致严重的内部碎片。vLLM提出的PagedAttention通过按需分配小块物理内存有效缓解了碎片问题，成为了当前LLM服务系统的行业标准。

但是，PagedAttention在尝试运行时分配物理内存时，改变了KV cache在虚拟内存中的连续性布局（从连续变为非连续）。这种设计带来了显著的编程复杂性和性能开销。本文旨在解决这一核心问题，其研究目标是在保留KV cache虚拟内存连续性的同时，缓解物理内存的碎片化。

本文的主要创新点如下：
1. 提出了vAttention，这是一种通过CUDA虚拟内存管理（VMM）API将虚拟内存和物理内存的分配进行解耦的内存管理方法。它在虚拟内存中保持KV cache的连续性，同时支持物理内存的动态分配。
2. 引入了多种针对LLM的特定优化，以解决CUDA虚拟内存支持的局限性，例如将内存分配与计算重叠、机会性地提前分配页面，以及推迟内存回收，从而隐藏了按需内存分配的延迟。
3. 修改了开源的CUDA统一虚拟内存驱动，增加了对更小的64KB页面的支持，以解决CUDA仅支持2MB大页导致的内部碎片问题。
4. vAttention提供了一个更简单、可移植且高性能的替代方案。它开箱即用地支持各种未修改的注意力算子（如FlashAttention-2、FlashInfer和FlashAttention-3），与基于PagedAttention的内核相比，LLM服务吞吐量最高提升了$1.23\times$。

系统/库与PagedAttention相关的问题
vLLM：开创了PagedAttention。尽管处于积极维护的代码库中，vLLM的PagedAttention内核比FlashAttention-2慢高达$2.8\times$。此外，改变块大小会使内核执行时间改变高达$1.9\times$。
FlashAttention-2：基于PagedAttention的Prefill内核比非分页内核慢高达37%，Decode内核慢高达12%。最初添加分页支持的尝试未能通过单元测试。
FlashAttention-3 / cuDNN-9中的SDPA：发布时，针对NVIDIA Hopper架构的最先进注意力内核不支持PagedAttention。
TensorRT-LLM：在Python前端中，服务吞吐量下降超过10%。建议使用C++前端。即使使用C++，在某些情况下PagedAttention的延迟也高出高达5%。
FlashInfer：基于PagedAttention的Prefill内核比非分页内核慢高达42%。

表 1. PagedAttention方法要求应用程序显式管理动态分配的物理内存，包括重写注意力内核。这些例子突出了与此方法相关的复杂性、性能和维护挑战。

背景知识与设计原则

大语言模型与KV Cache。LLM推理包含并行处理提示词的Prefill阶段和逐个生成输出token的Decode阶段。模型计算给定序列的查询（query）、键（key）和值（value）向量：
$q_i = W_q x_i, \quad k_i = W_k x_i, \quad v_i = W_v x_i$
生成的$k_i$和$v_i$被追加到先前token的向量中，产生两个矩阵$K, V \in \mathbb{R}^{L' \times (H \times D)}$，其中$L'$是当前请求的上下文长度，$H$是KV头数，$D$是每个头的维度。注意力计算如下：
$Attention(q_i, K, V) = softmax(\frac{q_i K^T}{scale}) V$
推理引擎必须将$k_i$和$v_i$存储在内存中以供跨迭代重用，这被称为KV cache。

PagedAttention的缺陷。PagedAttention在用户空间实现了请求分页，这带来了几个问题：
1. 需要重写注意力内核。传统的注意力算子假设输入张量$K$和$V$在内存中是连续的。PagedAttention要求修改内核以处理非连续的块，这阻碍了系统快速跟进社区对注意力算子的最新优化（如FlashAttention-2）。
2. 在服务框架中增加冗余。服务系统需要跟踪KV cache块的虚拟内存地址并在运行时传递给注意力内核，这实际上是在重复操作系统用于虚拟到物理地址转换的工作。
3. 性能开销。在GPU上，查找Block-Tables和执行额外分支增加了$7\%-13\%$的指令量，导致显著的性能下降。在CPU上，准备Block-Tables的过程在Python运行时中也会带来高达$10\%$的延迟。

图 1. PagedAttention涉及两层内存管理：一层在用户空间，一层在OS内核空间。
图 2. PagedAttention在Prefill内核中的开销。
图 3. vLLM的分页Decode内核延迟对块大小很敏感。

LLM服务系统的关键观察。
1. KV cache内存需求在每次迭代中是可预测的。在Decode阶段，KV cache大小每次迭代均匀增加一个token。这允许系统提前确定是否需要额外内存。
2. KV cache不需要高内存分配带宽。单个token的内存占用通常只有几十到几百KB。即使在大批处理量下，内存分配速率最高也不超过750MB/秒，带宽需求在达到一定批处理大小后会饱和。

方法细节

vAttention采用分离的虚拟与物理内存分配策略。vAttention提前在虚拟内存中为KV cache分配一个巨大的连续缓冲区，同时将物理内存的分配推迟到运行时。这种设计在不产生物理内存碎片的情况下，保留了KV cache的虚拟连续性。由于现代64位系统为每个进程提供了128TB的用户可寻址虚拟内存，因此虚拟内存的碎片化和浪费不是问题。

预留虚拟内存。由于虚拟内存充足，vAttention预先分配足够大的尺寸来容纳需要支持的最大批处理大小（可配置）的KV cache。在执行此操作时，系统假设每个请求的上下文长度与模型支持的最大长度相同。

虚拟内存缓冲区的数量。服务框架为模型的每一层维护独立的$K$和$V$张量。因此，vAttention在每个工作节点（worker）上预留$2 \times N$个缓冲区，其中$N$是该工作节点管理的层数。

虚拟内存缓冲区的大小。每个缓冲区的最大大小为$BS = B \times S$，其中$B$是最大批处理大小，$S$是单个请求的每层$K$ cache（或$V$ cache）在工作节点上的最大大小。进一步地，$S = L \times H \times D \times P$，其中$L$是模型支持的最大上下文长度，$H$是工作节点上的KV头数，$D$是每个KV头的维度，$P$是基于模型精度的字节数（例如，对于FP16/BF16，$P = 2$）。例如，对于Yi-34B（FP16，两路张量并行TP-2），$N = 60, H = 4, D = 128, P = 2$，最大支持上下文长度$L = 200K$。对于此配置，$S = 200MB$。假设$B = 500$，每个工作节点的每个缓冲区最大大小为$100GB$。因此，60层所需的总虚拟内存为120个100GB的缓冲区（总计12TB），这完全在256TB的虚拟地址空间限制内。

图 5. vAttention中单个K cache（或V cache）张量的动态内存管理过程。

利用CUDA虚拟内存API解耦分配。标准的GPU内存分配接口cudaMalloc不支持按需分页，即它同时分配虚拟内存和物理内存。vAttention利用CUDA提供的低级VMM API【12，CUDA Toolkit Documentation: Virtual Memory Management，2024，URL】来实现虚拟和物理内存分配的解耦。分配的粒度取决于GPU使用的页面大小，虚拟内存缓冲区或物理内存句柄的大小必须是物理内存分配粒度的倍数。物理内存页面可以独立于其他子区域分配给（或从其取消分配）虚拟内存缓冲区中的子区域。

扩展PyTorch缓存分配器。KV cache是张量的集合。在当前的深度学习框架（如PyTorch）中，通过torch.empty等API分配的张量带有预先分配的物理内存，因为PyTorch缓存分配器依赖于cudaMalloc接口。vAttention依靠CUDA的低级API支持，扩展了PyTorch缓存分配器，允许应用程序为张量预留虚拟内存缓冲区，而无需提前提交物理内存。通过这些API分配的张量被称为虚拟张量（virtual tensors）。

图 6. LLM服务系统通过一组简单的API与vAttention交互进行KV cache管理。

请求级KV cache索引定位。一个虚拟张量代表最大批处理大小$B$下的某一层的$K$ cache（或$V$ cache）。在这些张量中，不同的请求占据不同的非重叠子区域（子张量）。vAttention使用一个在$0$到$B-1$范围内的唯一整数标识符reqId来定位请求的子张量。请求的子张量在整个批次的虚拟张量中的偏移量为$reqId \times S$，其中$S$是每个请求每层$K$ cache的最大大小。reqId由vAttention负责分配。

vAttention的初始设置与参数配置。vAttention被构建为一个Python库，内部使用CUDA/C++扩展与CUDA驱动程序交互。当服务框架启动时，每个模型工作节点加载vAttention库，并通过init API配置模型参数$N, H, D, P, B$以及首选的页面组大小。在内部，vAttention在该工作节点上预留$2 \times N$个虚拟张量。这些虚拟张量在服务应用程序的生命周期内被预留。此外，vAttention在初始化期间还在每个工作节点预先分配物理内存页面，但此时这些页面尚未映射到KV cache中。

调度新请求与分配标识符。当第一次调度一个新请求时，服务框架通过alloc_reqid API从vAttention获取一个新的reqId。该请求的所有后续内存管理操作都标记有此reqId。

在模型执行期间动态映射物理内存。在分发批次执行之前，框架需要确保每个活跃请求的KV cache子张量都有物理内存支持。为此，在将迭代的第一个内核分发到GPU之前，框架调用step API，指定每个请求当前的上下文长度。vAttention内部确保为每个活跃的reqId映射足够的物理页面，然后再将执行权交还给框架。如果vAttention无法满足内存需求，它将返回失败，框架可以抢占一个或多个请求以允许向前推进。

Prefill与Decode阶段的差异化内存映射逻辑。根据请求处于Prefill阶段还是Decode阶段，给定迭代可能需要映射不同数量的物理内存。Prefill阶段并行处理给定提示词的输入token，因此需要映射的物理内存量取决于正在调度的提示词token数量。如果模型某一层的全部提示词token的$K$ cache总大小为$s$，且页面组大小为$t$，则每个工作节点需要确保在给定reqId的$2 \times N$个KV cache子张量中至少映射$(s + t - 1) / t$个页面组。对于处于Decode阶段的请求，由于每次迭代仅产生一个输出token，因此每个虚拟张量最多需要一个新页面组。vAttention内部跟踪为每个请求映射的页面组数量，并仅在先前的页面组即将耗尽时才映射新页面组。

请求完成与连续批处理的支持。当请求达到用户指定或模型支持的最大上下文长度，或模型生成结束token时，请求终止。框架通过free_reqid API通知vAttention请求完成。内部，vAttention可能取消映射已完成请求的物理页面，或将其推迟到稍后释放。为了支持连续批处理（当批处理中间的请求退出时，会在KV cache的虚拟张量中产生未使用的空洞），vAttention利用了FlashAttention提供的丰富API支持（cache_batch_idx参数），允许Q和KV cache在批次维度上具有不同的大小并按任意顺序排列。当批次的请求组成发生变化时，vAttention更新运行中请求的cache_batch_idx，使得它们的Q张量根据其reqId映射到各自的KV cache。

# Algorithm 1 Using vAttention in a serving framework.
1: max_batch_size <- B
2: cache_seq_len <- [0]*B
3: req_batch_idx <- dict()
4: vattention.init(config_params)
5: while !request_pool.is_empty() do
6:     for Ri in new_requests do
7:         if can_schedule(Ri) then
8:             idx <- vattention.alloc_reqid()
9:             req_batch_idx[Ri] <- idx
10:            cache_seq_len[idx] <- prompt_len(Ri)
11:        end if
12:    end for
13:    vattention.step(cache_seq_len)
14:    model.forward()
15:    for Ri in active_requests do
16:        idx <- req_batch_idx[Ri]
17:        if is_complete(Ri) then
18:            cache_seq_len[idx] <- 0
19:            vattention.free_reqid(idx)
20:        else
21:            cache_seq_len[idx] += 1
22:        end if
23:    end for 
24: end while

隐藏内存分配的延迟：重叠Decode阶段的分配与计算。由于每次调用CUDA VMM API（如cuMemMap + cuMemSetAccess）需要约40微秒，为具有60层的Yi-34B增加一个请求的KV cache将带来约5毫秒的延迟。为了隐藏这一延迟，vAttention利用了内存需求的可预测性来将内存分配与计算重叠。在Decode阶段，请求每次迭代仅需要最多一个新页面组。vAttention跟踪当前上下文长度和已映射的物理内存量。当框架在迭代$i-1$中调用step API时，vAttention判定请求在迭代$i$中是否需要更多内存，并启动一个后台线程在迭代$i-1$执行期间为迭代$i$映射页面组。由于单次迭代延迟通常为几十到几百毫秒，后台线程有足够的时间在迭代开始前准备好物理内存映射。

隐藏内存分配的延迟：Prefill阶段的延迟回收与急切分配。为了避免在Prefill阶段从头分配物理内存，vAttention采用了延迟回收机制。如果请求R1在迭代$i$完成，而新请求R2在迭代$i+1$加入，vAttention简单地推迟回收R1的页面组，并将R1的reqId分配给R2。这样，R2可以直接重用已由物理页面支持的KV cache张量，只有当R2的上下文长度大于R1时才需要新的分配。此外，vAttention通过提前主动分配少量页面组来优化Prefill阶段（急切分配）。它试图在一个非活跃的reqId的虚拟张量中保持一定数量的页面组被映射。当新请求到达时分配该reqId，同时识别下一个要分配的reqId并急切地为其映射物理页面组。仅当vAttention中缓存的页面组数量降至特定阈值（如GPU内存的$10\%$）以下时，才触发内存回收。

缓解内部碎片：修改NVIDIA驱动以支持更小的页面组。NVIDIA GPU原生支持至少三种页面大小：4KB、64KB和2MB。然而，现有的CUDA VMM API（如cuMemCreate）仅支持以2MB大页的倍数分配内存，这会导致严重的内部碎片。由于CUDA VMM API在闭源的NVIDIA驱动程序中实现，vAttention转而在开源的NVIDIA统一内存驱动程序中实现了一组新的API，以模仿现有CUDA API的功能，但增加了对多页面大小的支持。新API（前缀为v）支持以64KB、128KB和256KB大小的页面组分配物理内存。例如，vMemMap结合了cuMemMap和cuMemSetAccess的功能，而vMemRelease结合了cuMemUnmap和cuMemRelease的功能。服务框架可以在初始化vAttention时配置所需的页面组大小。

实验环境

数据集：ShareGPT（用于平均解码长度分析，415 tokens）、arXiv-Summarization（用于长上下文离线/在线端到端评估，上下文长度64K-192K）、OpenChat（用于动态负载批处理大小测试）。
模型：Yi-6B（32层，4 KV头）、Llama-3-8B（32层，8 KV头）、Yi-34B（60层，8 KV头）。
硬件配置：
- Yi-6B：1张 NVIDIA A100 GPU (80GB)。
- Llama-3-8B 和 Yi-34B：2张通过NVLink连接的 NVIDIA A100 GPUs (每张80GB)，采用张量并行（TP-2）。
- FlashAttention-3 移植性实验：1-2张 NVIDIA H100 GPUs。
软件配置：以 vLLM v0.2.7 为通用服务框架。集成了 FlashAttention-2 v2.5.9 和 FlashInfer v0.4.0 作为注意力后端。底层使用CUDA 12.1、Python 3.10、PyTorch 2.3.0及定制的VMM APIs。

实验结果

1. Prefill阶段吞吐量评估 (7.1)
* 实验内容：比较四种配置在不同上下文长度下的Prefill吞吐量（Token/秒）：FA2_Paged, FI_Paged, FA2_vAttention, FI_vAttention。
* 实验结果：对于短上下文，FA2的Paged和vAttention吞吐量几乎相同，但由于避免了Block-Table的开销，FI_vAttention优于FI_Paged。对于长上下文（>16K），vAttention支持的内核显著优于分页版本。例如在192K上下文下，FA2_vAttention比FA2_Paged在Yi-6B、Llama-3-8B和Yi-34B上分别提升了$1.24\times$、$1.26\times$和$1.24\times$。
* 分析结论：vAttention在Prefill阶段的提升主要归功于虚拟连续KV cache带来了更快的注意力内核计算，因为长提示词的Prefill阶段主要被注意力计算所主导。
* 图表引用：Fig 7证实了vAttention在长上下文Prefill中的优势。
图 7. Prefill吞吐量。vAttention支持的系统优于FlashAttention-2和FlashInfer的分页版本。由于Prefill注意力的二次复杂度，较长上下文的吞吐量较低。

2. Decode阶段吞吐量评估 (7.2)
* 实验内容：在16K初始上下文长度下，比较vLLM默认解码内核、FI_Paged、FA2_Paged与FA2_vAttention随着批处理大小（1到32）变化的Decode吞吐量。
* 实验结果：FA2_vAttention的性能与表现最好的分页内核FA2_Paged持平，并显著优于FI_Paged和vLLM。与vLLM相比，FA2_vAttention在Yi-6B、Llama-3-8B和Yi-34B上分别最高提升了$1.99\times$、$1.58\times$和$1.53\times$。与FI_Paged相比，最高提升了$1.23\times$（Yi-6B，Batch 12）。
* 分析结论：vLLM的解码内核由于缺乏最新优化导致延迟极高。FA2_vAttention与FA2_Paged表现相似的原因是Decode注意力受限于内存带宽，内存停顿掩盖了分页支持所需的额外计算开销。
* 图表引用：Fig 8展示了不同批处理大小下的Decode吞吐量对比。
图 8. Decode吞吐量。FA2_vAttention与FA2_Paged持平（注意重叠的线），后者是所有基于PagedAttention的替代方案中最好的，同时优于FI_Paged和vLLM。

3. 端到端性能：离线场景评估 (7.3)
* 实验内容：处理arXiv-Summarization数据集中427个长上下文请求（64K-192K tokens），测量每分钟完成的请求数。
* 实验结果：FA2_vAttention分别比FA2_Paged提升了$1.18\times$ (Yi-6B)、$1.15\times$ (Llama-3-8B) 和 $1.13\times$ (Yi-34B)。比FI_Paged提升了$1.19\times$、$1.23\times$和$1.14\times$。
* 分析结论：vAttention的性能增益与上下文长度及Prefill:Decode token比例成正比。工作负载越受Prefill限制，vAttention提供的收益越高。
* 图表引用：Fig 9证实了离线场景下的吞吐量提升。
图 9. 处理arXiv-Summarization数据集时长上下文请求的离线推理吞吐量。

4. 端到端性能：在线场景评估 (7.4)
* 实验内容：在不同输入负载（QPS，泊松分布）下，运行512个请求，测量端到端请求执行延迟的CDF（累积分布函数）。
* 实验结果：FA2_vAttention一致优于所有基线。与FA2_Paged相比，中位数请求执行延迟在Yi-6B (QPS 0.25) 降低了高达42%，Llama-3-8B (QPS 0.3) 降低了28%，Yi-34B (QPS 0.1) 降低了29%。
* 分析结论：主要原因是vAttention能更快地计算新请求的Prefill阶段，从而大幅减少了排队延迟。
* 图表引用：Fig 10展示了在线推理中端到端请求执行延迟的CDF。
图 10. 变化负载下在线推理中端到端请求执行延迟的CDF。

5. 移植性验证：FlashAttention-3 (7.5)
* 实验内容：在H100 GPU上测试不支持PagedAttention的最新FA3内核，验证vAttention的开箱即用移植性。
* 实验结果：带有vAttention的FA3相比FA2_vAttention进一步提供了高达$1.35\times$ (Yi-6B) 的加速，而FA2_vAttention已经比FA2_Paged快了$1.15\times$。
* 分析结论：vAttention无需更改代码即可部署最新的硬件优化内核，证明了其卓越的移植性优势。
* 图表引用：Fig 11展示了H100 GPU上的离线推理吞吐量。

6. 消融实验 (7.6)
* 隐藏分配延迟：通过重叠内存分配与计算，有效隐藏了CUDA VMM API同步分配带来的5-15毫秒的延迟尖峰（Fig 12）。
* 延迟回收：同步分配64KB页面会导致高达$1.15\times$的Prefill开销。延迟回收机制消除了调用CUDA VMM API的需要，确保Prefill延迟不受内存分配影响（Fig 13）。
* 页面大小的影响：在LLM推理中，使用较小的64KB页面不会导致TLB抖动，注意力内核的执行延迟与使用2MB大页时基本不受影响（Fig 14）。同时，较小的页面组能够支持更大的最大批处理大小（Fig 15）。
* 内存分配带宽：即使使用64KB页面，vAttention每秒每GPU也能分配高达7.6GB的物理内存，远超Decode阶段750MB/s的最大内存分配率。
图 12. 有无重叠内存分配与计算情况下的Decode迭代延迟。

图 14. 页面大小对FlashAttention-2 Prefill（左）和Decode（右）注意力内核运行时间的影响。
图 15. 动态工作负载下使用不同页面组大小获得的最大批处理大小。

补充细节

通过统一内存管理KV cache。作者也考虑过利用cudaMallocManaged提供的统一内存来实现动态内存分配。然而，当前的统一内存支持不适合服务LLM：首先，它不支持部分释放，阻碍了单个请求物理内存的回收；其次，它缺乏对内存别名（aliasing）的支持，阻止了物理内存中KV cache内容的去重（当请求共享公共前缀时去重很有用）；此外，它默认分配2MB页面会导致严重碎片。vAttention在NVIDIA驱动中的更改基于统一内存，但通过额外API启用了部分释放和页面共享，并支持更小的页面。

通过张量切片（Tensor Slicing）减少碎片。为了在不修改NVIDIA驱动程序的情况下减少2MB页面引起的碎片，vAttention提供了一种替代方法：使用单个2MB页面存储给定请求所有层的KV cache token。通过分配形状为$[B, L, N, H, D]$的单一虚拟张量并在所有层上对其进行切片，碎片可以减少到先前设计的$1/N$。但这种方法导致每层的KV cache不再连续，需要注意力内核支持带步幅（strides）的内存寻址。由于早期版本的FlashInfer缺乏此类支持，为了支持未修改的内核，vAttention选择了在驱动中添加对更小页面的支持。这两种解决方案相互兼容。

编程工作量对比。PagedAttention需要巨大的编程工作量。例如，在vLLM中集成FlashInfer解码内核需要跨越15个文件修改600多行代码；在FlashAttention-2中实现初始分页支持也需要约280行代码更改。相比之下，vAttention使得在服务框架中只需几行代码更改即可用一个注意力内核替换另一个（Fig 16），内存管理将在底层透明地继续工作。

图 16. 当使用vAttention进行内存分配时，用FlashInfer替换FlashAttention-2的Prefill注意力内核所需代码更改的说明。

相关工作。近期的GMLake【45，GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching，2024，ASPLOS，URL】展示了使用CUDA虚拟内存支持可以缓解DNN训练作业中的碎片，增加训练批次大小。GMLake利用CUDA支持将多个较小的物理内存页面合并为单个虚拟连续对象。相比之下，vAttention专注于优化推理工作负载。

结论

本文提出了vAttention，用于LLM服务系统中的动态KV cache内存管理。vAttention的核心亮点在于，它在缓解物理内存碎片的同时，保留了KV cache在虚拟内存中的连续性。与流行的PagedAttention方法相比，vAttention不仅显著减轻了编程负担，还提高了系统的可移植性和整体推理性能，为未来集成和部署新的硬件优化内核铺平了道路。

附录

Artifact Appendix。vAttention的代码已开源，运行环境依赖于CUDA 12.1、Python 3.10和PyTorch 2.3.0。硬件要求包括两张NVLink连接的NVIDIA A100 GPU (80GB)。通过克隆GitHub仓库并使用Conda或提供的Docker镜像进行安装。附录中包含了复现论文关键结果（Figures 2, 3, 4, 6, 7, 8, 9, 11）的详细脚本命令说明。由于端到端实验（Fig 8, 9）可能需要超过24小时，脚本默认配置了较小的请求数以供快速验证。

引用汇总

以下是本文在背景、方法及相关讨论中引用的核心文献汇总及原文描述上下文：

【12，CUDA Toolkit Documentation: Virtual Memory Management，2024，URL】：在5.2节方法细节中引用，原文描述vAttention利用这些低级API解耦虚拟内存和物理内存的分配。
【45，GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching，2024，ASPLOS，URL】：在第9节相关工作中引用，原文描述GMLake使用CUDA虚拟内存支持将小页面合并为连续对象以缓解DNN训练中的碎片，而vAttention针对的是推理。