跳过主要内容
CookLLM logoCookLLM 文档
系统工程推理优化

KV Cache

会员专享

decode 每步重算历史 token 的 K/V 是纯浪费,KV cache 把它省掉,代价是显存

在权益中心获取代码

你已经会让 Transformer 做一次完整的前向传播:给定一段 token,一次算出每个位置的输出(见 Transformer 前向传播)。但真正生成文本时,模型不是一次吐出一整段,而是一个 token 一个 token 地挤出来。这个"逐 token"的过程藏着一笔巨大的浪费:每生成一个新 token,朴素实现都要把整段历史重新过一遍模型,而历史 token 的 K 和 V 从第一次算出来那一刻起就再也不会变。

这一章就来消除这份重复劳动。我们要搞清楚四件事:这笔浪费到底有多大、为什么历史的 K/V 一定不变、怎么把它们缓存下来复用、以及这份缓存会吃掉多少显存。KV cache 是几乎所有推理系统的地基,理解它,才能看懂后面 batching、显存管理、量化这些优化在对付什么。

生成是一步一步挤出来的

先回忆自回归生成的循环。给定已有的 token 序列,每一步做四件事:

  1. 把当前整个序列喂进模型,做一次前向传播
  2. 取最后一个位置的 logits
  3. 对词表做 softmax,选出下一个 token(贪心、采样、top-k/top-p 都行)
  4. 把这个 token 拼回序列末尾,回到第 1 步

写成公式,模型每一步都在估计一个条件分布:

P(xt∣x1,x2,…,xt−1)P(x_t \mid x_1, x_2, \ldots, x_{t-1})

注意条件是从 11 到 t−1t-1 的每一个 token,不只是上一个。Transformer 预测下一个词时会注意到全部历史,这正是它能在几千个 token 上保持连贯的原因,也正是每一步都得把整段序列重新过一遍的根源。

这个循环其实分成性质完全不同的两段:

  • Prefill(预填充):把整个 prompt 一次性喂进去,所有位置并行算完,产出第一个新 token。假设 prompt 有 NN 个 token,这一步一次处理 NN 个位置。
  • Decode(解码):之后每一步只新增一个 token,逐个往外挤。

Generation Flow

BentoLM generation uses one prefill pass, then token-by-token decode with KV cache.

1
Prompt tokens
plain text continuation or ChatML prompt
input_ids
2
Prefill
run the full prompt once and build KV cache
past_key_values
3
Decode loop
feed only the newest token at each step
logits
4
Sampling
temperature, top-k, top-p and repetition penalty
next token
5
Stop
finish on max_new_tokens or EOS
output ids

生成 100 个 token,等于 1 次 prefill 加 99 次 decode。prefill 只发生一次,decode 才是循环的主体,也是所有推理优化真正要对付的部分。

登录以继续阅读

这是一篇付费内容,请登录您的账户以访问完整内容。