vAttention 完全拆解：用 CUDA 虚拟内存保住连续 KV Cache

为什么 KV Cache 的地址布局会决定推理系统的形态

vAttention 由 Microsoft Research India、印度科学理工学院等机构的研究者于 2024 年 5 月首次发布。它追问的不是“如何再写一个更快的分页注意力 Kernel”，而是一个更靠近操作系统的问题：既然 GPU 已经有页表和虚拟内存，推理框架为什么还要自己维护一套块表，再要求注意力 Kernel 理解这套块表？

自回归推理包含 Prefill 和 Decode 两个阶段。Prefill 并行处理输入 token，Decode 每轮只为每个请求生成一个新 token。第 $i$ 个 token 的 Key、Value 会被后续所有 token 重用，因此每层都要维护随上下文增长的 KV Cache。若一个工作进程负责 $N$ 层，每层有 $H$ 个本地 KV 头，头维度为 $D$，每个元素占 $P$ 字节，请求当前长度为 $L$，其 KV Cache 大小近似为：

系数 2 对应 K 与 V。对于固定模型，$N$、$H$、$D$、$P$ 都不变，因此显存需求与 $L$ 线性增长。困难在于系统接收请求时通常不知道最终的 $L$：它可能很快遇到结束 token，也可能一直生成到上下文上限。

最直接的做法是按最大长度为每个请求预留连续显存。这样普通张量 Kernel 可以通过“基地址 + 线性偏移”访问 KV，却会产生严重的内部碎片。假设上限是 128K，而请求只走到 4K，那么绝大部分已经承诺的物理显存从未写入；更隐蔽的浪费发生在请求生命周期内，即便最终长度可以预知，增长过程中的尾部空间仍长期闲置。显存容纳的并发请求变少，连续批处理也就失去扩大 batch 的空间。

PagedAttention 的解决办法是把每个请求的 KV 拆成固定 token 数的小块，只在需要时分配物理块。它成功解决了容量浪费，却同时改变了 Kernel 看到的布局：一个请求的逻辑序列被散落在多个不连续的虚拟地址中。推理框架要维护“逻辑块号到物理块号”的 Block Table，Kernel 先查 Block Table，再由 GPU 页表把查到的虚拟地址翻译成物理地址。

图 1 来自 vAttention 论文。左侧是框架和分页 Kernel 显式完成的第一次寻址：根据 token 所在逻辑块查 Block Table。右侧是硬件本来就要完成的第二次寻址：GPU MMU 查询页表。图中真正值得关注的不是“多查了一张表”这么简单，而是它把分页语义扩散到了整个软件栈：

• 调度器必须分配、引用、复制和回收 KV 块；
• CPU 每轮要为动态 batch 整理 Block Table 元数据；
• 注意力 Kernel 必须接受块表、块大小等额外参数，并处理跨块边界；
• 新 Kernel 即使已支持连续 KV，也不能直接获得动态内存管理能力。

这不意味着 PagedAttention 的设计错误。它以用户态软件可控的方式解决了当时最紧迫的显存碎片问题，还自然支持前缀块共享、写时复制等能力。vAttention 的贡献是指出另一种分层边界：让框架只描述“每个请求现在需要多长的连续张量”，把虚拟地址到物理页的映射交还给系统层。

vAttention 的核心抽象：虚拟连续，物理按需

vAttention 的一句话定义是：先为所有可能的 KV Cache 保留足够大的连续虚拟地址区间，但不立即为整个区间承诺物理显存；请求增长时，才把空闲物理页映射到即将访问的那一段虚拟地址。

这里必须区分三个常被混用的动作：

1. 保留虚拟地址只是在进程地址空间中圈出一段号码范围，不等于消耗同等大小的 HBM。
2. 创建物理内存句柄才真正从 GPU 取得物理存储，但句柄尚未必能被 Kernel 访问。
3. 映射把物理内存挂到保留的虚拟地址上，并设置访问权限；此后普通 CUDA Kernel 才能用指针读写。

因此，“虚拟张量有 12 TB”与“GPU 装有 12 TB 显存”并不矛盾。论文以 Yi-34B、FP16、TP=2 为例：每个工作进程有 60 层、4 个本地 KV 头、头维 128、最大长度 200K。单层单个 K 或 V 的请求槽位上限为：

若最大 batch 为 $B=500$，一个层级 K 张量的虚拟区间为 $BS=100$ GB。K、V 各一份且共有 60 层，总虚拟地址预留约 12 TB，但任一时刻映射的物理显存只覆盖活跃请求实际走到的页。64 位进程通常有远大于 HBM 容量的用户虚拟地址空间，这正是方案能够成立的容量差。

对普通注意力 Kernel 而言，这仍然是形如 $[B,L,H,D]$ 的连续张量。请求槽位 reqId 的起始偏移可写成：

请求内部第 $t$ 个 token 的 K 地址继续是固定步长的线性索引。Kernel 不知道底层哪些页何时被映射，也不需要 Block Table。只要系统保证 Kernel 启动前，其将访问的地址范围已有物理页支撑，连续张量 Kernel 就能原样运行。

图 2 展示一个 K 张量中的两个请求槽位。阶段 (a) 只有虚拟地址，没有物理页；(b) 为 R1 映射第一页；(c) R1 扩到两页，R2 获得一页；(d) R1 结束后，页面暂不回收；(e) R3 复用 R1 的 reqId 和已经映射的页。连续性发生在虚拟地址层，按需性发生在物理地址层，两者不再互相排斥。

这也解释了 vAttention 与 Unified Memory 缺页迁移的差别。vAttention 不等待 Kernel 访问未映射地址后触发 fault，而是在已知下一轮上下文长度时主动完成映射。推理循环本来就知道本轮会写入多少 token，因此没有必要把正确性押在运行中的缺页异常和迁移策略上。

CUDA 虚拟内存如何支撑动态 KV Cache

CUDA 从 10.2 起提供底层虚拟内存管理接口。vAttention 使用的核心能力可以按以下顺序理解：

1
// 1. 查询物理分配的最小/推荐粒度
2
cuMemGetAllocationGranularity(&granularity, &prop, flags);
3

4
// 2. 仅保留连续虚拟地址，不分配同等大小物理显存
5
cuMemAddressReserve(&base, virtual_size, alignment, 0, 0);
6

7
// 3. 创建一段物理显存，得到与地址无关的 handle
8
cuMemCreate(&handle, page_size, &prop, 0);
9

10
// 4. 把 handle 映射到 base + offset
11
cuMemMap(base + offset, page_size, 0, handle, 0);
12

13
// 5. 为目标 GPU 设置读写权限
14
cuMemSetAccess(base + offset, page_size, &access, 1);

cuMemAddressReserve 类似操作系统中的地址空间保留；cuMemCreate 创建物理分配句柄；cuMemMap 建立映射；cuMemSetAccess 决定哪些设备可以访问。释放时则反向执行 cuMemUnmap、cuMemRelease 和 cuMemAddressFree。NVIDIA 的 CUDA Driver API 文档 特别指出，cuMemMap 本身不会自动让地址可访问，仍需正确设置访问权限。

从请求到页映射：完整运行流程

vAttention 向服务框架暴露的接口很少：初始化 KV 张量、分配请求槽位、释放槽位，以及根据当前序列长度推进映射。把一次请求按时间展开，可以看到内存层与调度层的边界。

为什么异步分配不会卡住 GPU

动态映射最危险的地方是尾延迟。论文给出的一个量级例子是：Yi-34B 有 60 层，扩展一个请求需要为 K、V 各层完成 120 次映射；若每次约 9 微秒，总计接近 1 毫秒。对追求稳定 TPOT 的在线服务，周期性出现 1 毫秒以上尖峰并不能忽略。

vAttention 利用了 Decode 的确定性：第 $i-1$ 轮开始前，系统已经知道第 $i$ 轮每个存活请求最多再写一个 token。于是 step_async 可以在后台线程为下一轮准备映射，同时 GPU 执行当前轮。其时间关系是：

1
CPU 主线程:  检查第 i-1 轮映射 ───── 发射第 i-1 轮 Kernel ───── 检查第 i 轮
2
后台线程:                    └── 为第 i 轮映射所需物理页 ──┘
3
GPU:                              执行第 i-1 轮计算

前台检查仍然不能省略。如果后台线程因调度抖动、驱动拥塞或突发 Prefill 没有按时完成，主线程必须同步补齐或拒绝本轮执行。异步路径优化常态延迟，同步检查守住内存安全。

图 3 来自论文消融实验，比较同步映射和与计算重叠两种执行方式。同步方案在跨页迭代出现明显尖峰；重叠后尖峰基本被当前轮 GPU 计算覆盖。图的意义并不是映射“免费”，而是每轮模型执行通常持续数到数十毫秒，而 KV 物理分配带宽需求相对低，后台线程有足够的隐藏窗口。

除了重叠，vAttention 还使用两种互补策略：

• 延迟回收：完成请求的页暂时留在槽位上，让后继请求直接复用，减少 unmap/map 往返。
• 提前分配：为空闲槽位预先映射一部分页，使突发新请求的 Prefill 不必在关键路径上从零准备。

两者本质上都是缓存。延迟回收提高命中率，但会暂时占用本可给其他槽位的物理页；提前分配降低入场延迟，但预测错误会形成闲置容量。因此系统设定低水位阈值：当空闲物理页不足，例如下降到 GPU 容量的一定比例时，后台线程才真正回收缓存页。这个策略把频繁的精确释放改成带迟滞的批量资源管理，类似 allocator 的缓存层。

连续 Kernel、分页 Kernel 与系统代价

把两种方案放在同一张表里，差异不在“是否分页”，而在分页机制由谁看见。

PagedAttention 常被描述为“像操作系统一样分页”，但从地址路径看，它更像在 CUDA 用户态实现一层面向 token block 的分段映射。这个抽象能表达前缀共享、非连续扩展和独立回收；代价是计算 Kernel 必须参与解释。

vAttention 则让 GPU MMU 处理最终的页翻译。Kernel 对 $K[t,h,d]$ 的线性索引不变，硬件页表把该虚拟地址导向任意物理页。这样可以直接使用 FlashAttention 或其他只接受连续 KV 的实现，新 Kernel 的 IO 分块、warp 调度和流水线优化不必再复制一份分页版本。

性能边界、适用条件与工程陷阱

论文在 Yi-6B、Llama-3-8B、Yi-34B 与 A100/H100 环境中分别测量 Prefill、Decode、在线和离线服务。理解结果时要先分清三种收益来源：

1. 内存容量收益：按需承诺物理页，避免按最大长度预留。
2. Kernel 收益：连续 KV 可以直接使用较快的非分页 Kernel。
3. CPU 收益：不再逐轮构造和搬运 Block Table。

第一项是 vAttention 与 PagedAttention 都追求的；后两项才是二者的主要差异。论文早期版本在长上下文微基准中报告，相比当时 vLLM 的 Decode Kernel，vAttention 最高达到 1.97 倍吞吐；但它与优化后的 FlashAttention 分页实现大致相当。这个对照说明：最高数字主要揭示特定基线 Kernel 的差距，不能归结为虚拟内存映射本身让算术更快。

在 Prefill 中，长上下文注意力的计算占比很高，连续版 FlashAttention/FlashInfer 避免分页路径，收益更明显。论文所测版本中，连续 FlashAttention 的 vAttention 配置相对分页版本在部分长上下文场景有显著优势；FlashInfer 的专用分页 Prefill 已经更强，差距较小。Prefill 结果也受 chunk 大小、GQA 头数和具体库版本影响。

图 4 来自论文较新修订版的离线 arXiv Summarization 轨迹实验。该轨迹的输入上下文很长，Prefill/Decode token 比较高，因而连续 Prefill Kernel 的优势会被放大。论文在三种模型上报告，FA2_vAttention 相对 FA2_Paged 的端到端吞吐提升约 1.13 至 1.18 倍，相对当时 FI_Paged 配置约 1.14 至 1.23 倍。图表用于说明趋势，部署决策仍需用当前框架、当前 Kernel 与真实流量重新测量。

如何把 vAttention 接入推理引擎

官方实现选择 Sarathi-Serve 作为宿主，但接入逻辑可以抽象为四个边界。已有文章 Sarathi-Serve 完全拆解 解释了分块 Prefill 与无停顿批处理；vAttention 在这里扮演的是 KV 内存后端，不改变调度器的 token 预算目标。

1
kv_tensors = vattention.init_kvcache(
2
    num_layers=local_num_layers,
3
    num_heads=local_num_kv_heads,
4
    head_size=head_dim,
5
    max_batch_size=max_running_requests,
6
    max_model_seq_len=max_context_len,
7
    dtype=dtype,
8
    page_size=page_size,
9
)
10
vattention.reserve_physical_pages(kv_budget_bytes)

1
req_id = vattention.alloc_new_batch_idx()
2
request_to_slot[request.id] = req_id
3

4
# 完成事件确认后
5
vattention.free_batch_idx(req_id)
6
del request_to_slot[request.id]

1
target_lengths = [0] * max_running_requests
2
for req in scheduled_requests:
3
    slot = request_to_slot[req.id]
4
    target_lengths[slot] = req.num_computed_tokens + req.num_scheduled_tokens
5

6
ok = vattention.step_async(target_lengths)
7
if not ok:
8
    preempt_and_retry()
9
run_model(kv_tensors, target_lengths)

小结

vAttention 的核心不是一种新的注意力公式，而是重新划分 KV Cache 内存管理的职责：框架预留稳定的连续虚拟张量，系统层按请求增长映射物理页，普通注意力 Kernel 继续使用线性指针。它用 GPU 已有的页表替代 Kernel 可见的 Block Table，从而降低新 Kernel 获得动态 KV 能力的适配成本。

这个设计能成立，依赖三个推理负载特征：Decode 的下一轮内存需求可预测；KV 的物理分配带宽远低于 HBM 计算数据流；一次迭代足够长，可以隐藏数微秒级映射调用。细粒度页控制容量浪费，后台映射隐藏延迟，延迟回收与提前分配提高页复用率，这四部分缺一不可。

它也不是 PagedAttention 的无条件替代。定制驱动、固定虚拟槽位、前缀共享、抢占换出和多租户隔离都会把复杂度带回系统工程。更准确的判断标准是：若部署最看重连续 Kernel 的快速接入和长 Prefill 性能，且能够控制驱动栈，vAttention 很有吸引力；若系统高度依赖成熟的块级共享、交换和跨请求组合，PagedAttention 仍拥有更完整的软件生态。

参考资料

1. vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention
2. Microsoft vAttention 官方实现
3. CUDA Driver API：Virtual Memory Management
4. NVIDIA Technical Blog：Introducing Low-Level GPU Virtual Memory Management
5. Efficient Memory Management for Large Language Model Serving with PagedAttention
6. FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
7. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving
8. Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve