PagedAttention 把序列的 KV Cache 划分为固定大小的块,通过块表记录这些块在缓存池中的位置。Attention kernel 根据块表读取 K/V,因此同一条序列的缓存可以分散存放。
它借鉴了操作系统的分页思想,但分页、地址翻译、换入换出分别承担不同职责。现代 GPU 已有硬件 MMU;PagedAttention 在硬件地址翻译之上,增加了一层面向 KV Cache 的软件映射。
KV Cache 如何按块分配
自回归生成持续追加 KV,最终序列长度事先难以确定。如果为每个请求预留一整段最大长度的连续缓存,尚未生成的 token 也会占用分配额度;不同大小的缓存释放后,还可能留下难以复用的空洞。
PagedAttention 允许引擎按固定大小的 KV 块分配空间。每个块容纳固定数量 token 的 K/V,序列增长到需要新块时,再从空闲池取块并更新块表。Prefill 可以一次需要多个块,具体分配量由本轮处理的 token 数和调度策略决定。
vLLM 原始内核设计文档使用过 BLOCK_SIZE=16 的示例。按从零开始的 token 编号展开:
| 逻辑块号 | 对应的 token 位置 | 存储块编号 |
|---|---|---|
| 0 | 0 至 15 | block_table[0] |
| 1 | 16 至 31 | block_table[1] |
| 2 | 32 至 47 | block_table[2] |
左侧块号表达序列中的顺序,右侧表项记录缓存池中的槽位编号。这些槽位可以彼此不相邻。块内 K/V 的具体排列由后端决定,K 和 V 也可以分别存放在两个缓存张量中。
“按需分配”发生在请求与缓存池之间。引擎可以提前申请一大块显存作为缓存池,再把其中的块按需分给请求,无须每生成一个 token 就调用一次显存分配接口。
这减少了请求级的预留浪费,并使同一缓存池中的空闲块可以复用。对于只在尾部追加的单条序列,末尾块仍可能没有填满。分页本身不压缩 K/V 数值,每个已缓存 token 所需的数据量保持不变。
OS 分页与换页的区别
操作系统的分页,把虚拟地址空间和物理内存划分成页,通过页表建立映射。程序使用连续的虚拟地址时,背后对应的物理页框可以分散。
换页则涉及数据搬运。启用磁盘交换区时,内存压力可能促使 OS 把可换出的匿名页写入交换区;程序再次访问已换出的页,会触发缺页异常,由 OS 将数据读回内存、恢复映射,再继续执行。
| 机制 | 处理的事情 | 是否必然访问磁盘 |
|---|---|---|
| 分页 | 按页组织内存并建立映射 | 否 |
| MMU 地址翻译 | 将访问指令中的虚拟地址转换为物理地址 | 否 |
| 磁盘 swap 换入换出 | 在内存和磁盘交换区之间搬运页面数据 | 是 |
所以,即使系统没有启用 swap,仍然可以使用分页。页面已驻留内存时,普通地址翻译也无需磁盘参与。
“活跃内容在内存,不活跃内容在磁盘”只能概括部分回收策略。较少使用的页可以继续留在内存;干净的文件缓存页可以直接丢弃,需要时从原文件重新读取。是否回收以及怎样回收,取决于内存压力、页面类型和 OS 策略。
GPU MMU 与块表是两层映射
现代 GPU 支持虚拟内存和硬件地址翻译。NVIDIA 的虚拟内存管理接口分别提供物理内存分配、虚拟地址预留以及二者之间的映射。
PagedAttention 的块表处理另一件事:序列的某个逻辑 KV 块,当前分配到了缓存池中的哪个槽位。
| 映射层 | 输入 | 输出 | 执行者 |
|---|---|---|---|
| KV 块映射 | 序列的逻辑块号 | 缓存池块号,并据此计算数据指针 | Attention kernel |
| 硬件地址映射 | 数据指针中的 GPU 虚拟地址 | 硬件物理地址 | GPU MMU |
块表由推理引擎维护,GPU 的硬件页表不会自动解释这张表。kernel 必须先查表算出要访问的指针,随后发出的内存访问仍由 GPU MMU 完成地址翻译。
源码中常见的 physical_block_number 表示缓存池中的存储块编号。这个名称使用的是 KV 管理层的“物理块”含义;由它计算出来的指针仍然处于 GPU 虚拟地址空间。
KV 块按 token 数组织,硬件内存页按字节地址组织,两者的大小和边界也无须一一对应。
“逐个 block 做 gather”具体执行什么
这里的 gather 指按索引从分散的位置读取数据。执行过程是:查到某个逻辑块的存储位置,计算块内所需 K/V 的地址,读取数据并参与 attention 运算。
下面摘取 vLLM v0.6.4 的 CUDA 内核展示这一过程。该版本代码用于说明原始实现的寻址细节;当前 vLLM 支持多种 attention 后端。
在 csrc/attention/attention_kernels.cuh 第 257 至 258 行,kernel 先读取块表:
const int64_t physical_block_number =
static_cast<int64_t>(block_table[block_idx]);
block_idx 是当前序列的逻辑块号,表项给出缓存池块号。
第 273 至 275 行,kernel 根据这个块号计算 K 的指针:
const cache_t* k_ptr =
k_cache + physical_block_number * kv_block_stride +
kv_head_idx * kv_head_stride + physical_block_offset * x;
k_cache 是缓存起点,kv_block_stride 是相邻存储块之间的元素跨度,其余部分定位到 KV head 和块内 token。随后还要根据该线程负责的向量片段补上偏移。在无需缓存类型转换的路径中,第 281 至 282 行完成读取:
k_vecs[j] = *reinterpret_cast<const K_vec*>(
k_ptr + offset1 * BLOCK_SIZE * x + offset2);
这就把数据从显存中的对应位置加载到线程使用的寄存器变量中。接下来与 Query 做点积;V 路径同样按块表定位数据,并参与注意力权重的加权求和。
整个过程可以直接在 attention kernel 内完成,无须先把整条序列的 K/V 复制成一个连续大数组。Gather 描述的是索引读取方式,也无须单独启动一个名为 gather 的算子。
“逐个”允许并行执行
按块寻址与 GPU 并行计算可以同时发生。原始设计文档给出的分工示例中,有 4 个 warp、6 个 KV 块:
| 执行单元 | 负责的逻辑块 |
|---|---|
| warp 0 | 0、4 |
| warp 1 | 1、5 |
| warp 2 | 2 |
| warp 3 | 3 |
多个 warp 可以同时处理不同块,每个 warp 内的线程又协作读取块内数据。跨块地址可以分散,块内布局仍可让相邻线程读取相邻地址,利用合并访存。
这里的 KV block 是缓存管理单位,CUDA thread block 是线程组织单位。两者都叫 block,但含义不同。
块表读取、地址计算和结果归并会产生开销。实际性能还取决于块大小、数据布局、并行划分和所用后端,不能仅由“发生了 gather”推断速度。
软件块表查找的固有开销
前文的 gather 中,每一步都要多做一些 OS 交给硬件完成的事情。
| OS 页表 | KV 块表 | |
|---|---|---|
| 执行者 | MMU 硬件电路 | kernel 内的 load 指令与整数运算 |
| 专用缓存 | TLB 保存近期翻译结果,命中时对指令流隐藏 | 无对应结构,块表靠通用 L1/L2 命中 |
| 代价形态 | 命中时该次访存不产生额外指令 | 每块额外取表项、算地址、判边界 |
vLLM 在遍历块的内层循环里读表项:v0.6.4 CUDA 内核 attention_kernels.cuh 第 257 至 258 行,当前 Triton 内核 triton_unified_attention.py 第 429 至 430 行。除取表项本身,还要把 int32 块号转为 int64 再乘 kv_block_stride(v0.6.4 第 229 至 231 行的注释解释了溢出风险),以及为块内未填满部分做分支判断。
原论文 §7.1 有量化结果:
our GPU kernels involve extra overheads of accessing the block table, executing extra branches, and handling variable sequence lengths. As shown in Fig. 18(a), this leads to 20–26% higher attention kernel latency, compared to the highly-optimized FasterTransformer implementation.
该区间只覆盖 attention 算子,Linear 等其他算子不受影响。论文未在该节写出实测的块大小与硬件,默认块大小 16 记录在 §7.2。
vLLM 内核里的缓解写法
合并访存。论文 §5.1:“To ensure coalesced memory access, we assign a GPU warp to read each block.” 当前 Triton 内核一次向量化读取整个 tile 的块号(第 429 至 430 行),同一个块内重复的表项合并为一次访存;Intel Xe2/Xe3 的 tensor descriptor 路径因 BLOCK_SIZE % TILE_SIZE == 0 的保证,退化为标量取表项(第 440 行)。
在线 softmax。每处理一个 tile 就更新运行最大值与分母,点积、softmax、V 累加在同一趟遍历里完成(第 566 至 567 行):
M, L, P, alpha = softmax_step(S, M, L)
acc = acc * alpha[:, None]
长序列再拆成多个 segment 并行,最后由 reduce_segments(第 690 行)合并各段的部分和。
块表本身在 vLLM 内核中没做预加载。v0.6.4 预加载到 shared memory 的是 Query(第 180 行 __shared__ Q_vec q_vecs[...]),块表只按序列取出基址(第 207 行),表项仍在循环内读。块表体量小、复用密集,实际收益来自 L1/L2 命中。
顺带一句版本差异:main 分支的 csrc/attention/ 已只剩类型头文件,paged attention 内核改由 FlashAttention、FlashInfer、Triton 等后端提供。
Decode 阶段为何常看不出开销
CUDA 最佳实践文档:
If the GPU must wait on one warp of threads, it simply begins executing work on another
以及
executing other warps when one warp is paused or stalled is the only way to hide latencies and keep the hardware busy
Decode 每步只处理一个 query token,耗时由读取历史 KV 的显存带宽主导。此时 SM 上驻留多个 warp,一部分 warp 等待 KV 返回时,另一部分正好执行取表项与地址计算。额外开销与访存重叠,不直接叠加到总时长。
成立条件有两个:显存带宽确为瓶颈;驻留 warp 数足够。占用率偏低时掩盖效果消失,查表与分支会显性拉长 kernel。
块表自身的显存
vLLM V1 的块表张量形状为 [max_num_reqs, max_num_blocks_per_req],表项 torch.int32,启动时按最大宽度一次性分配(vllm/v1/worker/block_table.py:112-116),另有同长度的 CPU 侧 np.int32 计数字段(第 117 行)。
以 max_model_len=32768、block_size=16 计:每请求 32768 / 16 = 2048 表项,占 8 KiB;100 并发占 800 KiB。块大小降到 8 时表项翻倍,共 1.6 MiB。较短的序列同样按 max_num_blocks_per_req 占位,这属于为固定形状付出的容量代价。相比之下,每请求的 KV 体量大几个数量级,块表占比很小。
末尾块的内部碎片
序列长度没有整除 block_size 时,末尾块的空槽位不能给其他请求使用。浪费上限为 block_size - 1 个 token。以 Llama 2 70B(GQA,80 层、8 个 KV head、head_dim 128、FP16)为例,每 token KV 占 327680 B(320 KiB):
block_size |
每请求浪费上限 | 均匀分布下的平均浪费 |
|---|---|---|
| 16 | 15 token,4.69 MiB | 8 token,2.5 MiB |
| 32 | 31 token,9.69 MiB | 16 token,5 MiB |
单请求内,这部分占比随序列变长下降。全局看,浪费量乘以并发数,且短请求多时占比上升,对应原论文 §7.2 中 Alpaca 轨迹下大 block_size 表现变差的原因。
block_size 会影响什么
block_size 决定每个 KV 块容纳多少个 token,类似 OS 页大小决定每页容纳多少字节。对于相同模型、缓存精度和布局,增大 block_size 会增大单个 KV 块的字节数。
| 方面 | 小 KV 块 | 大 KV 块 | OS 类比 |
|---|---|---|---|
| 尾部空间浪费 | 最后一个块的空槽较少 | 空槽可能更多,挤占其他请求的容量 | 不足一页的数据仍占整个页 |
| 块表与管理开销 | 相同序列需要更多块,表项和分配管理更多 | 块数与表项更少 | 小页需要更多页表项 |
| GPU 访存与计算 | 每块数据少,块级准备工作占比可能较高,并行能力可能利用不足 | 每块可处理更多 token,但收益受 kernel 实现限制 | 大粒度可以摊薄管理开销,具体硬件机制有差异 |
| 共享与写时复制 | 共享边界更细,按块复制的数据量较小 | 共享边界更粗,修改共享块时可能复制更多数据 | 页大小影响共享与写时复制粒度 |
| 换入换出,若另行启用并按 KV 块搬运 | 搬运粒度细,但小传输的调度开销可能更多 | 容易形成大传输,但可能搬运额外数据 | 类似换页粒度与 I/O 开销的权衡 |
对于单条、仅尾部追加的序列,块大小为 B 时,最后一个块最多闲置 B-1 个 token 槽位。短请求多、并发高时,尾部浪费对容量的影响更明显。
OS 大页还可以让单个 TLB 表项覆盖更大的地址范围。增大 KV 块只改变软件层的缓存组织,不会直接增大 GPU 硬件页,也不会直接扩大 TLB 覆盖范围。
块数增加同样不必然增加 kernel 启动次数:一次 kernel 可以处理多个 KV 块。启用 offload 后,传输实现也可以把多个小块合并搬运,块数与传输次数无须一一对应。
“块级开销难以摊薄”是什么意思
处理一个 KV 块时,kernel 需要查块表、计算块起始地址,再读取和处理块内数据。这些块级准备工作可以服务本块中的多个 token。块内 token 增多时,准备工作无须同比增加,因此平均每个 token 分担的块级开销可以降低。
原论文讨论了每块 16 和 32 个 token 的配置。只比较分块数量,对于长度为 T 且能被 32 整除的同一序列:
| 配置 | 块数 | 一份块级准备工作服务的 token 数 |
|---|---|---|
| 每块 16 个 token | T / 16 |
16 |
| 每块 32 个 token | T / 32 |
32 |
后者需要处理的块数减半,有效 KV 数据总量保持不变。在按块遍历的实现中,更少的块通常意味着更少的块级准备工作,这就是“摊薄”。
这里比较的是块级工作量。线程间如何复用表项、查表是否命中缓存以及实际执行多少条指令,取决于 kernel 实现。KV 数据读取、点积等工作仍然存在,整体延迟不会因此必然减半。
原论文中的块大小实验
原始 PagedAttention 论文 §7.2 中,ShareGPT 工作负载在每块 16 至 128 个 token 时表现较好;序列更短的 Alpaca 工作负载在 16、32 时表现较好,更大的块明显降低性能。论文将其归因于短序列下增大的内部碎片;同时指出块过小可能无法充分利用 GPU 的并行能力。
这些结果对应论文当时的实现和实验条件。块大小的选择需要同时考虑空间浪费与处理效率,也受到请求长度分布和 attention 后端的影响。
KV 块驻留在哪里,显存不足怎么办
在未启用 KV offload 的常规 GPU 推理路径中,已经保留的 KV 块驻留在 GPU 显存缓存池。PagedAttention 允许它们在池中分散存放,并通过块表找到它们。
Gather 在此访问的是显存中的数据,不会因为某个 KV 块较少使用,就自动把它搬到 CPU 内存或磁盘。对于运行中的标准全注意力序列,每次 decode 通常都要读取其全部历史 KV,这与“历史越久越少访问”的直觉也有区别。
显存不足时的处理属于调度与缓存管理策略。vLLM V1 的调优文档规定默认抢占模式为 RECOMPUTE:引擎暂停部分请求,释放可回收的 KV 空间供其他请求使用;等空间足够,再通过重新计算恢复所需 KV。重算会增加请求延迟。
推理系统也可以额外实现 KV offload,把缓存保存在 CPU 内存或其他存储层,需要时再传回 GPU。它与 PagedAttention 的分块组织可以配合使用,但需要传输、容量管理和调度机制。
因此,“KV 分页”说明缓存如何组织和寻址;“KV 换入换出”说明数据是否在不同存储层之间迁移。单独采用 PagedAttention,并不意味着引擎会自动执行 OS 式磁盘换页。
参考资料
- vLLM:PagedAttention 发布介绍,固定大小 KV 块、按需分配和块表映射。
- PagedAttention 原论文 §7.2,块大小对内部碎片、GPU 并行利用和性能的影响。
- vLLM 原始 Paged Attention 内核设计,块内布局与 warp 分工;页面标注为历史设计文档。
- vLLM v0.6.4 CUDA 内核,第 227 至 308 行,K 路径的块表读取、地址计算和点积。
- Linux 内存管理概念,虚拟内存、页表、匿名内存与页面回收。
- NVIDIA GPU 虚拟内存管理,GPU 虚拟地址与物理内存映射。
- vLLM V1 抢占策略,KV 空间不足时的抢占与默认重算行为。
- PagedAttention 原论文 §7.1,块表访问、额外分支与变长序列带来的 kernel 延迟实测。
- CUDA C++ Best Practices Guide,warp 切换与占用率对内存延迟的掩盖。
- vLLM V1 Triton 统一 attention 内核,向量化取块号与在线 softmax。
- vLLM V1 块表缓冲分配,int32 表项与按最大宽度预分配。