Skip to content
kefan.life
Go back

PagedAttention 详解:缓存的存储结构

谈到 PagedAttention1 时,几乎所有介绍都会说它”借鉴了操作系统虚拟内存的分页机制”。灵感来源确实如此,但这篇文章会深入比喻背后的结构和设计思路,看看工程实现里到底是怎么借鉴的。

KV Cache 的显存困境

自回归生成中,每个 token 在每个 layer 产生一对 K, V 向量,需要缓存以避免重复计算(背景见 《为什么 KV 缓存没有 Q》)。对一个正在服务的请求,KV Cache 显存占用随生成长度线性增长,且在生成完成前无法释放。

传统做法是为每个请求预分配 max_seq_len 长度的连续显存。这带来三层浪费:

浪费类型原因
预分配浪费(reservation)实际生成长度远小于 max_seq_len,多余显存空置
内部碎片(internal fragmentation)预分配的固定大小 slot 内部有空隙
外部碎片(external fragmentation)请求结束后释放的显存块大小各异,无法被新请求完整使用

vLLM 论文在 OPT-13B 上的实测数据:有效利用率仅 20.4%(ShareGPT 数据集)到 38.2%(Alpaca 数据集)。超过六成的 KV Cache 显存被浪费了。

内存结构

PagedAttention 把 KV Cache 的管理分成两部分:Block Pool 和 Block Table 负责分配与映射,Attention Kernel 按映射结果读取离散的 KV blocks。

结构示意图

Block Pool

Block Pool 管理 KVCacheBlock 及其 block ID,同时维护 free queue、ref_cnt 和 prefix cache 索引。真正的 K/V 数据位于各 layer 的 KV cache tensor 中。一个 block ID 是调度器分配的统一 block 编号,同一 KV cache group 内的各 layer 用它索引各自对应的 KV cache block。

因此,block_id = 3 表示 block table 指向 ID 为 3 的 KV cache block。同一 group 内的各 layer 都用这个 ID 索引自己的 KV cache。各 layer 的 K/V 仍然分开存放,共享的是 block ID 和 block table。每个 block 容纳固定数量的 token slots,记作 BB,常见取值是 16 或 32。对普通 full attention layer 来说,这个 block 包含该 layer 中 BB 个 token、所有 KV heads 的 K/V 数据。

这套结构与操作系统分页相似:

操作系统分页PagedAttention
物理内存KV cache tensors
物理页帧号(PFN)block ID(KV cache block 的整数索引)
页表Block Table(per-request、per-KV-cache-group 的 block ID 数组)
地址翻译:PFN × 页大小 + 偏移地址翻译:cache_base + block_id × block_stride + offset

OS 的地址翻译由硬件 MMU 完成,应用程序只看到连续的虚拟地址。PagedAttention 则把查表和 block 内偏移计算写进 Attention Kernel,由软件显式完成。

这也是它与 glibc malloc 的区别。malloc 返回连续虚拟地址,调用者无需关心底层分页;PagedAttention 传给 Kernel 的是 block table,Kernel 必须按 token 位置查表,找到对应的 KV cache block。它借用了页表的间接寻址思路,但没有 MMU 替它屏蔽地址转换。

Block Table

每个请求在每个 KV cache group 上维护一张 block table。对于只包含 Full Attention 的普通模型,attention layers 属于同一个 group,因此共享一张表;hybrid attention 模型可能有多个 group,分别维护映射。

以一个已生成 35 个 token 的请求为例(B=16B = 16),某个 group 的 block table 为 [7, 3, 12]

下标block ID对应 token 位置
07token 0–15(满)
13token 16–31(满)
212token 32–34(3/16 已填充)

Attention Kernel 按顺序遍历历史 KV 时,block_idx 是 block table 的数组下标,对应上表中一段最多 BB 个 token 的位置;数组元素是实际 block ID:

  1. 查物理位置block_table[block_idx] → 得到 pool 中的 block ID,算出显存地址
  2. 算 token 位置token_idx = block_idx * B + offset → 得到这个 token 在序列中的逻辑位置

新分配的 block ID 会追加到对应 block table 的末尾,因此数组顺序与 token 在请求中的顺序一致。同一 group 内的各 layer 随后用相同的 block ID 索引自己的 KV cache。

另外需要区分两个阶段使用的不同数据结构:

软件地址翻译

连续 KV Cache 只需要通过 base + token_offset 定位。但使用 Paged Attention 后,相邻 token 段对应的 KV cache blocks 可能拥有任意 block ID,Kernel 需要先查表,再计算 block 内偏移。

下面以 vLLM 自研 CUDA Kernel 的 attention_kernels.cuh 为例。注意:实际部署中,vLLM 通常会根据硬件和配置选择 FlashAttention、FlashInfer 等 backend;具体 KV layout 会变化,但 block table 的间接寻址作用不变。另外为了突出本文主线,下面先假设没有 context parallelism,且 KV cache group 的 block size 也与 Kernel 使用的 block size 相同。

读取历史 K/V 时,Kernel 主要接收:

attention_kernels.cuh 使用下面的 K/V 布局:

key_cache:   [num_blocks, num_kv_heads, head_size / x, block_size, x]
value_cache: [num_blocks, num_kv_heads, head_size,     block_size]

K 和 V 在数学上都是每个 token、每个 KV head 一个 head_size 维向量,两种布局的元素数也相同;形状不同是为了适配两步计算的访存方向:计算 QKTQK^T 时需要沿 head_size 读取同一 token 的 key,因此 key 将 head 维拆成每组 x 个连续元素;计算 softmax(QKT)V\operatorname{softmax}(QK^T)V 时则要在固定 head 维上遍历 token,因此 value 把 block_size 放在最内层。这是针对 Kernel 访存模式的布局优化,不是 K/V 的语义维度或大小不同。

cache_t 表示 KV Cache 在显存中的存储类型,可以是 FP16、BF16 或 FP8。x = 16 / sizeof(cache_t),表示 16-byte 访存单元包含多少个 cache 元素。FP16/BF16 中每个元素占 2 bytes,因此 x = 8

源码遍历 block_idx,再用 block_table[block_idx] 取得实际 block ID。如果从单个 token 位置 i 理解这个过程,对应关系是:

block_idx             = i / block_size
physical_block_offset = i % block_size
physical_block_number = block_tables[seq_id][block_idx]

随后再定位 KV head 和 head 维内部的数据。为了区分 KV cache tensor 和指向其首元素的基地址,下面将源码中的 k_cache 指针写作 key_cache_base,key 读取可以简化为:

# kv_head_idx 选择当前 query head 对应的 KV head。
# physical_block_number 选择当前 layer KV cache 中的 block。
# physical_block_offset 选择 block 内的 token slot。

k_ptr = key_cache_base
      + physical_block_number * kv_block_stride
      + kv_head_idx * kv_head_stride
      + physical_block_offset * x

# key_cache 将 head_dim 拆成 [head_size / x, x]。
addr = k_ptr + block_size * x

Value 的布局把 head_size 放在 block_size 之前。同理,下面用 value_cache_base 表示源码中 v_cache 指向的基地址:

v_ptr = value_cache_base
      + physical_block_number * kv_block_stride
      + kv_head_idx * kv_head_stride

addr = v_ptr + head_dim_idx * block_size + physical_block_offset

kv_block_stridekv_head_stride 直接取自 tensor stride。对上面的连续布局,它们分别是:

kv_block_stride = num_kv_heads * head_size * block_size
kv_head_stride  = head_size * block_size

Kernel 每处理 block table 中的一项,都要先取出 block ID,再在对应 KV cache block 内执行普通的偏移计算。只接受连续 KV 的 Attention Kernel 无法直接读取这种布局;支持 PagedAttention 的 backend 必须把 block table 纳入寻址过程。

一个具体的例子

假设 block_size = 16head_size = 128num_kv_heads = 32,KV Cache 是 FP16,因此 x = 8。请求 ss 的 block table 为 [7, 3, 12]。现在 Kernel 要访问位置 i = 20kv_head_idx = 5 的 key。下面的数字都是以 cache_t 元素为单位的偏移,换成字节偏移时还要乘以 sizeof(cache_t)

block_idx             = 20 / 16 = 1
physical_block_offset = 20 % 16 = 4
physical_block_number = block_tables[s][1] = 3

kv_block_stride = 32 * 128 * 16 = 65536
kv_head_stride  = 128 * 16 = 2048

k_ptr = key_cache_base + 3 * 65536 + 5 * 2048 + 4 * 8
      = key_cache_base + 206880

physical_block_number 决定读取哪个 KV cache block,kv_head_idx 决定读取哪个 KV head,physical_block_offset 决定 block 内的 token slot。

写入新 K/V

读取历史 KV 时,Attention Kernel 需要遍历 block table。写入本轮新产生的 K/V 时,vLLM 会提前把目标位置压缩成 slot_mapping

block_table_idx = position / block_size
offset_in_block = position % block_size
block_id        = block_tables[seq_id][block_table_idx]

slot_mapping = block_id * block_size + offset_in_block

vLLM 在 cache_kernels.cu 中再将 slot_mapping 拆回 block ID 和 block 内偏移:

block_id       = slot_mapping / block_size
offset_in_block = slot_mapping % block_size

这样一来,读写两条路径各自保留最适合的输入形式:Attention 读取需要整张 block table 来遍历历史上下文,cache 写入只需要知道每个新 token 的目标 slot。

使用流程

分配与增长

请求到达后,分配器先查缓存再分配。在线服务中大量请求共享相同前缀(如 system prompt),重复计算这些 KV 是浪费。

  1. 将 prompt tokens 按 BB 切分为 token blocks
  2. 对每个完整 token block 计算链式 content hash:hi=hash(tokensi,  hi1)h_i = \text{hash}(\text{tokens}_i,\; h_{i-1})。相同前缀会产生相同的 hash 序列
  3. 用 hash 在 block cache 中逐级查找:
    • 命中:复用已有 block,增加 ref_cnt,并把对应 block ID 填入该请求的 block table
    • 未命中:从 free list 分配新 block
  4. Prefill:只对未命中的 token blocks 计算 K, V,并通过 slot_mapping 写入各 layer 的 KV cache;命中的 blocks 直接参与注意力计算
  5. Decode:每生成一个 token,将其 K, V 写入当前 block 的下一个空位,写满时从 free list 取新 block,追加到对应 group 的 block table
  6. 一个 block 填满且其中的 token 已经 finalized 后,就可以按 content hash 加入 block cache,不必等待整个请求结束

system prompt 等高频前缀首次计算后,完整 blocks 会逐步进入缓存。后续请求可以复用这些 block IDs,prefill 只计算未命中的尾部 token。显存也随序列增长逐步分配,不需要一次预留 max_seq_len

回收与抢占

这里假设 KV Cache 只存放在 GPU。vLLM 也提供 KV offloading,SGLang 的 HiCache 则进一步利用 CPU 内存等外围存储扩展缓存层级,这些机制不在本文展开。

性能提升的来源

原始 PagedAttention 论文中的实验显示,其自研 Kernel 单次执行比当时 FasterTransformer 的优化 Kernel 慢约 20–26%。block table 查找和离散 block 访问会带来额外开销。注意:这是论文的口径,不能直接当作当前 FlashAttention、FlashInfer 等 backend 的性能结论。

吞吐量提升来自系统层面:更高的显存利用率 → 更大的 batch size → 更高的 GPU 利用率。当有效利用率从 20%–38% 提升到接近 100% 时,同等显存下能并发服务的请求数大幅增加。单次 Kernel 变慢了,但并发度变高了,总吞吐提升 2–4 倍甚至更多。

Block 大小 BB 的选择是一个经典 tradeoff:

论文实测 B=1B = 1 时性能灾难性下降:内部碎片为零,但注意力计算退化为纯随机访存。实践中 B=16B = 16 是常见选择。

vAttention:真正的透明映射

PagedAttention 的策略是重写 Kernel 来适应分散的内存,vAttention2 正好反过来:让分散的物理内存看起来连续,让 Kernel 不用改

vAttention 使用 CUDA Virtual Memory Management API(cuMemMap),将不连续的物理显存页映射到一段连续的虚拟地址空间。Attention Kernel 看到的仍然是连续内存,FlashAttention 等标准 Kernel 无需修改。这才是名副其实的 OS 虚拟内存类比:硬件级地址映射,对上层完全透明。

两种方案代表了一个经典的系统设计 trade-off,透明性与可控性

PagedAttentionvAttention
地址翻译软件显式(CUDA Kernel 代码)硬件透明(CUDA VMM)
Kernel 侵入性高:必须重写 Attention Kernel无:标准 FlashAttention 直接可用
分配粒度灵活(block 通常几十 KB)受限(CUDA VMM 最小页 2MB)
跨进程共享兼容标准 CUDA IPC不兼容,需 POSIX fd + UDS 替代路径
驱动依赖需要特定 CUDA 驱动版本
生态适配FlashAttention / FlashInfer 已深度适配需要生态全面重新适配

vAttention 追求透明性:上层不感知分页,Kernel 不用改。代价是控制权下沉到硬件和驱动层,分配粒度、跨进程共享、调度策略都不再由应用层决定。

vLLM 核心维护者在 GitHub issue #176123 中明确拒绝了 vAttention 集成,理由涵盖上述多个维度。其中 CUDA IPC 冲突尤为致命:PD-sep 架构需要跨进程零拷贝共享 GPU 显存,cudaIpcGetMemHandle 不支持 VMM 分配的内存,替代方案链路复杂且不成熟。加上 cuMemMap 自身的调用开销(约 120 次调用 ~5ms)和 2MB 最小页粒度对细粒度调度的破坏,vAttention 在工业部署中的收益远不及其引入的系统复杂度。

这里很难说哪个方案更好。vAttention 引入的 2MB 页粒度、驱动依赖和 IPC 不兼容,同样是沉重的负担。更准确的说法是:现有生态选择了 PagedAttention 这条路径。CUDA 的 VMM API 成熟较晚,FlashAttention、FlashInfer 等高性能 Kernel 已围绕显式分页 KV 布局构建了整套生态,PD-sep 等分布式架构也依赖标准 CUDA IPC。在这个既成事实上,vAttention 的切换成本远大于收益。如果 CUDA 从一开始就提供成熟的虚拟化地址管理,故事可能完全不同。

结语

PagedAttention 的本质是一个用户态的显存分页系统:Block Pool 负责分配 block ID,Block Table 按 token 顺序保存请求使用的 block ID,Attention Kernel 遍历 block table,查表后读取对应 KV cache block 内的数据。分配、寻址、共享和回收都由软件显式管理。额外的查表与 Kernel 复杂度换来了更高的显存利用率、更细的调度粒度,以及与分布式架构的兼容。

参考资料

  1. Kwon, W. et al., “Efficient Memory Management for Large Language Model Serving with PagedAttention”, SOSP ‘23. arXiv:2309.06180

  2. Prabhu, S. et al., “vAttention: Dynamic Memory Management for Serving DNN Inference with LLMs”, 2024. arXiv:2405.04437

  3. vLLM Issue #17612, “Support vAttention”. github.com/vllm-project/vllm/issues/17612


Share this post on:

Previous Post
Roofline 分析:瓶颈的判定与局限
Next Post
为什么KV缓存没有Q