推理与 KV Cache 01：从一次生成到显存容量

KV Cache 保存各层历史 Token 的 K/V，让后续生成复用已有状态；缓存容量随层数、KV Head 数、精度与序列长度增长。本文从隐藏状态和 Q/K/V 串起 Prefill、Decode，再用 GQA 算例计算容量。讨论范围为标准因果 Decoder-only Transformer，MLA、滑动窗口和混合模型需按各自结构计算。

输入一句“为什么电灯泡通电会发光”，模型先读完输入，再逐个输出 Token。对系统工程师来说，值得追踪的是：每一步新算了什么，哪些结果会留下来，下一步又读取哪些数据。

从 Token 到隐藏状态

Tokenizer 把文本转换成 Token ID。Token 与汉字、单词并非一一对应。模型将 ID 映射为向量，再让向量依次经过多层 Transformer。

隐藏状态 X，就是模型在某一计算位置上对 Token 的数值表示。 输入端的向量主要来自词嵌入；经过前面的层后，它还携带已经聚合的上下文信息。同一个 Token 在不同句子、不同层中的表示可以不同。

设某层实际送进 Q/K/V 投影的输入为 $X$ 。采用行向量约定，省略归一化、偏置和位置编码等细节：

Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V

$W_Q,W_K,W_V$ 是训练得到的权重；Q/K/V 是这一次输入产生的中间结果。Q 提供当前查询，K 提供匹配特征，V 提供参与聚合的信息。注意力先计算 Q 与 K 的匹配程度，再据此聚合 V。

图 1：Q/K/V 由隐藏状态经过可训练投影产生。Q 参与当前计算；K/V 形成可供后续 Token 读取的历史状态。

层数与 Token 位置是两个独立维度。例如，第 3 层、位置 10 的 Q，读取的是第 3 层可见位置的 K/V。第 2 层的计算结果会先形成隐藏状态，再进入第 3 层，产生这一层自己的 Q/K/V。

因此，可以先把缓存想成 KVCache[layer][position]。这是帮助理解的逻辑索引，实际内存布局留到第三篇。

Prefill 与 Decode 的交界

假设输入已经分成 4 个 Token，记为 $p_0$ 到 $p_3$ 。先忽略前缀缓存命中、推测解码和特殊调度。

阶段	本次送入模型的 Token	本次新增的 KV	得到的输出
Prefill	$p_0,p_1,p_2,p_3$	四个输入位置在各层的 KV	第一个输出 Token $y_0$
第一次 Decode	$y_0$	$y_0$ 在各层的 KV	$y_1$
第二次 Decode	$y_1$	$y_1$ 在各层的 KV	$y_2$

Prefill 逐层处理整段输入。同一层的多个输入位置可以并行计算，但因果掩码使每个位置只能看到自身及之前的 Token。最后一个输入位置的最终隐藏状态经过输出映射和采样，得到 $y_0$ 。

Decode 处理一个新位置时，每层生成新 Q/K/V，把当前 K/V 纳入可见集合，再用 Q 读取该层的历史与当前 K/V。具体内核可以融合这些操作。

TTFT 是请求开始到收到首个输出 Token 的时间，通常包含排队、输入处理、Prefill 和首 Token 输出等开销。TPOT 描述后续输出 Token 的平均耗时；常见统计方式是将首 Token 之后的生成时间除以后续 Token 数量。比较系统时需要对齐测量口径。

历史 KV 为什么还能继续使用

对固定模型和固定输入前缀，历史位置在因果注意力中看不到后来追加的 Token。后续内容的到来不会改变历史位置已经计算出的状态，因此这些 K/V 可以继续复用。

未来 Token 会生成自己的 Q。过去的 Q 完成当时的查询后，通常不再被未来步骤读取；过去的 K/V 则继续作为查询对象保留。这解释了标准自回归推理为何主要持久缓存 K/V。

缓存省去了历史位置的重复计算。全注意力 Decode 仍要读取可见的历史 K/V；长上下文带来的访问量并没有消失。Hugging Face 的缓存说明给出了按层追加与复用的实现示例。

用 GQA 理解容量里的 Head 数

设模型隐藏维度为 4096，Query Heads 为 32，Head Dimension 为 128。Q 投影的输出宽度为 $32\times128=4096$ 。若模型使用 8 个 KV Heads，K 和 V 的输出宽度分别为 $8\times128=1024$ 。

每 4 个 Query Heads 共享一个 KV Head。这种结构称为 GQA。以下比较保持层数、Head Dimension 和精度相同：

结构	Query Heads	KV Heads	相对于 MHA 的 KV 容量
MHA	32	32	1
GQA	32	8	1/4
MQA	32	1	1/32

这些 Head 数由模型架构和权重形状确定，不能在部署时任意改成另一个数。隐藏维度与 Q 投影宽度相等是这里的模型设定，也不应推广为所有架构的硬性等式。

从一个 Token 算到整个请求

先取所有层结构一致、K/V 同维度同精度的模型。每个 Token 在所有层的 KV 有效数据量为：

s_{\mathrm{token}}=2LH_{kv}db

$L$ 是层数， $H_{kv}$ 是 KV Head 数， $d$ 是每个 Head 的维度， $b$ 是每个元素的字节数；系数 2 对应 K 和 V。

教学算例取 $L=40$ 、 $H_{kv}=8$ 、 $d=128$ 、BF16 的 $b=2$ ：

s_{\mathrm{token}}=163840\ \mathrm{B}=160\ \mathrm{KiB}

一个已经处理 8192 个 Token 的序列，需要 $8192\times160\ \mathrm{KiB}=1.25\ \mathrm{GiB}$ 的 KV Payload。10 条同长度、没有共享的序列合计 12.5 GiB。这里 $1\ \mathrm{KiB}=1024\ \mathrm{B}$ ， $1\ \mathrm{GiB}=2^{30}\ \mathrm{B}$ 。

s_{\mathrm{token}}=\sum_l H_{kv,l}\left(d_{k,l}b_{k,l}+d_{v,l}b_{v,l}\right)

下一步要解决动态分配：序列不断增长，请求长短不同，部分前缀还可以共享。把刚才的连续容量拆成可管理的 Block，就进入了 Paged KV Cache。

下一篇：AI Infra 学习 02：KV Cache 如何分页、共享与回收。