基础原理第 10 篇，共 13 篇

Decoder-only 完整手算：D=2

保留可手算的完整 Attention、LM Head 与增量缓存数值例子

更新于 2026-10-02阅读约 8 分钟

本文用一个 D=2、Vocab=6 的完整数值例子，逐项核对 Attention、LM Head 和缓存。这里 H 表示 hidden size；其他文章用 D 表示 hidden size、H 表示头数。本文的数值与其他文章里 D=8 的“小猫”例子无关，不要混用。

接着读：D=8 的训练和生成对照。

本文用一个可手算的微型模型，贯穿 Decoder-only LLM 的完整链路：

文本 / 对话消息
  → Chat Template
  → Token IDs
  → Token Embedding
  → Decoder Blocks
  → Final Hidden States
  → LM Head
  → Logits
  → Softmax / Sampling
  → Next Token
  → KV Cache 增量 Decode

Decoder-only 表示模型只使用带 Causal Mask 的 Transformer Decoder 主干，没有独立 Encoder 和 Cross-Attention。整体架构见 Transformer 总览。

1. 完整张量主线

先看一遍不展开内部的维度变化：

input_ids
[B, S]
   ↓ Embedding Lookup
hidden_states
[B, S, H]
   ↓ L × Decoder Block
final_hidden_states
[B, S, H]
   ↓ LM Head（H → V）
logits
[B, S, V]
   ↓ 取最后位置
next_token_logits
[B, V]
   ↓ logits processors + sampling
next_token_id
[B, 1]
   ↓ 追加历史，进入下一轮 Decode

符号	含义
`B`	batch size
`S`	当前输入序列长度
`H`	hidden size
`V`	vocabulary size
`L`	Decoder Block 层数

2. 微型模型与词表

假设模型只有：

参数	数值
Batch size `B`	1
Sequence length `S`	3
Hidden size `H`	2
Attention heads `A`	1
Head dimension `D`	2
Decoder layers `L`	1
Vocabulary size `V`	6

词表：

Token ID	Token
0	`<bos>`
1	`我`
2	`爱`
3	`猫`
4	`狗`
5	`<eos>`

输入文本是：

我爱

经过 Chat Template 和 Tokenizer 后，假设得到：

<bos> 我 爱

input_ids = [[0, 1, 2]]

\mathrm{input\_ids}\in\mathbb{N}^{1\times3}

Tokenizer 只负责字符串与 token id 之间的转换，不产生语义向量。分词过程见分词与对话格式。

3. Token IDs → Embedding

Embedding 矩阵的形状是：

W_E\in\mathbb{R}^{V\times H}=\mathbb{R}^{6\times2}

假设这三个 token 查表得到：

<bos> → [1, 0]
我    → [0, 1]
爱    → [1, 1]

那么：

X= \begin{bmatrix} 1&0\\ 0&1\\ 1&1 \end{bmatrix}

包含 batch 维后：

hidden_states = [
    [
        [1.0, 0.0],  # <bos>
        [0.0, 1.0],  # 我
        [1.0, 1.0],  # 爱
    ]
]

input_ids.shape     = [1, 3]
hidden_states.shape = [1, 3, 2]

位置信息在哪里

GPT-2 把学习式 Position Embedding 直接加到 Token Embedding。
Llama/Qwen 等现代模型通常不在这里直接相加，而是在 Attention 内对 Q/K 应用 RoPE。

两者不改变本文的外部主形状 [B, S, H]。

4. 通过一层 Decoder Block

现代 Pre-Norm Decoder Block 可概括为：

R_1=X+\operatorname{Attention}(\operatorname{Norm}(X))

R_2=R_1+\operatorname{MLP}(\operatorname{Norm}(R_1))

Norm、Attention、MLP 和残差连接都不改变 Block 对外的主形状：

[B, S, H] → [B, S, H]

4.1 Q、K、V 投影

为了可手算，假设 Norm 后数值不变，并且：

W_Q=W_K=W_V=I

因此：

Q=K=V=X

拆成 Attention Head 后：

Q.shape = [B, A, S, D] = [1, 1, 3, 2]
K.shape = [1, 1, 3, 2]
V.shape = [1, 1, 3, 2]

省略 batch 和 head 维：

Q=K=V= \begin{bmatrix} 1&0\\ 0&1\\ 1&1 \end{bmatrix}

真实的 GQA 模型中，Q Head 数可以大于 K/V Head 数；具体形状见 Attention 张量形状。

4.2 Attention Scores

Scaled Dot-Product Attention：

Scores=\frac{QK^T}{\sqrt D}

这里 $D=2$ ，因此：

Scores= \frac{1}{\sqrt2} \begin{bmatrix} 1&0&1\\ 0&1&1\\ 1&1&2 \end{bmatrix}

近似为：

Scores\approx \begin{bmatrix} 0.707&0&0.707\\ 0&0.707&0.707\\ 0.707&0.707&1.414 \end{bmatrix}

scores.shape = [B, A, S, S] = [1, 1, 3, 3]

第一个 S 表示 Query 位置，第二个 S 表示可以关注的 Key 位置。

4.3 Causal Mask

Decoder-only 模型不允许当前位置读取未来 token：

M= \begin{bmatrix} 0&-\infty&-\infty\\ 0&0&-\infty\\ 0&0&0 \end{bmatrix}

加上 Mask：

Scores+M= \begin{bmatrix} 0.707&-\infty&-\infty\\ 0&0.707&-\infty\\ 0.707&0.707&1.414 \end{bmatrix}

它表达：

<bos> 只能看 <bos>
我    可以看 <bos>、我
爱    可以看 <bos>、我、爱

4.4 Softmax 与 Attention Output

对 Scores 每一行做 Softmax：

A\approx \begin{bmatrix} 1&0&0\\ 0.330&0.670&0\\ 0.248&0.248&0.503 \end{bmatrix}

再计算：

O=AV

得到：

O\approx \begin{bmatrix} 1&0\\ 0.330&0.670\\ 0.751&0.751 \end{bmatrix}

attention_output.shape = [1, 3, 2]

最后位置的 Attention Output 已融合 <bos>、我、爱 三个位置的信息。

4.5 Residual、Norm 与 MLP

Attention Output 经过输出投影后与残差相加，然后进入 MLP。假设 MLP 中间维度为 $4H=8$ ：

MLP input:  [1, 3, 2]
   ↓ H → 4H
intermediate: [1, 3, 8]
   ↓ activation / gate
   ↓ 4H → H
MLP output: [1, 3, 2]

真实数值由训练出来的 Attention、Norm 和 MLP 权重决定。为了继续演示，假设经过全部 Decoder Block 与 Final Norm 后得到：

final_hidden_states = [
    [
        [1.0, 0.0],  # 编码 <bos>
        [0.0, 1.0],  # 编码 <bos> 我
        [1.0, 1.0],  # 编码 <bos> 我 爱
    ]
]

final_hidden_states.shape = [1, 3, 2]

真实模型会把这一 Block 重复 L 次，形状一直保持 [B, S, H]，但每层都在更新表示内容。

5. LM Head：Hidden Size → Vocabulary Size

LM Head 是一个对每个序列位置共享的线性投影：

W_{LM}\in\mathbb{R}^{V\times H}=\mathbb{R}^{6\times2}

假设：

W_{LM}= \begin{bmatrix} -1&-1\\ 2&0\\ 0&2\\ 1.5&1.5\\ 0.5&0.5\\ -0.5&-0.5 \end{bmatrix}

各行依次对应 <bos>、我、爱、猫、狗、<eos>。

计算：

Logits=H_{final}W_{LM}^T

形状变化：

[1, 3, 2] @ [2, 6]
        ↓
[1, 3, 6]

数值为：

Logits= \begin{bmatrix} -1&2&0&1.5&0.5&-0.5\\ -1&0&2&1.5&0.5&-0.5\\ -2&2&2&3&1&-1 \end{bmatrix}

logits.shape = [B, S, V] = [1, 3, 6]

每一行的位置对齐是：

Logits 行	Hidden State 看过的内容	用来预测
位置 0	`<bos>`	`我`
位置 1	`<bos> 我`	`爱`
位置 2	`<bos> 我爱`	下一个 token

一些模型会让 LM Head 和 Token Embedding 共享权重（weight tying），但是否共享不改变 [B,S,H] → [B,S,V] 的张量流。

6. 从最后位置选出下一个 Token

生成阶段只需要最后一行：

next_token_logits = logits[:, -1, :]

next_token_logits.shape = [1, 6]

数值：

[-2, 2, 2, 3, 1, -1]

Token	Logit	Softmax 概率（近似）
`<bos>`	-2	0.4%
`我`	2	19.4%
`爱`	2	19.4%
`猫`	3	52.7%
`狗`	1	7.1%
`<eos>`	-1	1.0%

如果使用 Greedy Decoding：

argmax → token id 3 → 猫

如果使用 Sampling，则按调整后的概率分布抽样，因此也有可能选到其他 token。

选中 猫 后，逻辑序列变成：

<bos> 我 爱 猫

7. 训练时为什么使用所有位置

对训练样本：

<bos> 我 爱 猫

可以看成：

输入：  <bos>  我   爱
标签：    我    爱   猫

也就是标签整体向左错位一格：

shift_logits = logits[:, :-1, :]
shift_labels = input_ids[:, 1:]

在更完整的样本中，猫 所在位置还可以用来预测 <eos>。训练使用所有可用位置计算 Cross Entropy Loss：

\mathcal L=-\sum_t\log P(x_{t+1}\mid x_{\le t})

对比：

阶段	LM Head 可产生的形状	实际使用
训练	`[B,S,V]`	所有有标签的位置
Prefill 生成	`[B,S,V]`	通常只需要 `[:, -1, :]`
增量 Decode	`[B,1,V]`	当前唯一新位置
投机验证	`[B,K,V]` 或等价布局	多个草稿位置

生产推理引擎可以不 materialize Prefill 的全部 [B,S,V]，而只对最后 hidden state 执行 LM Head，直接得到 [B,V]。

8. Prefill 后如何进入 KV Cache Decode

8.1 Prefill 结束时

对 prompt：

<bos> 我 爱

输入形状：

input_ids = [1, 3]

在每个 Decoder Layer 中建立：

K_cache.shape = [B, A, S, D] = [1, 1, 3, 2]
V_cache.shape = [1, 1, 3, 2]

Prefill 最后位置的 logits 选出 猫。此时 cache 中仍然只有 prompt 的 K/V；猫 只是刚被选中，还没有通过下一次 Transformer forward。

8.2 下一轮只输入新 Token

decode_input_ids = [[3]]  # 猫

decode_input_ids.shape = [1, 1]

经过 Embedding 和 Q/K/V 投影：

new_hidden.shape = [1, 1, 2]
Q_new.shape      = [1, 1, 1, 2]
K_new.shape      = [1, 1, 1, 2]
V_new.shape      = [1, 1, 1, 2]

追加 K/V：

历史 K_cache [1, 1, 3, 2]
新的 K_new   [1, 1, 1, 2]
          ↓ 沿 sequence 维追加
新 K_cache     [1, 1, 4, 2]

新 Query 需要读取全部四个 Key：

attention_scores.shape = [1, 1, 1, 4]

这一轮只产生新位置的 hidden state：

final_hidden.shape = [1, 1, 2]
logits.shape       = [1, 1, 6]

它用来预测 猫 之后的 token。如果选中 <eos>，生成结束。KV Cache 的分层结构见 KV Cache。

9. 端到端形状总表

阶段	张量	形状
Tokenize	`input_ids`	`[1,3]`
Embedding	`hidden_states`	`[1,3,2]`
Q/K/V Projection	`Q,K,V`	`[1,1,3,2]`
Attention Scores	`QK^T`	`[1,1,3,3]`
Attention Output	`attn_output`	`[1,3,2]`
MLP Up Projection	`intermediate`	`[1,3,8]`
MLP Down Projection	`mlp_output`	`[1,3,2]`
Final Norm	`final_hidden_states`	`[1,3,2]`
LM Head	`logits`	`[1,3,6]`
最后位置	`next_token_logits`	`[1,6]`
Token Selection	`next_token_id`	`[1,1]`
下一轮 Decode	`decode_input_ids`	`[1,1]`
Decode Query	`Q_new`	`[1,1,1,2]`
更新后 Cache	`K_cache,V_cache`	`[1,1,4,2]`
Decode Logits	`logits`	`[1,1,6]`

10. 对应到真实 LLM

本文微型模型	真实 Decoder-only LLM
`H=2`	hidden size 常为数百到数千
`V=6`	vocabulary 常为数万到十多万
`L=1`	通常堆叠数十层 Decoder Block
`A=1`	多头 Attention，可用 MHA、GQA 或 MQA
省略位置旋转	Llama/Qwen 等在 Q/K 上应用 RoPE
普通 MLP	现代模型常用 SwiGLU
单请求	服务引擎会做 continuous batching 和调度

尺寸会放大，但核心链路不变：

[B,S] → [B,S,H] → [B,S,H] → [B,S,V]
                                    ↓
                              选最后位置
                                    ↓
                                  [B,V]
                                    ↓
                                  [B,1]

11. 常见混淆

LM Head 只能输出最后一个位置吗？

不是。LM Head 可对 [B,S,H] 的所有位置做同一个 H→V 投影，得到 [B,S,V]。普通生成只需要最后位置，所以引擎可以只计算或只保留 [B,V]。

最后位置的 logits 是“最后 token 的概率”吗？

不是。它是“读完截至最后 token 的上下文后，下一 token 的分布”。

为什么训练用所有位置，生成只用最后位置？

训练已经知道整条正确序列，Causal Mask 允许各位置并行做 next-token prediction。生成时未来 token 尚不存在，只能使用最后位置预测下一个，再把结果追加回去。

选中的新 token 立即在 KV Cache 里吗？

不是。选中 token 只产生了 id。它要在下一次 Decode forward 中经过 Embedding 和每层 K/V Projection，才会被追加到各层 cache。

Hidden States 和 KV Cache 是同一份数据吗？

不是。Hidden State 是当前层的 token 表示；K/V 是该层将 Hidden State 经投影后为 Attention 保存的历史状态。每层都有自己的 K/V Cache。

Transformer 结构：Transformer 总览
Attention 细节：Attention 与因果遮罩
KV Cache：KV Cache

← Attention