大模型推理解析：从 Token、QKV 到 GQA 与 KV Cache

当我们向大模型输入“北京是中国的”，模型并不是一次把“首都。”两个词一起吐出来。Decoder-only Transformer 做的是下一个 Token 预测：先根据整个输入预测“首都”，再把“首都”作为新的输入预测“。”，最后根据“。”预测 <EOS> 并停止。

本文用一个简化的三层 Transformer，把这条链路完整拆开。重点不是记住某个模型的具体数字，而是理解每个阶段的张量代表什么、维度为什么能够对齐，以及 GQA 和 KV Cache 为什么能显著改善推理效率。

为了让图和公式保持一致，全文采用下面的示例配置：

符号	示例值	含义
(B)	1	Batch Size
(S)	4	输入序列长度
(H)	4096	Hidden Size
(L)	3	Transformer 层数，仅用于示意
(n_q)	32	Query Head 数量
(n_{kv})	8	Key/Value Head 数量
(d)	128	每个 Head 的维度

并假设文本被分成四个 Token：

t1 = 北京
t2 = 是
t3 = 中国
t4 = 的

真实模型的 Tokenizer 可能采用不同切分方式，聊天模型还会加入 system、user、assistant 等模板 Token，因此真实序列长度通常更长，但计算过程完全相同。

一、总览：从输入到生成“首都。”

图 1：从输入、Embedding、三层 Transformer 到逐 Token Decode 的完整链路。

整条链路可以压缩成：

文本
→ Tokenizer
→ Token IDs
→ Embedding
→ Transformer × L
→ Final RMSNorm
→ LM Head
→ Logits
→ 采样下一个 Token

1. Tokenizer 与 Embedding

Tokenizer 先把文本映射为整数 ID：

“北京是中国的”
→ [北京, 是, 中国, 的]
→ input_ids，形状 [1, 4]

Embedding 不是复杂计算，而是根据 Token ID 从词嵌入表中查出对应行。假设隐藏维度为 4096，那么每个 Token 会变成一个 4096 维向量：

H⁰.shape = [B, S, H] = [1, 4, 4096]

这里的四行分别对应“北京”“是”“中国”“的”。每一行都是该 Token 当前的隐藏表示。

2. 为什么模型中间始终保持 4096 维

三层模型的主干状态依次是：

H⁰ [1,4,4096]  Embedding 输出
H¹ [1,4,4096]  第1层输出
H² [1,4,4096]  第2层输出
H³ [1,4,4096]  第3层输出

形状没有改变，变化的是向量内部的数值和语义。第一层从词向量中建立初步上下文；后面的层再以此前的隐藏状态为输入，重新生成本层自己的 Q、K、V，并提取更复杂的关系。

需要特别区分两种流动方向：

• 隐藏状态 (H) 沿网络深度向上流动：(H^0 \rightarrow H^1 \rightarrow H^2 \rightarrow H^3)。
• KV Cache 留在各自所在的层，沿生成时间增长；第 2 层不会直接使用第 1 层的 K/V。

3. 第一个生成 Token 从哪里来

输入的四个 Token 经过全部层后得到：

H³.shape = [1,4,4096]

再经过 Final RMSNorm 和 LM Head：

[1,4,4096] × [4096,Vocab]
→ logits [1,4,Vocab]

训练时每个位置都承担下一个 Token 的预测目标：

“北京”的最终向量 → 预测“是”
“是”的最终向量   → 预测“中国”
“中国”的最终向量 → 预测“的”
“的”的最终向量   → 预测下一个 Token

推理时只取最后位置的 Logits。假设采样结果是“首都”，这时“首都”只是刚被选出来，还没有经过 Transformer。下一轮才会把“首都”送入模型，用来预测“。”。

二、拆开单个 Transformer 层

图 2：单层内部从 RMSNorm、QKV 投影到 Attention、MLP 和残差连接的形状变化。

对第 (l) 层，输入和输出可以写成：

输入：H^(l-1) [1,4,4096]
输出：H^l     [1,4,4096]

1. RMSNorm：稳定每个 Token 的数值尺度

RMSNorm 对每个 Token 的隐藏维度独立操作，不混合不同 Token：

x_norm = weight * x / sqrt(mean(x²) + eps)

输入输出形状相同：

[1,4,4096] → [1,4,4096]

与 LayerNorm 相比，RMSNorm 不减去均值，主要控制向量的整体尺度。现代 Llama、Qwen 等模型通常采用 Pre-Norm，即先归一化再进入 Attention 或 MLP。

2. 从一个隐藏状态产生 Q、K、V

RMSNorm 后的张量分别乘三组线性投影矩阵：

Q = X × WQ
K = X × WK
V = X × WV

在示例 GQA 配置中：

Q [1,32,4,128]
K [1, 8,4,128]
V [1, 8,4,128]

可以按下面的索引理解：

Q[batch, q_head,  token_position, head_dimension]
K[batch, kv_head, token_position, head_dimension]
V[batch, kv_head, token_position, head_dimension]

Q、K、V 的直觉含义是：

• Q（Query）：当前 Token 想找什么。
• K（Key）：每个历史 Token 用什么特征接受匹配。
• V（Value）：匹配到某个 Token 后，真正读取什么内容。

3. RoPE：将位置写入 Q 和 K

RoPE 不把位置向量直接加到 Embedding 上，而是在每一层中旋转 Q 和 K 的二维分量。第 (p) 个位置使用与位置相关的角度：

[x', y'] =
[x cosθ - y sinθ, x sinθ + y cosθ]

RoPE：

• 作用于 Q 和 K，通常不作用于 V；
• 不改变张量形状；
• 使点积能够自然表达两个 Token 的相对距离。

因此位置 0 的“北京”和位置 3 的“的”，即使某些内容特征相似，经过 RoPE 后也会带有不同的位置信息。

4. GQA：32 个查询头共享 8 组资料库

普通 MHA 中，32 个 Q Head 各自拥有一组 K/V；GQA 则让每四个 Q Head 共用一个 KV Head：

Q Head  0～3  → K/V Head 0
Q Head  4～7  → K/V Head 1
...
Q Head 28～31 → K/V Head 7

映射关系是：

kv_head = floor(q_head / 4)

这里始终只有 8 个独立 KV Head。可以为了矩阵对齐把每组 K/V 想象成逻辑重复四次，但高效内核通常不会在显存中真的复制四份。

以 Q Head 0 为例：

Q0 [4,128] × K0ᵀ [128,4]
→ Scores0 [4,4]

Softmax(Scores0) [4,4] × V0 [4,128]
→ O0 [4,128]

Q0～Q3 虽然共享 K0/V0，但 Q 不同，所以注意力权重不同，最终输出也不同：

O0 = softmax(Q0 K0ᵀ) V0
O1 = softmax(Q1 K0ᵀ) V0
O2 = softmax(Q2 K0ᵀ) V0
O3 = softmax(Q3 K0ᵀ) V0

全部 32 个 Q Head 都会产生一个 128 维输出：

O [1,32,4,128]
→ transpose + concat
→ [1,4,4096]

所以 GQA 压缩的是 K/V 的独立 Head 数量和 KV Cache，而 Attention 输出仍由 32 个 Query Head 恢复为 (32 \times 128 = 4096) 维。

5. 残差连接与 SwiGLU MLP

多头输出拼接后经过 (W_O) 投影，并与进入 Attention 前的主状态相加：

U = H^(l-1) + Attention(RMSNorm(H^(l-1)))

然后进入第二个 RMSNorm 和 SwiGLU MLP：

Gate = X × W_gate
Up   = X × W_up
Z    = SiLU(Gate) * Up
FFN  = Z × W_down
H^l  = U + FFN

Attention 负责在 Token 之间读取信息，MLP 主要对每个 Token 内部的特征进行非线性变换。两次残差连接保证原有信息可以绕过子层直接向后传播。

三、注意力的 4×4 矩阵表示什么

图 3：一张注意力矩阵中，行是 Query Token，列是 Key Token。

单个 Head 中：

Q [4,128] × Kᵀ [128,4]
→ Scores [4,4]

这个 4×4 不是某个 Token 的向量，而是四个 Token 两两之间的匹配分数：

行 i：第 i 个 Token 正在查询
列 j：第 j 个 Token 被查询
格 (i,j)：q_i · k_j / sqrt(128)

例如第 4 行第 3 列表示：

“的”的 Query 与“中国”的 Key 的匹配分数

1. 因果 Mask

Decoder-only Transformer 不能让当前 Token 看到未来：

北京 只能看 北京
是   可以看 北京、是
中国 可以看 北京、是、中国
的   可以看 北京、是、中国、的

未来位置会在 Softmax 前被设置为负无穷，Softmax 后权重变成 0。

2. 注意力权重如何发挥作用

假设“的”这一行 Softmax 后得到：

[北京:0.25, 是:0.10, 中国:0.55, 的:0.10]

那么输出向量是：

o_的 =
0.25 v_北京 +
0.10 v_是 +
0.55 v_中国 +
0.10 v_的

因此可以把注意力理解成：

Q/K 决定去哪里读、读多少；
V 提供真正被读出的内容；
O 是读取结果的加权汇总。

3. 为什么需要多头

如果只有一个 4096 维大 Head，所有特征只能共同产生一套 4×4 注意力分布。32 个 128 维 Head 则会产生 32 张不同的 4×4 矩阵：

Head 0 可能更关注“北京”
Head 1 可能更关注“中国”
Head 2 可能更关注“是”
其他 Head 可能编码位置、语法或混合关系

这些只是理解上的示意。真实模型不会由人手工指定每个 Head 的职责，而且一个 Head 的功能通常不是完全可解释或单一的。

四、Prefill、Decode 与 KV Cache

图 4：Prefill 一次处理历史输入，Decode 每轮只新增一个 Token 在每层的 K/V。

1. Prefill：输入 Token 一次并行通过模型

第一次收到提示词时，四个输入 Token 一起进入模型：

input_ids [1,4]

在每一层中都会计算四个位置的 Q/K/V，并保存该层的 K/V：

Layer 1 Cache：K¹/V¹ [1,8,4,128]
Layer 2 Cache：K²/V² [1,8,4,128]
Layer 3 Cache：K³/V³ [1,8,4,128]

不同层的缓存不能互换。第二层 K/V 是从 (H^1) 计算出来的，不是第一层 K/V 的复制品。

Prefill 完成后，使用最后位置“的”的最终隐藏状态预测“首都”。

2. Decode：每轮只处理最新生成的 Token

选择“首都”之后，逻辑上下文是：

北京 / 是 / 中国 / 的 / 首都

但是物理输入只有：

[首都]

在第 1 层中计算“首都”的新 (q_5^1,k_5^1,v_5^1)，让 (q_5^1) 查询第 1 层历史 K Cache，再把新 K/V 追加到第 1 层缓存。第 1 层输出 (h_5^1) 进入第 2 层，第 2 层再生成自己独立的新 Q/K/V，直到第 3 层。

这一轮每层的形状变化为：

新Q：        [1,32,1,128]
历史+新K/V：[1, 8,5,128]
Scores：     [1,32,1,5]
层输出：     [1, 1,4096]

三层结束后，根据“首都”的最终隐藏状态预测“。”。

下一轮只输入“。”：

新Q：        [1,32,1,128]
历史+新K/V：[1, 8,6,128]
Scores：     [1,32,1,6]

最后根据“。”的最终隐藏状态预测 <EOS>，推理引擎检测到结束标记后停止。

3. 为什么缓存 K/V，不缓存 Q

未来的新 Token 只需要用自己的 Query 查询历史：

q_new × [k_1, k_2, ..., k_new]

再根据注意力权重读取：

[v_1, v_2, ..., v_new]

未来不会再使用旧 Token 的 Query，因此历史 Q 没有复用价值。由于因果注意力保证旧 Token 不会看到未来，历史 Token 在各层中的 K/V 也不会因为新 Token 到来而改变，所以可以安全缓存。

4. KV Cache 节省了什么

没有 KV Cache 时，每生成一个 Token 都需要重新计算完整前缀：

输入4个Token → 预测第5个
输入5个Token → 预测第6个
输入6个Token → 预测第7个
...

有 KV Cache 后：

Prefill：历史前缀只计算一次
Decode：每轮只计算一个新Token，并复用历史K/V

但 Decode 并不是固定成本。上下文越长，新 Query 每层需要读取的历史 K/V 越多，因此单 Token 的 Attention 读取量会逐渐增大。

五、把三个阶段的张量放在一起

阶段	本轮实际输入	逻辑上下文长度	每层 Q	每层 KV Cache	Attention Scores
Prefill	北京、是、中国、的	4	[1,32,4,128]	[1,8,4,128]	[1,32,4,4]
Decode 1	首都	5	[1,32,1,128]	[1,8,5,128]	[1,32,1,5]
Decode 2	。	6	[1,32,1,128]	[1,8,6,128]	[1,32,1,6]

对于任意 Decoder-only Transformer，都可以用四个问题检查自己是否理解了推理链路：

1. 当前跨层传递的是哪一个隐藏状态 (H^l)？
2. 当前层生成并缓存的是哪一层的 K/V？
3. 当前 Query 对应哪个 Q Head，又映射到哪个 KV Head？
4. 当前是并行处理提示词的 Prefill，还是每轮一个新 Token 的 Decode？

六、总结

以“北京是中国的”续写为“首都。”，真正发生的是：

1. 四个输入Token一起完成Prefill；
2. 每一层建立自己的历史KV Cache；
3. 使用“的”的最终隐藏状态预测“首都”；
4. 只把“首都”送入全部层，结合缓存预测“。”；
5. 只把“。”送入全部层，结合缓存预测<EOS>；
6. 检测到<EOS>，停止生成。

其中：

• RMSNorm 稳定每个 Token 的隐藏向量尺度；
• RoPE 把位置信息写入每层的 Q/K；
• 多头 Attention 提供多套独立的信息读取方式；
• GQA 用较少的独立 KV Head 服务较多的 Q Head；
• KV Cache 避免在自回归生成中反复计算历史前缀；
• LM Head 把最后位置的 4096 维隐藏状态映射为整个词表上的下一个 Token 概率。

理解这条链路后，模型参数量、上下文长度、KV Head 数量、Batch Size、显存占用和推理速度之间的关系也就更容易分析了。