vLLM:KV Cache技术原理解析

vLLM:KV Cache技术原理解析
Chase Woo在使用 vLLM 等框架部署大模型时,总有一个技术概念逃不掉:KV Cache。在 vLLM 等高性能推理框架中,KV Cache 是加速大模型推理、提升吞吐量的核心机制。
用一句话总结KV Cache就是:在处理生成请求时,缓存并复用之前 Token 产生的中间计算结果(Key 和 Value 张量),使得后续计算无需重复执行相同的矩阵运算,从而实现加速推理的效果。
一、大模型的基本框架
首先我们要知道大模型在推理时的两个阶段:Prefill 阶段和 Decode 阶段。
Prefill 是大模型计算提示词 QKV 张量的阶段。
大模型在这个阶段需要将所有提示词进行一次全量的并行计算。由于自注意力的计算复杂度是 $O(n^2d)$(n 是 Token 的个数,d 是 Token 的隐藏维度),所以 Prefill 的耗时取决于提示词的长度,且是平方级增加。显卡在计算这么大的张量计算时会占用大量算力,所以 Prefill 是计算密集型任务。Decode 是推理时的吐字阶段。
大模型是自回归的,推理时每次只能生成一个 Token,每一个 Token 都是前面所有的 Token 的一次概率计算,直到终止符或达到输出长度上限。
KV Cache 的优化原理就在于此:由上图可知,输入 Input 时,需要进行一次完整的自注意力计算,输出时的 Output 张量用到了 Input 的张量。当模型生成第 t 个 Token 时,它需要将当前 Token 与前面 t-1 个 Token 进行自注意力计算。如果在 Decode 阶段重新计算前面所有 Token 的自注意力,单步复杂度依然极高。但我们将之前已经计算过的 Token 的 Key 和 Value 张量保存在显存中。当预测下一个 Token 时,只需计算当前新 Token 的 $Q_t, K_t, V_t$,然后将其 $K_t, V_t$ 拼接到缓存中,再用 $Q_t$ 与完整的 $K, V$ 缓存进行注意力计算即可,这样就省下了大量重复计算。由于显卡需要不停地读写缓存,而计算只占较小一部分,所以 Decode 阶段是访存密集型任务。
二、vLLM 的 PagedAttention
早期的 HuggingFace Transformers 库通常会为每个请求预先分配最大长度(如 2048)的连续显存。但模型实际生成的长度往往不可知,导致大量显存闲置,产生严重的内部碎片。即使空闲显存总量足够,但如果不是连续的大块,也无法分配给新的请求。
vLLM 引入“页(Page)”或“块(Block)”的概念。将逻辑上连续的 KV Cache 划分为固定大小的块(每个块一般包含 16 个 Token 的 KV 数据)。逻辑上相邻的 Token,在物理显存上的 Block 可以是完全分散的。vLLM 通过一张 Block Table(块表) 来记录逻辑块和物理块的映射关系。这彻底消除了静态分配内存带来的浪费问题,使得 vLLM 能够批处理更多的请求。
关于 Block 设计的几个核心疑问:
- 为什么要划分 Block ,为什么不直接用 Token ?
GPU 计算极其依赖连续的内存访问来实现高带宽,使用 Token 作为颗粒度的话,就会导致显存地址很分散,这样会大大影响数据加载的速度。假如说有160个 Token 并使用 Token 作为颗粒度,最多会有160处不连续的显存读取,严重拖慢访存速度,换做是 Block 就最多只有10次不连续访问了。 - 为什么 Block Size 要设置为16,而不是越大越好呢?
Block 越大,显存分布越连续,访存效率更高,但多数情况下,请求的序列长度并不会被 Block Size 整除。如果余数是1,那也要为它单独申请一个 Block,那 Block Size 越大,最后一个 Block 中浪费的显存(内部碎片)也就越多。 - 为什么 Block Size不设为动态的呢?
如果根据请求动态分配不同大小的 Block,显存分配又会退化回自由申请的状态。随着请求的不断创建和销毁,显存池中会产生大量的外部碎片,最终可能导致有效显存利用率连 30% 都不到。固定大小的 Block 才能使 PagedAttention 达到最高效率。
三、KV Cache 的预热与分配
我们知道,在每次启动 vLLM 的时候,都需要指定一些参数,比如 gpu_memory_utilization, enable_prefix_caching 等,这些参数会控制内存占用或优化推理。vLLM 在加载完模型权重后,并不会立刻启动服务,而是会进行一次模拟推理。
其目的是为了探测模型在最大并发和最大序列长度下,各种中间激活值究竟会占用多少显存。同时,vLLM也会计算可用 Block 数量:
剩余可用显存** = GPU 总显存 - 模型权重占用 - 探测出的峰值激活占用 - 预留安全空间
然后将剩余显存除以单个 Block 的大小,计算出系统可以维护的 num_gpu_blocks。同时 vLLM 也会在 CPU 内存中分配一部分空间作为 num_cpu_blocks,用于后续由于显存不足发生抢占时的换页操作(Swapping)。
最后,vLLM 会根据探测出的 Block 数量,一次性在 GPU 和 CPU 上申请两块巨大的、连续的 Tensor,这就是整个系统生命周期内的 KV Cache 内存池。此后处理请求时,vLLM 不再频繁向 CUDA 申请释放内存,而是完全接管并自己维护这片内存池的分配与回收。
四、 KV Cache 的前缀缓存(Prefix Caching)
前面提到了 vLLM 是通过一个 Block Table 来维护并使用 KV Cache 的。当开启 Automatic Prefix Caching(自动前缀缓存,APC) 功能后,vLLM 会对 Block 的分配引入哈希(Hash)机制。
当分配 Block 时,vLLM 会为每一个 Block 都计算 Hash 值,其基础内容是该块内包含的 Token IDs。为了体现上下文的顺序关系,从提示词的第二个 Block 开始,vLLM 会将上一个 Block 的 Hash 值也作为当前 Block Hash 计算的一部分。这样一来,每一个 Block 的 Hash 都隐式地包含了其前面所有 Token 的顺序信息,形成了一条类似哈希链(Hash Chain)的结构。
- 前缀复用带来的极大优势:
当新请求到达时,系统只需按照同样的规则计算出前缀的 Hash 值,然后去现有的 KV Cache 池中进行匹配。如果匹配命中(例如多个请求使用了相同的系统提示词),即可直接将这些 Block 的物理引用加入当前请求的 Block Table 中,实现零计算复用。这能跳过共享前缀的 Prefill 计算阶段,显著降低首字延迟(TTFT),尤其适合多轮对话、RAG(检索增强生成)或使用超长系统提示词的场景。 - 哈希链的脆弱性与淘汰机制:
这种机制也有其局限性。只要提示词中的某一个历史 Token 发生改变,后续整条链路的 Hash 都会随之改变。因为后续节点的有效性严格依赖于前置节点。此外,为了维持这套缓存系统,当 KV Cache 内存池被占满但又需要分配新 Block 时,vLLM 会基于最近最少使用(Least Recently Used, LRU) 等规则,淘汰那些当前未被任何请求引用且使用时间最久的历史 Block,为新请求腾出空间。
总结
我对 vLLM 也只是一知半解,还有很多东西要学。vLLM 自从融资之后,版本发布非常频繁,半年干了之前几年的活。导致我们也要不停地学习,先这么记录一下吧,之后可能还会写 vLLM 的其他博客。







