1. 为什么KV Cache是大模型推理的命门
第一次认真翻看推理框架源码的人,大概率会在attention计算那块卡住——明明训练时一个完整的前向传播写得清清楚楚,怎么到了推理阶段,代码里凭空多出来一个叫past_key_values的东西,还越滚越大?这个past_key_values就是KV Cache,它不是什么锦上添花的优化技巧,而是自回归生成能跑得起来的前提。
先把场景摆清楚。大模型生成文本是逐token往外蹦的,每蹦一个token,都要拿当前这个token的向量去和前面所有token做注意力计算。假设上下文已经有1000个token,现在要生成第1001个,注意力机制需要这1001个位置各自的Key和Value。如果不用缓存,那就得把前1000个token重新过一遍模型,把它们的K、V重新算出来——注意,是每一轮生成都重算一遍。生成1000个token,就要做1000次全量前向,计算量是O(n²)级别的重复劳动。
KV Cache干的事情非常朴素:把每一层、每一个注意力头在之前所有位置算出来的Key和Value张量存下来,生成新token时只算当前这个token的Q、K、V,然后把新的K、V追加到缓存里,注意力计算直接用"当前Q"去查"历史K、V + 当前K、V"。这样每生成一个token,计算量从"重算全部历史"降到"只算当前一步",复杂度从O(n²)的重复计算变成O(n)的增量计算。
一句话概括:KV Cache是用显存换计算。它把推理从"每次都要重新读完整本书"变成"读过的内容记在便签上,只读新的一页"。
但这里有个容易被忽略的点:KV Cache缓存的是Key和Value,不是Query。原因是Query只跟"当前要生成的位置"有关,历史位置的Query在后续生成中再也用不到了;而Key和Value是"被查询的对象",历史位置的K、V会被未来的每一个新token查询,所以必须留着。这个不对称性是理解KV Cache的第一道门槛,很多人第一次看代码时想不通"为什么只缓存K和V",想通这一点后面就顺了。
从工程角度看,KV Cache带来的显存开销是实打实的。以常见的配置估算:层数L、注意力头数H、头维度d、序列长度S、批次B、精度字节数b,KV Cache的显存占用约为2 × L × H × d × S × B × b。拿一个70亿参数、32层、32头、头维度128的模型举例,单条4096长度的序列,FP16精度下KV Cache大约是2 × 32 × 32 × 128 × 4096 × 2字节,算下来约2GB。注意这只是单条序列,如果并发32路,直接飙到64GB——比模型权重本身还吃显存。这就是为什么长上下文和高并发场景下,KV Cache的管理策略直接决定了推理服务的吞吐上限。
理解了"为什么需要"和"代价是什么",后面聊的所有优化手段——PagedAttention、GQA、量化、滑窗——本质上都是在回答同一个问题:怎么在保住KV Cache收益的前提下,把它的显存和带宽开销压下去。
2. KV Cache在注意力计算里的真实数据流
光知道"缓存K和V"还不够,得把数据流走一遍,才能明白后面那些优化到底在动哪块。
2.1 从单头注意力到缓存张量的形状
标准多头注意力里,输入hidden state经过三个线性投影得到Q、K、V,形状都是[batch, num_heads, seq_len, head_dim]。训练时seq_len是整个序列,一次性算完。推理时是自回归的,第一次prefill阶段输入整个prompt,seq_len等于prompt长度;之后的decode阶段每次只输入一个token,seq_len等于1。
KV Cache要存的就是每一层、每个decode step累积下来的K和V。实际实现里,缓存张量通常预分配成[batch, num_heads, max_seq_len, head_dim],然后用一个seq_len指针记录当前已经填到哪了。每来一个新token,就把新算出的K、V写到seq_len位置,指针加一。这样避免了每步都重新分配内存,是工程上的常规做法。
这里有个细节值得说:prefill阶段和decode阶段的注意力计算逻辑其实不一样。prefill时Q、K、V的seq_len都等于prompt长度,是一次标准的全注意力;decode时Q的seq_len是1,K、V的seq_len是"历史长度+1",是一个"单query对多key"的注意力。很多框架会把这两个阶段分开写kernel,就是因为形状差异太大,混在一起写效率上不去。
2.2 缓存命中时注意力到底怎么算
decode阶段拿到当前token的Q(形状[batch, num_heads, 1, head_dim]),和缓存里的K(形状[batch, num_heads, S, head_dim])做矩阵乘,得到注意力分数[batch, num_heads, 1, S],softmax之后再去乘缓存的V,得到输出[batch, num_heads, 1, head_dim]。整个过程里,历史K、V只读不写,当前K、V写一次。
这个"只读历史、只写当前"的模式,决定了KV Cache的性能瓶颈不在计算,而在显存带宽。decode阶段每生成一个token,都要把整个KV Cache从显存读一遍——序列越长,读的数据量越大,而计算量几乎不变。这就是所谓的"memory-bound":GPU的算力大量闲置,时间都花在搬数据上了。理解了这一点,就能明白为什么量化KV Cache(把FP16压成INT8)能提速——不是因为算得快了,而是因为读的数据少了。
2.3 一个具体的形状推演
假设模型配置是:32层,32个注意力头,头维度128,FP16精度,batch=1,序列长度2048。单层单头的KV Cache是2 × 2048 × 128 × 2字节 = 1MB,32头就是32MB,32层就是约1GB。生成第2049个token时,需要读取这1GB数据。如果显存带宽是2TB/s,光读KV Cache就要0.5毫秒,而实际计算可能只要几十微秒。这个比例直观说明了为什么长上下文推理这么吃带宽。
实操提示:如果你在做推理性能分析,先算一下KV Cache的显存占用和每步的读取量,再对比GPU带宽,基本就能判断出瓶颈在哪。这个估算比盲目上profiler快得多。
3. PagedAttention如何解决显存碎片
KV Cache的显存管理在早期实现里非常粗暴:给每条序列预分配一块连续的最大长度显存。问题是,实际生成的序列长度是动态的,预分配4096但只用了500,剩下3500就浪费了;而且不同序列长度不一,连续分配容易产生碎片,显存利用率经常只有20%到40%。PagedAttention就是冲着这个浪费来的。
3.1 操作系统虚拟内存的类比
PagedAttention的核心思想直接借用了操作系统的虚拟内存分页机制。它把KV Cache切成固定大小的block(比如16个token一块),每条序列的KV Cache不再要求连续存放,而是由一组block拼起来,再用一张block table记录"逻辑位置到物理block"的映射。生成新token时,如果当前block没满就追加,满了就从显存池里再拿一块。
这个设计的好处立竿见影:block大小固定,显存池可以预先切好,分配和回收都是O(1);序列之间可以共享block(比如相同的prompt前缀),进一步省显存;碎片问题基本消失,因为所有block等大,不存在"这块剩300、那块剩500"的尴尬。
3.2 block table与注意力kernel的配合
有人会问:KV Cache不连续了,注意力计算怎么找到对应的K、V?答案是block table在kernel里被当作索引表用。计算注意力时,kernel根据block table把逻辑上连续的token映射到物理上分散的block,逐个block读取K、V。虽然读取地址不连续,但因为是按block批量读,访存效率依然可以接受。
实测数据上,PagedAttention在vLLM里的表现是把显存利用率从传统实现的百分之三四十提到百分之九十以上,吞吐量提升2到4倍。这个提升在长上下文、高并发场景下尤其明显,因为那时候碎片浪费最严重。
3.3 前缀共享带来的额外收益
PagedAttention还有个副产品:如果多条请求有相同的系统提示词或相同的few-shot示例,它们的KV Cache前缀可以指向同一批block,物理上只存一份。这在"同一个系统提示+不同用户问题"的对话服务里非常常见,能省下可观的显存。实现上靠的是block的引用计数,引用归零才回收。
注意:前缀共享要求前缀完全一致,包括tokenize后的结果。如果你的系统提示词里有动态内容(比如时间戳),共享就失效了。做服务优化时,把系统提示词里的动态部分挪到用户消息里,能显著提高共享率。
4. 分组查询注意力与KV Cache的显存账
前面聊的都是"怎么管"KV Cache,GQA(分组查询注意力)解决的是另一个维度的问题:从源头减少KV Cache的"头数"。
4.1 MHA、MQA、GQA的取舍
标准多头注意力(MHA)里,Q、K、V的头数相同,比如都是32头。KV Cache要存32头的K和V。多查询注意力(MQA)把K、V压成1头,所有Q头共享同一组K、V,KV Cache直接降到1/32,但效果会掉。GQA是折中:把32个Q头分成8组,每组4个Q头共享一组K、V,KV Cache降到1/4,效果损失很小。
这个设计的逻辑很直白:Q头负责"从不同角度提问",K、V负责"提供被查询的信息"。多个Q头共享K、V,相当于多个提问者查同一份资料,信息量够用,但省了存储。实践中GQA已经成为主流大模型的标配,Llama 2 70B、Llama 3、Qwen系列都用了GQA。
4.2 显存账怎么算
还是用之前的配置:32层、32头、头维度128、FP16、序列2048。MHA下KV Cache约1GB。换成GQA,8组KV头,KV Cache变成约256MB。如果并发32路,MHA要32GB,GQA只要8GB——这直接决定了单卡能扛多少并发。
| 注意力类型 | KV头数 | KV Cache相对大小 | 效果损失 |
|---|---|---|---|
| MHA | 32 | 1x | 无 |
| GQA(8组) | 8 | 1/4 | 很小 |
| MQA | 1 | 1/32 | 较明显 |
4.3 和PagedAttention的叠加效果
GQA减少的是"每个token的KV数据量",PagedAttention减少的是"分配浪费"。两者叠加,显存效率提升是乘性的。一个用GQA的模型配上PagedAttention,相比MHA配连续分配,显存利用率能差出一个数量级。这也是为什么现在的推理框架基本都同时支持这两项。
实操心得:选模型时,如果推理是你的主要场景,优先选带GQA的版本。同样参数量下,GQA版本在长上下文和高并发下的吞吐优势非常明显,而效果差距在多数任务上肉眼难辨。
5. 长上下文场景下KV Cache的压缩与淘汰
上下文长度从4K涨到128K甚至更长,KV Cache的显存需求是线性增长的,但显存是固定的。这时候就得做取舍:要么压缩,要么淘汰。
5.1 KV Cache量化
最直接的压缩是降精度。FP16的KV Cache占2字节,量化到INT8占1字节,显存直接减半;再激进一点到INT4,减到1/4。量化会带来精度损失,但KV Cache对量化的容忍度比模型权重高——因为注意力分数经过softmax后,对K、V的微小扰动不那么敏感。实测中INT8量化KV Cache在多数任务上几乎无损,INT4在长上下文任务上会有可感知的下降。
实现上,量化KV Cache需要在写入时做量化、读取时做反量化,这会增加一点计算开销,但省下的带宽通常能覆盖掉。关键是量化粒度的选择:per-tensor量化简单但精度差,per-channel或per-token量化精度好但实现复杂。主流框架一般提供per-head或per-token的量化选项。
5.2 滑窗注意力与稀疏化
另一条路是"不存那么多"。滑窗注意力只保留最近W个token的KV,更早的直接丢弃。这对局部依赖强的任务(比如代码补全、流式对话)效果不错,但对需要长距离回忆的任务(比如长文档问答)会丢信息。
更精细的做法是稀疏化:根据注意力分数判断哪些历史token重要,只保留重要的。比如H2O这类方法,维护一个"重击者"集合,注意力分数高的token的KV长期保留,低的淘汰。这类方法在效果和显存之间找平衡,但实现复杂度高,调度开销也不小。
5.3 淘汰策略的坑
淘汰KV Cache最怕的是"淘汰了后面又需要"。比如对话里用户前面提了一个关键约束,中间聊了很多无关内容,最后又回到那个约束——如果中间把约束的KV淘汰了,模型就"忘"了。所以淘汰策略要么保守(保留比例高,省不了多少),要么配合"重要token识别"(比如系统提示词、用户明确强调的内容强制保留)。
实操提示:做长上下文服务时,别一上来就激进淘汰。先量化到INT8,这步几乎无损;还不够再考虑滑窗,滑窗窗口别小于任务的最长依赖距离。淘汰策略一定要在你的真实任务分布上验证,别只看通用benchmark。
6. 推理框架里KV Cache的配置与调优
理论聊完,落到实操。现在主流推理框架对KV Cache都有比较成熟的配置项,调对了能明显改善吞吐和延迟。
6.1 关键配置项
以常见的推理服务配置为例,和KV Cache相关的参数主要有几个:max_model_len决定预分配的最大序列长度,直接影响显存占用;gpu_memory_utilization控制框架能用多少显存,留出余量给KV Cache;block_size是PagedAttention的block大小,一般16或32;kv_cache_dtype控制KV Cache精度,可选auto、fp8、int8等;max_num_seqs限制并发序列数,间接控制KV Cache总量。
这些参数不是独立的。max_model_len调大,单序列KV Cache变大,能并发的序列就少;gpu_memory_utilization调高,留给KV Cache的空间多,但可能触发OOM。调优的本质是在"单序列长度""并发数""显存安全边际"三者之间找平衡。
6.2 一个调优实例
假设单卡80GB显存,模型权重占40GB,剩下40GB给KV Cache和激活值。如果max_model_len设8192,单序列KV Cache(GQA,8组KV头,32层,头维度128,FP16)约2 × 32 × 8 × 128 × 8192 × 2 = 1GB。40GB理论上能放40条,但激活值和其他开销要占一部分,实际可能设max_num_seqs=32比较稳。
如果业务上大部分请求长度不超过2048,那把max_model_len降到4096,单序列KV Cache减半,并发能翻倍。这就是"按业务实际长度配置"的价值——别为了偶尔的超长请求,让所有请求都按最大长度预留。
6.3 监控与排障
上线后要盯几个指标:KV Cache使用率(用了多少block)、请求排队时间(并发满了会排队)、preemption次数(显存不够时框架会抢占序列,把KV Cache换出,这个次数高说明配置偏紧)。如果preemption频繁,要么降max_num_seqs,要么开KV Cache量化,要么加卡。
踩坑记录:我见过有人把
gpu_memory_utilization设到0.98,结果模型加载后KV Cache空间不够,服务起来就OOM。留5%到10%的余量是常规做法,别把显存榨干。
7. 关于KV Cache的几个常见误解
最后澄清几个我在交流中反复遇到的理解偏差,这些点想清楚了,对推理优化的判断会准很多。
第一个误解是"KV Cache越大越好"。KV Cache本身不产生价值,它只是让生成不用重算。缓存越大意味着序列越长或并发越高,但如果业务不需要长上下文,盲目调大max_model_len只是浪费显存、降低并发。配置要跟着业务走。
第二个误解是"KV Cache能跨请求复用"。默认情况下不能,每条请求的KV Cache是独立的。只有前缀完全一致时,PagedAttention的前缀共享才能复用。所以别指望"上一个用户问过类似问题,这个用户就能加速"——除非前缀真的相同。
第三个误解是"量化KV Cache一定掉效果"。INT8量化在绝大多数任务上效果损失在噪声范围内,真正需要担心的是INT4及以下。而且KV Cache量化和权重量化是两回事,可以独立选择。实践中"权重FP16 + KV Cache INT8"是很常见的组合。
第四个误解是"decode阶段慢是因为算力不够"。前面算过,decode是memory-bound,瓶颈在显存带宽。换更强的GPU如果带宽没提升,decode速度可能纹丝不动。优化decode要从减少数据搬运入手——量化、GQA、PagedAttention都是这个思路。
把这些想明白,再看推理框架里那些参数和优化选项,就不会觉得是一堆玄学开关,而是每个都有明确的物理意义和适用边界。KV Cache这个看似简单的缓存机制,撑起了整个大模型推理的性能天花板,值得花时间吃透。