KV cache 显存管理:为什么 PagedAttention 是 vLLM 的核心
到目前为止,我们已经搭出了一个可以持续推进多个请求的推理引擎。
一个请求进入系统后,会被抽象成 Sequence。Scheduler 决定它什么时候 prefill,什么时候 decode。ModelRunner 负责执行模型。Sampler 负责采样下一个 token。Engine.step 则像心跳一样不断推进整个系统,并把增量输出交给服务层。
这个结构已经很接近一个在线 LLM serving 系统。
但它还有一个核心问题没有解决。
KV cache 怎么管理?
前面的文章里,我们使用 past_key_values 来表示一个请求的 KV cache。这样做非常适合理解 KV cache 的语义:prefill 创建 cache,decode 读取并追加 cache,请求结束后释放 cache。
但如果我们要实现一个高吞吐推理引擎,仅仅知道“有 KV cache”是不够的。真正困难的是:当系统里同时有大量请求,每个请求长度不同、生成进度不同、结束时间不同,KV cache 应该如何分配、增长、复用和释放?
这就是 vLLM 最核心的设计之一。
vLLM 的 PagedAttention 并不是改变了 Transformer 的数学公式。它解决的是一个更系统层的问题:如何让 KV cache 的显存管理像操作系统管理虚拟内存一样灵活。
这一篇我们不急着写 block manager,也不急着实现 attention kernel。我们先把问题讲清楚:为什么普通 KV cache 管理方式会浪费显存,为什么连续内存不适合在线推理,为什么分页思想可以解决这个问题,以及 PagedAttention 到底把 attention 计算中的哪一层抽象改掉了。
KV cache 为什么会变成显存管理问题
先回到 KV cache 本身。
在 decode 阶段,每生成一个新 token,模型都要把这个 token 在每一层对应的 K 和 V 保存下来。下一个 token 生成时,会拿自己的 Q 去访问历史所有 K 和 V。
所以 KV cache 会随着上下文长度不断增长。
一个请求的上下文长度等于:
prompt tokens 数量 + generated tokens 数量
如果 prompt 有 2048 个 token,模型生成了 512 个 token,那么这个请求已经需要保存 2560 个 token 的 KV cache。
这还只是一个请求。
在线服务里,系统可能同时处理几十、几百甚至更多请求。每个请求都有自己的 KV cache。请求越多,上下文越长,显存占用越大。
KV cache 的大小大致可以用这个公式理解:
KV cache 显存 ≈ 层数 × token 数 × KV head 数 × head_dim × 2 × dtype_size
这里的 2 表示 K 和 V 两份。
这个公式背后有两个重要结论。
第一,KV cache 显存和 token 数线性相关。
第二,KV cache 显存和并发请求数线性相关。
这意味着 LLM serving 的显存压力不只是来自模型权重,还来自正在服务的请求本身。
模型权重加载后相对固定。
KV cache 则是动态变化的。
请求进入,KV cache 增加。
请求生成,KV cache 继续增加。
请求结束,KV cache 释放。
新的请求又进来,再重新申请 cache。
所以在推理服务里,KV cache 不是一个普通张量,而是一种动态显存资源。
这也是为什么它需要专门的管理模块。
最直观的做法:给每个请求分配一段连续 cache
如果我们从零开始实现,很容易想到一种简单方案。
每个请求进来时,根据最大上下文长度,为它分配一段连续 KV cache。
比如系统支持最大上下文长度 4096。那就给每个请求准备 4096 个 token 的 cache 空间。
这样实现最简单。
每个请求的第 t 个 token,就写到自己的第 t 个位置。
逻辑非常清楚:
sequence A:
cache[0], cache[1], cache[2], ...
sequence B:
cache[0], cache[1], cache[2], ...
decode 时也很简单。
当前请求长度是多少,就读取前面多少个位置的 K 和 V。
但这个方案很快会暴露一个大问题:浪费太严重。
假设一个请求最多允许生成到 4096 token,但实际只用了 512 token 就结束了。那么剩下 3584 个 token 的 cache 空间在这个请求生命周期内都被浪费了。
更现实的是,线上请求长度高度不确定。
有的请求只问一句话。
有的请求带着长文档。
有的请求生成几十个 token。
有的请求生成上千个 token。
如果我们按照最大长度预分配,绝大多数请求都会浪费大量 cache 空间。
这类浪费不是小问题。
在 LLM serving 中,KV cache 往往会成为限制并发数的关键因素。浪费的 cache 越多,同时能服务的请求越少,系统吞吐就越低。
所以第一个矛盾出现了:
为了实现简单,我们想提前分配连续大块空间
为了提高并发,我们又不能接受大量未使用空间
这正是连续分配的问题。
按需增长也不够:碎片化会出现
也许我们可以换一种方式。
不要一开始就给每个请求分配最大长度,而是按需增长。
请求刚进来时,只根据 prompt 长度分配 cache。每生成一个新 token,再继续追加空间。
这样看起来可以减少预分配浪费。
但它会带来新的问题:显存碎片化。
我们用一个简单例子理解。
假设显存中有一段可以放 KV cache 的空间,一开始是连续空闲的。
请求 A 进来,占用一段空间。
请求 B 进来,占用一段空间。
请求 C 进来,占用一段空间。
过了一会儿,请求 B 结束了,它释放中间那段空间。
现在显存变成这样:
A 占用 | 空闲 | C 占用 | 空闲
这时来了一个新请求 D。它需要一段比较大的连续空间。
虽然总空闲空间可能足够,但这些空闲空间分散在不同位置,不一定能组成一个足够大的连续块。
这就是碎片化。
如果 KV cache 要求每个请求都占用连续内存,碎片化就会让显存越来越难用。
更麻烦的是,LLM 请求的生命周期非常动态。
请求不断进入。
请求不断生成。
请求不断结束。
KV cache 不断增长和释放。
这和操作系统里的内存管理很像。操作系统也不能要求每个进程都拿到一整段连续物理内存。否则内存碎片会非常严重。
所以操作系统引入了分页。
进程看到的是连续虚拟地址。
但实际物理内存可以是不连续的页。
vLLM 的 PagedAttention 借鉴的正是这个思想。
分页思想:逻辑连续,物理不连续
分页思想的核心很简单:
逻辑上连续。
物理上不一定连续。
放到 KV cache 里,可以这样理解:
对于一个 Sequence 来说,它的 token 序列是连续的。
token 0, token 1, token 2, token 3, ...
从 attention 的语义看,它当然应该像一个连续上下文。
第 t 个 token 生成时,它要看到前面所有 token 的 K 和 V。
但这些 token 的 KV cache 在物理显存中不一定要连续存放。
我们可以把 KV cache 切成固定大小的 block。
例如每个 block 存 16 个 token 的 K 和 V。
那么一个长度为 40 的 Sequence 需要 3 个 block。
block 0: token 0 到 token 15
block 1: token 16 到 token 31
block 2: token 32 到 token 39
逻辑上,这 40 个 token 是连续的。
物理上,这 3 个 block 可以放在显存中的不同位置。
Sequence 只需要维护一张 block table,记录自己的第几个逻辑 block 对应哪个物理 block。
Sequence A 的 block table:
logical block 0 -> physical block 7
logical block 1 -> physical block 2
logical block 2 -> physical block 15
attention 计算时,不再假设 KV cache 是一整段连续数组,而是通过 block table 找到每个 token 对应的物理位置。
这就是 PagedAttention 的核心变化。
它让请求看到连续上下文,同时允许底层 KV cache 使用不连续的物理 block。
为什么固定大小 block 可以减少浪费
使用 block 后,请求不需要一次性拿到最大上下文长度的空间。
它只需要随着 token 增长逐步申请 block。
比如 block size 是 16。
一个请求当前长度是 17,那么它需要 2 个 block。
第一个 block 存 16 个 token。
第二个 block 只用了 1 个 token。
这时浪费最多发生在最后一个 block。
也就是说,block 机制把浪费从“最大长度减实际长度”缩小到了“最后一个 block 里的剩余空间”。
如果 block size 是 16,一个请求最多浪费 15 个 token 的 cache 空间。
相比连续预分配里可能浪费几千个 token,这个差距非常大。
当然,block size 也不是越小越好。
block 太小,block table 会更长,管理开销更大,attention kernel 访问 block 的次数也更多。
block 太大,内部碎片会变多。
所以 block size 是一个取舍。
它要在管理开销和空间浪费之间平衡。
这也是系统设计里经常出现的主题:没有免费的优化,只有更合适的取舍。
PagedAttention 改变了什么
现在我们可以更准确地说,PagedAttention 改变的不是 attention 的数学含义,而是 attention 访问 KV cache 的方式。
普通 attention 可以假设 K 和 V 是连续的。
K cache: [token 0, token 1, token 2, token 3, ...]
V cache: [token 0, token 1, token 2, token 3, ...]
PagedAttention 下,K 和 V 被放进多个物理 block。
logical token 0 到 15 -> physical block 7
logical token 16 到 31 -> physical block 2
logical token 32 到 47 -> physical block 15
attention kernel 在读取第 t 个历史 token 的 K 和 V 时,需要先通过 block table 找到它属于哪个物理 block,再找到 block 内偏移。
可以把映射关系理解成:
logical token index
-> logical block index
-> physical block id
-> offset inside block
-> KV cache address
对于 block size 为 16 的情况:
logical_block_id = token_index // 16
block_offset = token_index % 16
physical_block_id = block_table[logical_block_id]
这就是分页思想在 attention 中的体现。
从模型语义上看,当前 token 仍然能访问完整历史上下文。
从显存管理上看,历史上下文被拆成了多个可复用的 block。
这样,KV cache 的分配和释放就变得灵活很多。
Block manager 是 PagedAttention 的前置条件
很多人一提到 PagedAttention,就会立刻想到 CUDA kernel。
但在我看来,理解 PagedAttention 更好的入口不是 kernel,而是 block manager。
因为如果没有 block manager,PagedAttention 没有意义。
PagedAttention 假设 KV cache 已经被组织成 block,并且每个 Sequence 都维护了自己的 block table。
那么谁负责这些事情?
就是 block manager。
它要负责:
维护所有空闲 physical blocks
为新请求分配 blocks
在 decode 时为请求追加 block
维护 Sequence 的 block table
在请求结束时释放 blocks
判断当前显存还能不能接收新请求
换句话说,block manager 是 KV cache 的内存分配器。
PagedAttention 是基于这种 block 化内存布局进行 attention 计算的方法。
两者关系可以这样理解:
block manager 负责存在哪里
PagedAttention 负责怎么读取
这一点非常重要。
如果只讲 PagedAttention kernel,很容易把问题理解成某个底层算子的优化。
但实际上,vLLM 的核心价值在于把推理系统里的 KV cache 生命周期管理和 attention 计算方式结合了起来。
也就是说,PagedAttention 不只是 kernel 优化,它是系统设计。
对 Scheduler 的影响
引入 block 管理之后,Scheduler 也会发生变化。
在前面的版本里,Scheduler 只看两个预算:
token budget
sequence 数量
但现在它还必须关心 KV cache block 是否足够。
因为一个 waiting 请求要做 prefill,就需要为 prompt tokens 分配 cache block。
一个 running 请求要做 decode,也可能需要为新 token 分配新的 cache slot。如果当前 block 已满,还需要申请一个新 block。
所以调度前必须先问一个问题:
如果这批请求被执行,KV cache 空间够不够?
这会改变 Scheduler 的逻辑。
以前,一个请求只要 token budget 放得下,就可以调度。
现在,一个请求还必须满足 cache budget。
比如 waiting 队列里有一个 prompt 长度为 4096 的请求。token budget 允许它进入,但剩余 KV cache blocks 不够,那么它仍然不能被调度。
又比如 running 队列里有很多请求都即将写满当前 block。下一轮 decode 会导致它们同时申请新 block。如果 free blocks 不够,Scheduler 就不能盲目把它们全都调度进去。
这说明 Scheduler 和 block manager 必须协作。
Scheduler 负责决策。
Block manager 负责资源判断和分配。
Engine 执行完成后,再根据 Sequence 是否 finished 释放资源。
推理系统的主循环因此变成:
Scheduler 选择候选请求
Block manager 判断 cache 资源
ModelRunner 执行 prefill 或 decode
Sequence 更新状态
Finished 请求释放 block
从这一刻开始,调度不再只是计算调度,也变成了显存资源调度。
对 Sequence 的影响
Sequence 也不再适合直接保存 past_key_values。
在 PagedAttention 版本里,Sequence 更应该保存自己的 block table。
也就是说,它不再持有一整块 KV cache 张量,而是持有一组 block 引用。
概念上可以这样看:
Sequence:
prompt_token_ids
generated_token_ids
block_table
context_len
status
stop_reason
其中 block table 描述这个 Sequence 的逻辑上下文如何映射到物理 KV cache blocks。
这会让 Sequence 变轻。
它不再拥有 cache 本体。
cache 本体由全局 KV cache pool 管理。
Sequence 只保存引用。
这种设计和操作系统里的进程页表很像。
进程看到自己的虚拟地址空间。
物理内存由操作系统统一管理。
Sequence 看到自己的逻辑 token 序列。
物理 KV cache blocks 由 block manager 统一管理。
这种间接层正是系统变灵活的关键。
对 ModelRunner 的影响
ModelRunner 的输入也会变。
在使用 past_key_values 时,ModelRunner 只需要把每个请求对应的 cache 对象传给模型。
但在 PagedAttention 版本里,它需要向 attention kernel 提供更多结构化信息。
例如:
当前输入 token
每个 Sequence 的上下文长度
每个 Sequence 的 block table
KV cache 的物理 block pool
当前 token 应该写入的 slot
其中 block table 用来读取历史 KV。
slot 信息用来写入当前 token 的 KV。
这也是为什么 PagedAttention 的实现比普通 attention 更复杂。
普通 attention 可以直接按连续地址读 K 和 V。
PagedAttention 必须在读取时进行逻辑地址到物理地址的转换。
但这份复杂度换来的,是更高的显存利用率和更灵活的调度能力。
对于在线 serving 来说,这是非常值得的。
PagedAttention 带来的系统视角变化
到这里,我们可以从更高层次理解 PagedAttention 的意义。
在没有 PagedAttention 之前,KV cache 更像是每个请求自己的私有连续数组。
请求之间的 cache 管理相对割裂。
系统很难高效复用碎片空间。
在 PagedAttention 之后,KV cache 变成了一个全局资源池。
所有请求都从这个资源池里申请 fixed size blocks。
请求结束后,blocks 回到资源池。
新请求可以复用这些 blocks。
Sequence 通过 block table 维护自己的逻辑上下文。
attention 通过 block table 读取正确的物理 cache。
这带来的变化非常大。
显存浪费从大块预留变成最后一个 block 的内部碎片。
请求结束后释放的 block 可以快速复用。
不同请求的 KV cache 不要求物理连续。
Scheduler 可以根据 free block 数量做更精确的准入控制。
系统可以承载更多并发请求。
这就是为什么说 PagedAttention 是 vLLM 的核心。
它不是某个孤立的小优化,而是改变了整个推理引擎对 KV cache 的组织方式。
这一篇先不实现 kernel
讲到这里,可能会很想直接写 PagedAttention。
但我们还不应该这么做。
因为在实现 attention kernel 之前,必须先有 block manager。
如果没有 block manager,就没有 block table。
没有 block table,PagedAttention 的输入就不完整。
所以更合理的路线是:
先实现 block manager
再实现基于 block table 的 KV cache 布局
再实现教学版 PagedAttention
最后再考虑高性能 kernel
下一篇文章,我们会先实现 block manager。
它不涉及复杂 attention 计算,但它会决定整个系统如何分配、追加和释放 KV cache blocks。
它要回答这些问题:
一个 block 存多少 token?
free block list 如何维护?
prefill 时需要分配多少 block?
decode 时什么时候需要追加 block?
Sequence 的 block table 如何更新?
请求结束后如何释放 block?
只有这些问题解决了,PagedAttention 才真正有落脚点。
小结
这一篇我们进入了 vLLM 最核心的思想:KV cache 的分页式管理。
KV cache 是 LLM 推理服务中最重要的动态显存资源。它会随着请求进入而创建,随着 decode 增长,随着请求结束释放。简单的连续分配方式容易造成大量预留浪费,而按需连续增长又容易遇到显存碎片化。
PagedAttention 的核心思想是:让 Sequence 在逻辑上拥有连续上下文,但在物理上使用不连续的 KV cache blocks。
Sequence 通过 block table 记录逻辑 block 到物理 block 的映射。attention 计算时,再通过这个映射找到正确的 K 和 V。
这样,KV cache 就从每个请求私有的大块连续内存,变成了一个全局可复用的 block pool。
这带来的不只是显存节省,还有调度能力的提升。Scheduler 可以根据 free blocks 做准入控制,finished 请求可以快速归还 blocks,新请求可以复用这些 blocks。整个推理系统因此变得更适合高并发在线服务。
下一篇文章,我们会真正实现 block manager。
这会是 mini vLLM 从“使用 KV cache”走向“管理 KV cache”的关键一步。