KV缓存:大模型首token慢、后续token快的工程真相

📅 发布时间:2026/10/20 23:14:47
KV缓存:大模型首token慢、后续token快的工程真相
如果你用过ChatGPT这类生成式AI大概率有过这种体验把问题敲进去界面先转个一两秒第一个字拖半天可一旦它开口后面的文字就像提前背好了稿子一样往外冒。很多人的直觉是——模型太大了算力不够。这个解释只说对了最表面的一层。真正让首token慢、后续token快的是Transformer自注意力机制背后的一个工程级环节KV缓存。想弄明白这件事得先把大模型的吐字节奏拆开看。大模型不是一次性把整段回答写好的它是一个token一个token往外蹦的。所谓token可以理解为文字切片英文单词大致一两个token一个中文则经常一个字对应一个到两个token。每生成一个新token模型都要把之前所有内容当作上下文再预测下一个最可能的token。这个重复回看的动作就是性能差异的根源。1. 大模型吐字这件事为什么是两段式节奏1.1 开场迟钝的真相prefill阶段在消化整段上下文当用户把问题发给大模型模型最先经历的阶段叫prefill预填充也叫上下文处理阶段。这个阶段里模型要把你输入的整个prompt所有token做一次完整前向计算从词嵌入映射到每一层Transformer的注意力运算再到前馈网络最后产出一个概率分布挑出第一个token。prefill阶段最大的特点是并行。GPU可以同时处理prompt里所有token的矩阵乘法把整段上下文一次性读完。但并行不等于免费你输入的文字越多需要算的矩阵就越大。同样是问一句话问今天天气怎么样和上传一篇几千字的材料前者首token延迟可能只有几十毫秒后者却是几秒起步。所以你会发现一个现象短问题几乎秒回长文档问答则卡顿明显——这个卡顿发生在第一个字出来之前正是prefill在处理你的全部输入。1.2 后半程丝滑的原因decode阶段只推最新token第一个token出来后模型进入decode解码阶段。这个阶段每生成一个新token都要基于之前已经生成的全部内容来决定下一个词。听起来像每走一步都要把所有历史内容重读一遍但工程上并没有这么笨。decode阶段的关键特征是模型权重不会变已经确认的历史token也不会变。唯一需要处理的是新token如何与全部历史进行注意力交互。于是推理框架会把历史token的某些中间产物提前存起来每走一步只需要做一次轻量级计算——把新token和已存好的历史中间结果做交互。生成速度因此从重新处理全部历史降为只处理最新token后者在工程上几乎和上下文长度解耦。一个直观类比prefill像是通读所有参考文献并做批注耗时耗力decode像是在批注之上逐字续写每次只需要查最近的一条批注就能接着写。KV缓存就是你提前写好的批注卡片。1.3 算力太重只是表象真正的瓶颈是显存带宽和缓存命中很多人把大模型推理慢归结为模型太大、算力太重。这个说法不能说错但它解释不了第一个token慢但后面快的现象。同一个模型同样参数量为什么开头要卡一下之后却一路丝滑如果纯粹是算力瓶颈每个token的耗时应该差不多才对。真正的区别在于prefill阶段做的是大规模并行矩阵乘法考验的是FLOPS每秒浮点运算次数属于算力密集decode阶段每一步只生成一个token要频繁把模型权重和KV缓存从显存搬到计算单元属于访存密集。对大模型推理来说绝大部分时间其实消耗在带宽上而不是算力上。每生成一个token理论上前馈网络都要把上百GB的权重过一遍这比计算本身更耗时。所以准确说法是首token慢主要是prefill阶段处理整个prompt的工程量积累后续丝滑是因为decode阶段借助KV缓存把每步重算量降到了最低。把锅全部甩给模型算力你就错过了真正可以做优化的地方——显存、缓存、内存布局。2. KV缓存到底存了什么拆开自注意力看重复劳动2.1 自注意力就是一场全员互查要理解KV缓存得先回到Transformer的核心自注意力Self-Attention。自注意力的直观含义是序列里的每个token都要和序列里所有token做一次相关性计算然后根据相关性高低去融合信息。具体一点模型把每个token拆成三个向量。Query可以理解为我要找什么Key是我被什么特征标记Value是我能提供什么内容。当前token拿自己的Query去依次和所有token包括自己的Key做点积得到一系列注意力分数再用softmax归一化最后用这些分数对Value做加权求和得到当前token的新表示。也就是说当模型生成第10个token时它心里其实把前9个token挨个看了一遍才知道应该接什么。问题在于前9个token的内容从第一步起就没有变过它们对应的Key和Value按理说也不需要变。但如果不做任何存储模型每走一步都会把前面所有token的K、V从头再算一遍——这就是重复劳动。2.2 没有缓存解码会陷入O(N²)泥潭我们算一笔账。假设当前已经生成到第t个token如果不缓存历史K、V模型必须把序列长度为t的所有token重新做一遍自注意力这一步的计算量约为O(t²)。从t1一直生成到N个token总计算量是ΣO(t²)约等于O(N³)。大模型日常动辄生成几百上千个token如果每步都从头重算早就不可用了。KV缓存出现后每次decode只需要为最新token算一次K和V历史token的K、V直接复用。单步计算量从重算整个序列降为只新增一个位置的交互。有人会问注意力分数不还是要和全部历史Key做点积吗没错这部分复杂度仍然随序列长度线性增长。但相比把每个历史token的K、V都重算一遍已经是天壤之别。这也是为什么在推理框架里打开KV缓存和关闭KV缓存解码速度可以差出一个数量级。试着把开源模型的生成参数里use_cache设为False很快就能感受到什么叫一步一卡。KV缓存不是锦上添花而是大模型能流畅对话的基石。2.3 KV缓存的工作原理把K和V变成可重复使用的便签KV缓存的实现逻辑其实不复杂。模型在prefill阶段处理完整prompt时会把每一层Transformer算出的Key张量和Value张量存进显存里的两块区域。到了decode阶段每生成一个新token模型会先算出这个token的K、V把K、V追加进缓存然后拿这个新token的Query与缓存里所有历史Key做点积得到注意力分数再对缓存里所有Value做加权求和。一个标准推理循环大致是这样输入promptprefill算出第一个token并把每一层的K、V存入缓存把第一个token作为输入结合缓存里的历史K、V算出第二个token同时把第二个token的K、V追加进缓存重复第2步直到遇到结束符或达到最大长度。这个过程中KV缓存的大小只和上下文长度与模型结构有关不随生成轮次数减少——它是只增不减的。所以当你把max_new_tokens设得很大或者上下文窗口很长显存里的KV缓存会像滚雪球一样越滚越大。这也是长文本生成时显存容易爆掉的直接原因。3. 工程落地课显存里那一大块KV缓存怎么算、怎么管3.1 用公式估算KV缓存体积KV缓存的体积可以用公式直接估算。假设模型有L层Transformer每层注意力头数为H每个头维度为DKV缓存用bf16存储每个数占2字节缓存总大小是存储字节数 2(K和V两份) × L × H × D × seq_len × 2(字节)注意H×D通常等于模型的hidden_size。以常见的7B模型为例层数L32hidden_size4096seq_len2048那么2 × 32 × 4096 × 2048 × 2 ≈ 1.07GB也就是说在2048上下文下光KV缓存就要吃1GB显存。如果把上下文拉长到32768这个数字直接变成约17GB——已经是7B模型bf16权重约14GB的1.2倍还多。当KV缓存逼近甚至超过权重体积时它就成了决定能不能跑起来的主要因素。如果再叠加并发请求压力更大。一个服务部署8路并发每个请求都得维护自己的KV缓存显存消耗就是单路的8倍。很多推理框架在并发一高时容易OOM原因往往不是模型权重放不下而是KV缓存占用了大量算不到的显存。3.2 长上下文为什么贵缓存线性涨注意力平方涨KV缓存随序列长度线性增长这个好理解。但注意力计算还有另一层开销。对于每个新token它要和历史所有token的Key做点积这个计算量随序列长度线性增长而整个生成过程每一次都要扫过所有历史累计开销就是平方级。序列越长每生成一个token就越吃力。这也是长上下文贵的本质。短prompt可能在几十token内就结束长文档问答动辄上万token。在同样流量下长上下文的单个请求会在decode阶段造成远高于短请求的注意力计算量。所以上下文长度翻倍成本不止翻倍不是一句口号它是KV缓存和注意力耗能双重叠加的结果。从产品角度看支持长上下文更像是一种能力储备。真正高频使用超长上下文的用户占比不高但每个长请求消耗的资源可能是普通请求的几十倍。推理系统如果不在缓存管理层做优化长上下文很容易把整台服务拖垮。3.3 首token慢的隐藏元凶之一解码前的缓存预分配回到最初的问题第一个token为什么慢。除了prefill计算量大还有一个工程细节常被忽略——KV缓存的预分配。KV缓存是动态增长的。如果每生成一个token就临时申请一段显存会产生大量内存碎片和分配开销。主流推理框架通常会在prefill结束、decode开始前根据max_new_tokens预先把整段缓存空间锁定。这意味着第一个token生成时系统可能已经为未来数千甚至上万token的KV缓存做了最大规模的显存预留和初始化。预留本身也有成本。在PagedAttention这类方案出现之前显存分配还要考虑连续性找一块足够大的连续显存可能触发额外整理或切换。所以你会发现即使模型不大、prompt很短输入后系统也经常卡一下——那未必是模型在思考而是推理引擎在为后续生成准备草稿纸。理解了这一点看首token延迟时就不会只盯着参数量了。3.4 应对生产压力的四大招GQA、PagedAttention、前缀共享、投机解码第一招GQA分组查询注意力。既然KV缓存大头是K和V能不能让多个Query头共享同一组K、VGQA就是这么设计的。把原来的多个Query头分成若干组每组用一套共享的K、VKV缓存体积直接除以组数。现在很多新发布的模型默认采用GQA其中一个重要原因就是省显存、降带宽。第二招PagedAttention。把KV缓存按固定大小的页block管理类似操作系统的分页机制。需要时才分配下一页不必一次性找一大块连续内存大幅降低碎片率并发吞吐常常成倍提升。我在实际部署中把推理框架从静态缓存切到支持PagedAttention的方案后同样的GPU能多塞好几路并发收益非常直观。第三招前缀共享。聊天系统的系统提示词通常很长且固定。如果多个请求共享这段前缀前缀部分的KV缓存也可以共享。类似系统提示词多轮历史的公共前缀只需算一次、存一份后续请求直接复用。这也是很多推理引擎专门做前缀缓存prefix caching功能的原因。第四招投机解码。decode阶段每步串行能不能先让一个小模型快速猜几个token再由大模型一次性验证投机解码就是这么干的小额草案由轻量模型产生大模型用一次前向批量确认或修正因为KV缓存参与验证批量确认的开销远小于逐个串行。这个方案能显著提升单请求吞吐代价是需要额外维护一个小模型适合对延迟极度敏感的场景。这些优化有一个共同点都没有增加模型一层算力而是从减少重复计算、减少显存冗余下手。理解了KV缓存就理解了这些框架为什么要花那么多心思去折腾内存管理。4. 亲手测一测用代码感受prefill与decode的差距4.1 一个简单的测量实验设计想直观感受prefill和decode的差距不需要几十B的模型。用几B参数的开源模型就够。核心思路是分别记录首token耗时和后续每个token的耗时。以常见的transformers库为例大致逻辑是这样import time from transformers import AutoTokenizer, AutoModelForCausalLM model_name 某开源因果语言模型路径 tokenizer AutoTokenizer.from_pretrained(model_name) model AutoModelForCausalLM.from_pretrained(model_name, device_mapauto) prompt 请详细解释一下什么是注意力机制并举三个例子。 inputs tokenizer(prompt, return_tensorspt).to(cuda) # 先把模型预热 _ model.generate(**inputs, max_new_tokens4, use_cacheTrue) # 用一个较短的prompt测首token耗时 start time.time() outputs model.generate(**inputs, max_new_tokens32, use_cacheTrue) first_token_time time.time() - start严格来说上面的计时包含整个生成过程并不精确。更精确的做法是用流式生成接口拿到第一个token的时间戳gen model.generate(**inputs, max_new_tokens128, streamTrue, use_cacheTrue) start time.time() first_token next(gen) ttft time.time() - start多测几次取平均就能得到比较稳定的参考值。然后做两组对比第一组是固定prompt长度看首token耗时和后续每token耗时第二组是不断加长prompt看首token耗时如何变化。实测下来的典型结果是prompt从几十token涨到几千token首token耗时可能从几十毫秒涨到几秒而后续每token耗时变化则小得多。这个现象正好对应prefill处理prompt与decode依赖KV缓存的差异。还有一个巧办法观察KV缓存的作用同样的prompt分别用use_cacheTrue和use_cacheFalse生成再对比相同token数量下的耗时。打开缓存时速度远快于关闭缓存差距会随着上下文长度增加而越拉越大。4.2 从现象看本质你的显存正在被谁吃掉实验过程中可以同步监控显存。生成几百个token后打开nvidia-smi观察显存曲线你会发现显存占用在生成过程中稳步抬升而不是一开始就恒定。抬升的这部分几乎就是KV缓存一个token一个token累积出来的。如果把上下文长度从2K改成8K相同模型显存占用会明显上涨。而如果改用支持PagedAttention的推理框架显存分配会呈现按块增长的阶梯状而不是一次性大块预留。对照这些现象你就能明白为什么总有人抱怨上下文长度开大了就OOM——模型权重是固定的真正持续膨胀的是KV缓存。从运维角度看这类监控特别重要。如果服务上线后显存占用率缓慢爬升然后触发OOM第一反应不应该是一味加显存而是先看是不是单请求缓存累积、再查并发数、最后看是否需要给KV缓存做量化或压缩。4.3 基于KV缓存的调参避坑清单写几条我实际踩过的坑。第一别随便关use_cache。调试或测试时可能会手滑关闭缓存生成速度立刻断崖式下降。如果发现模型又慢又卡第一件事就是检查缓存是不是被关了。第二max_new_tokens不是越大越好。它直接影响KV缓存预分配规模。请求频繁的场景下把max_new_tokens精确调到业务需要的上限能省下一大块闲置显存。很多框架预分配会按这个上限来宁小勿大。第三长上下文要提前做显存估算。参考前面的公式把KV缓存量、模型权重、激活值内存加一起再定并发数。别等线上OOM了再去减并发那会儿通常已经被用户投诉一轮了。第四看延迟不能只盯着首token时间。首token快不代表整体快。要提升实际用户体验还要关注每秒输出token数decode吞吐和并发下的饱和度。两个指标是一对跷跷板要一起看否则很容易被某个指标误导。5. 顺着KV缓存往下走现代推理框架还在优化什么5.1 预填充和解码为什么要分开调度prefill是计算密集型长矩阵乘法把GPU算力跑满访存量相对可控decode是访存密集型每生成一个token都要把大权重和长KV缓存搬一遍算力反而闲置。这两种计算特性完全不同硬放在同一套调度逻辑里效率必然被拖累。所以很多推理系统会把prefill和decode分开处理。比如用同一GPU的不同计算流或者把prefill请求集中到高算力设备decode请求交给高带宽设备。即使在同一设备上也会在请求调度时把正在prefill的长请求和正在decode的短请求错开让两种计算互补尽量不互相阻塞。从KV缓存的角度看decode阶段的访存开销和KV缓存长度强相关prefill阶段则需要为第一个token建立并分配KV缓存。拆开调度之后两边都可以针对自己的瓶颈做优化。这也是现代推理框架吞吐高出经典实现一大截的重要原因。5.2 KV缓存的量化、压缩与复用KV缓存虽然省下了重复计算但自己却成了显存和带宽的大户。于是工程师们又把矛头指向了它。一是量化。把KV缓存从bf16或fp16压到int8甚至int4。由于K、V经过注意力计算后很多数值对精度不太敏感量化后缓存的显存占用直接减半甚至减到四分之一带宽压力也同步降低。代价是精度略有损失需要在实际任务里评估影响。二是压缩。有些方案会丢弃一部分不那么重要的历史K、V只保留对后续预测影响大的token。典型做法是滑动窗口保留最近若干token或者根据注意力分数选择top-k的token作为缓存。这相当于给KV缓存做剪枝。做得好可以极大延长有效上下文做不好则容易丢关键信息属于高收益高风险。三是复用。多轮对话里系统提示词和历史对话的前缀是重复的。如果框架支持前缀缓存不同请求可以共享同一段KV新请求只需要计算新增的那部分。在多用户、多轮产品的场景里复用带来的提升通常是数倍级。这些优化再次印证了一件事大模型推理的战场早就从算法精度转移到了内存与带宽的经营。谁能把KV缓存管理得更精细谁就能在同样的显卡上服务更多请求、给用户更快的响应。说个我个人的体会。一次压测里我盯着显存曲线看了一下午才真正理解KV缓存对大模型推理意味着什么——它不是某一个API的优化选项而是整个推理系统的地基。常规讨论总是习惯用参数量大来解释延迟可把KV缓存那部分算清楚之后很多反直觉的现象都能解释通了。如果你也被首token慢这类问题困扰建议别急着换更大显卡先把模型结构和缓存策略摸一遍说不定一张老卡也能把吞吐撑上去。这就是我踩过坑之后最想分享的结论。