# 第 11 课：Attention 原理与源码
这一课只解决一个问题：一个 token 如何通过 Attention 从上下文 token 里选择并汇总信息，然后仍然回到 `[batch, seq, hidden_size]` 主干。
## 目录
- [0. 本节主线](#l11-mainline)
- [1. 本节要懂的 5 个层次](#l11-levels)
- [2. Attention 一次完整原理讲解](#l11-complete-principle)
- [3. Attention.forward 一次完整源码走读](#l11-forward-walkthrough)
- [4. 源码变量字典](#l11-variable-dict)
- [5. 关键源码对照](#l11-source-map)
- [6. 实验验证](#l11-experiment)
- [7. 本节检查](#l11-check)
- [8. 下一课](#l11-next)
## 0. 本节主线
Attention 的完整逻辑是：
```text
每个 token 先有自己的 hidden 向量
-> 每个 hidden 向量生成 Q/K/V 三种角色
-> 当前 token 的 Q 和所有可见 token 的 K 计算匹配分数
-> causal mask 把未来 token 屏蔽掉
-> softmax 把分数变成注意力权重
-> 用注意力权重对所有可见 token 的 V 加权求和
-> 得到当前 token 汇总上下文后的新表示
-> 所有 token、所有 head 并行做这件事
-> 合并多头并用 o_proj 回到 hidden_size
```
一句话：
```text
Q/K 决定看谁，V 决定拿什么，softmax 决定拿多少，causal mask 决定不能看未来，o_proj 决定回到主干。
```
这节课分成两大块：
```text
先完整讲一遍 Attention 原理；
再完整走一遍 MiniMind 的 Attention.forward 源码。
```
## 1. 本节要懂的 5 个层次
| 层次 | 要理解什么 | 源码位置 |
|---|---|---|
| Attention 解决什么问题 | token 需要从上下文 token 里取信息 | `model/model_minimind.py:178-194` |
| 一次完整 Attention 怎么算 | Q/K 打分，softmax，汇总 V | `model/model_minimind.py:111-134` |
| Q/K/V 从哪里来 | 同一份 hidden_states 投影成三种角色 | `model/model_minimind.py:91-116` |
| mask 和多头怎么进来 | causal mask 禁止看未来，多头并行做 attention | `model/model_minimind.py:123-131` |
| 怎么回到主干 | repeat_kv 对齐 K/V heads，o_proj 输出回 hidden_size | `model/model_minimind.py:86-89`, `model/model_minimind.py:132-134` |
学完本节，你应该能不看源码说出一遍 Attention 的完整计算过程，然后再能对着 `Attention.forward` 逐行解释每个变量的 shape 和作用。
## 2. Attention 一次完整原理讲解
先不要看源码。我们用一个序列来理解 Attention：
```text
位置: 0 1 2 3
token: A B C D
```
假设现在看位置 2，也就是 token `C`。
因为这是 causal language model，`C` 只能看：
```text
A, B, C
```
不能看未来的 `D`。
### 2.1 每个 token 先有 hidden 向量
第 9 课讲过，进入 block 时，每个 token 都已经有自己的 hidden 向量：
```text
h_A, h_B, h_C, h_D
```
这些向量都在同一个 hidden 空间里，比如每个都是 64 维或 768 维。
如果模型只对每个 `h_i` 单独做 MLP，它就不能直接知道别的 token 里有什么信息。Attention 要补上的就是：
```text
让 h_C 能从 h_A、h_B、h_C 里按需取信息。
```
### 2.2 为什么要有 Q/K/V
Attention 把每个 hidden 向量投影成三种角色：
```text
q_i = W_Q h_i
k_i = W_K h_i
v_i = W_V h_i
```
对 token `C` 来说：
```text
q_C：C 当前想找什么信息
k_A/k_B/k_C：A/B/C 各自提供什么可匹配的特征
v_A/v_B/v_C：A/B/C 真正能贡献给 C 的内容
```
所以 Q/K/V 不是三个不同输入，而是同一批 hidden 向量的三种投影角色：
```text
Q/K 用来决定“看谁”；
V 用来提供“拿什么”。
```
### 2.3 Q 和 K 怎么决定“看谁”
`C` 要决定看 A、B、C 各多少，就用 `q_C` 分别和它们的 key 做点积：
```text
score(C, A) = q_C · k_A / sqrt(head_dim)
score(C, B) = q_C · k_B / sqrt(head_dim)
score(C, C) = q_C · k_C / sqrt(head_dim)
```
点积越大，表示越匹配。这里的分数不是概率，只是未归一化的匹配分数。
除以 `sqrt(head_dim)` 是为了缩放分数尺度。否则 `head_dim` 越大，点积越容易变得很大，softmax 会过早变得极端。
### 2.4 causal mask 为什么出现
如果不加限制，`C` 也会和 `D` 算：
```text
score(C, D)
```
但训练语言模型时，位置 `C` 不能提前看到未来 token `D`。否则第 10 课讲的 next-token prediction 就会作弊。
所以 causal mask 会把未来位置变成：
```text
score(C, D) = -inf
```
softmax 之后：
```text
weight(C, D) = 0
```
也就是未来 token 完全不参与信息汇总。
### 2.5 softmax 后为什么乘 V
对 A、B、C 的分数经过 softmax 后，变成权重：
```text
weight(C, A), weight(C, B), weight(C, C)
```
这些权重表示：
```text
C 应该从 A/B/C 各拿多少信息。
```
然后真正汇总的是 value：
```text
new_C =
weight(C, A) * v_A
+ weight(C, B) * v_B
+ weight(C, C) * v_C
```
这就是 Attention 的核心公式：
```text
Attention(Q, K, V) = softmax(QK^T / sqrt(head_dim) + mask) V
```
### 2.6 从一个 token 推广到所有 token
上面只讲了 `C` 如何看 A/B/C。实际模型会对每个位置都做同样的事：
```text
A 看 A
B 看 A/B
C 看 A/B/C
D 看 A/B/C/D
```
所以 attention scores 会形成一个矩阵：
```text
[query_position, key_position]
```
如果有 batch 和 heads，就变成：
```text
[batch, heads, query_seq, key_seq]
```
这就是为什么实验里 `scores.shape=(2, 4, 5, 5)`。
### 2.7 多头 attention 是什么
如果只用一个 head，模型只有一种“看上下文”的方式。
多头 attention 是把 hidden 向量拆成多个子空间：
```text
hidden_size = num_heads * head_dim
```
每个 head 都独立算一套：
```text
QK 分数 -> softmax -> 加权 V
```
最后把多个 head 的结果合并，再映射回 `hidden_size`。
所以多头不是多了几层，也不是复制 batch；它是在 hidden 维度内部并行做多个 attention 视角。
### 2.8 这一节原理到此为止
到这里你应该已经有完整观念：
```text
Attention 不是神秘模块；
它就是让每个 token 用 Q 去匹配可见 token 的 K，
再按匹配权重汇总这些 token 的 V。
```
下一步再看源码，源码只是把这套过程用张量并行实现。
## 3. Attention.forward 一次完整源码走读
这一节先完整走一遍源码，不拆成零散片段。
### 源码证据 A：Attention.forward 全流程
文件：`model/model_minimind.py:111-134`
看它是为了理解：MiniMind 如何把上一节的完整 Attention 原理实现成张量计算。
源码：
```python
def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
bsz, seq_len, _ = x.shape
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xq, xk = self.q_norm(xq), self.k_norm(xk)
cos, sin = position_embeddings
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
if past_key_value is not None:
xk = torch.cat([past_key_value[0], xk], dim=1)
xv = torch.cat([past_key_value[1], xv], dim=1)
past_kv = (xk, xv) if use_cache else None
xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
if self.flash and (seq_len > 1) and (not self.is_causal or past_key_value is None) and (attention_mask is None or torch.all(attention_mask == 1)):
output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)
else:
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(output))
return output, past_kv
```
这段源码可以分成 8 步读。
### 第 1 步：确认输入 shape
```python
bsz, seq_len, _ = x.shape
```
这里的 `x` 是 block 传进 Attention 的 hidden states，来自：
```python
self.input_layernorm(hidden_states)
```
它的 shape 是：
```text
[batch, seq, hidden_size]
```
这一步只取出 batch 和 seq 长度，为后面 reshape 做准备。
### 第 2 步：从 x 投影出 Q/K/V
```python
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
```
这里对应原理里的：
```text
q_i = W_Q h_i
k_i = W_K h_i
v_i = W_V h_i
```
同一份 `x` 通过三组 Linear 变成三种角色：
```text
xq: query
xk: key
xv: value
```
### 第 3 步：把 Q/K/V reshape 成多头
```python
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
```
这一步把最后一维拆开：
```text
Q: [batch, seq, num_attention_heads, head_dim]
K: [batch, seq, num_key_value_heads, head_dim]
V: [batch, seq, num_key_value_heads, head_dim]
```
MiniMind 允许 Q heads 和 K/V heads 不同，这是 GQA 的设计。
### 第 4 步：Q/K norm 和 RoPE
```python
xq, xk = self.q_norm(xq), self.k_norm(xk)
cos, sin = position_embeddings
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)
```
`q_norm` 和 `k_norm` 对每个 head 的向量做 RMSNorm，稳定 Q/K 的尺度。
RoPE 作用在 Q/K 上，是为了让 Q/K 匹配时带上位置信息。注意 V 不参与 RoPE，因为位置信息主要影响“怎么匹配”，不是“拿回什么内容”。
RoPE 细节放第 12 课，这里只记住：
```text
RoPE 修改 Q/K，不修改 V。
```
### 第 5 步：KV cache 先跳过细节
```python
if past_key_value is not None:
xk = torch.cat([past_key_value[0], xk], dim=1)
xv = torch.cat([past_key_value[1], xv], dim=1)
past_kv = (xk, xv) if use_cache else None
```
这段用于生成阶段的 KV cache。
训练时通常一次性输入整段序列，可以先不看。第 12 课会专门讲：
```text
为什么生成时要缓存历史 K/V；
为什么不用每步重算全部历史。
```
本节只要知道：
```text
cache 缓存的是 K/V，不是 Q。
```
### 第 6 步：转成 attention 计算布局，并 repeat K/V
```python
xq, xk, xv = (
xq.transpose(1, 2),
repeat_kv(xk, self.n_rep).transpose(1, 2),
repeat_kv(xv, self.n_rep).transpose(1, 2)
)
```
原来多头 shape 是：
```text
[batch, seq, heads, head_dim]
```
attention 矩阵乘法更方便的布局是：
```text
[batch, heads, seq, head_dim]
```
所以要 `transpose(1, 2)`。
K/V 先经过 `repeat_kv`，是因为 MiniMind 使用 GQA：
```text
Q heads 多；
K/V heads 少；
计算前把 K/V heads 复制到和 Q heads 一样多。
```
### 第 7 步：计算 scores、mask、softmax、加权 V
普通分支是最容易理解的：
```python
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
if self.is_causal:
scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
if attention_mask is not None:
scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
```
这几行就是 Attention 原理公式：
```text
scores = QK^T / sqrt(head_dim)
weights = softmax(scores + mask)
output = weights V
```
shape 是：
```text
xq: [batch, heads, query_seq, head_dim]
xk^T: [batch, heads, head_dim, key_seq]
scores: [batch, heads, query_seq, key_seq]
xv: [batch, heads, key_seq, head_dim]
output: [batch, heads, query_seq, head_dim]
```
causal mask 的作用是：
```text
未来 key position 的分数变成 -inf；
softmax 后未来位置权重变成 0。
```
flash attention 分支：
```python
F.scaled_dot_product_attention(...)
```
做的是同一件事，只是调用 PyTorch 优化实现。
### 第 8 步：合并 heads，o_proj 回到 hidden_size
```python
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(output))
return output, past_kv
```
attention 内部输出是：
```text
[batch, heads, seq, head_dim]
```
先转回：
```text
[batch, seq, heads, head_dim]
```
再合并：
```text
[batch, seq, heads * head_dim]
```
最后 `o_proj` 输出：
```text
[batch, seq, hidden_size]
```
这样 Attention 才能回到 `MiniMindBlock` 的 residual stream：
```python
hidden_states += residual
```
### 完整源码走读到此为止
你读 `Attention.forward` 时，不要把它看成很多互不相关的小技巧。它就是这条链：
```text
x
-> Q/K/V
-> 多头 shape
-> Q/K 加位置
-> K/V cache 可选拼接
-> QK 打分
-> mask
-> softmax 权重
-> 汇总 V
-> 合并 heads
-> o_proj 回主干
```
## 4. 源码变量字典
这一节把 `Attention.forward` 里的关键变量统一列出来，方便你回看源码。
| 变量 | shape | 含义 |
|---|---|---|
| `x` | `[batch, seq, hidden_size]` | Attention 输入 hidden states |
| `xq` | `[batch, seq, q_heads, head_dim]` 后转 `[batch, q_heads, seq, head_dim]` | query，用来匹配 key |
| `xk` | `[batch, seq, kv_heads, head_dim]` 后转 `[batch, q_heads, seq, head_dim]` | key，被 query 匹配 |
| `xv` | `[batch, seq, kv_heads, head_dim]` 后转 `[batch, q_heads, seq, head_dim]` | value，被权重加权汇总 |
| `scores` | `[batch, heads, query_seq, key_seq]` | 每个 query 位置看每个 key 位置的分数 |
| `attention_mask` | 可广播到 scores | padding 等外部 mask |
| `output` | 先 `[batch, heads, seq, head_dim]`，最后 `[batch, seq, hidden_size]` | Attention 汇总后的结果 |
| `past_kv` | `(xk, xv)` | 生成时缓存历史 K/V |
## 5. 关键源码对照
前面已经完整走读过源码。这里只把几个容易混的点再对照一次。
### 5.1 Q/K/V 从哪里来
文件：`model/model_minimind.py:91-116`
看它是为了理解：Q/K/V 是同一份 `x` 的三套投影，不是三个外部输入。
源码：
```python
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
...
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
```
这段代码说明：
- 输入都是同一个 `x`。
- Q/K/V 是三套投影参数学出来的三种表示。
- Q/K 负责匹配，V 负责被汇总。
### 5.2 为什么 scores 是 `[batch, heads, seq, seq]`
文件：`model/model_minimind.py:123-128`
看它是为了理解：query position 和 key position 如何形成二维注意力矩阵。
源码：
```python
xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
```
这段代码说明：
- `xq` 是 `[batch, heads, query_seq, head_dim]`。
- `xk.transpose(-2, -1)` 是 `[batch, heads, head_dim, key_seq]`。
- 相乘后得到 `[batch, heads, query_seq, key_seq]`。
### 5.3 causal mask 为什么加在 scores 上
文件：`model/model_minimind.py:128-131`
看它是为了理解：未来 token 为什么不会进入 softmax 权重。
源码：
```python
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
if self.is_causal: scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
```
这段代码说明：
- causal mask 是在 softmax 前加到 scores 上。
- 未来位置变成 `-inf`。
- softmax 后未来位置权重为 0。
- 最后才用权重乘 V。
### 5.4 GQA 和 o_proj 如何收尾
文件：`model/model_minimind.py:86-89`, `model/model_minimind.py:132-134`
看它是为了理解：K/V heads 如何对齐 Q heads，以及多头结果如何回到主干。
源码：
```python
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
bs, slen, num_key_value_heads, head_dim = x.shape
if n_rep == 1: return x
return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim))
output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(output))
return output, past_kv
```
这段代码说明：
- `repeat_kv` 把较少的 K/V heads 扩到 Q heads 数量。
- 多头输出合并后经过 `o_proj`。
- 最终输出回到 `[batch, seq, hidden_size]`。
## 6. 实验验证
### 实验 A：追踪完整 Attention 流程
这个实验验证：
```text
Q/K/V 来自同一个 hidden_states；
scores 是 QK 匹配矩阵；
causal mask 让未来位置权重为 0；
softmax(scores) @ V 后再 o_proj；
手算结果和模块 forward 完全一致。
```
运行：
```bash
cd /home/sun/minimind
PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_attention_shapes.py
```
记录：
```text
x.shape =
q_proj(x).shape =
k_proj(x).shape =
v_proj(x).shape =
scores.shape =
attention weights for batch 0, head 0 =
future_attention_mass_for_batch0_head0 =
manual_vs_module_max_abs_diff =
```
你应该看到：
```text
x.shape=(2, 5, 64)
q_proj(x).shape=(2, 5, 64)
k_proj(x).shape=(2, 5, 32)
v_proj(x).shape=(2, 5, 32)
scores.shape=(2, 4, 5, 5)
future_attention_mass_for_batch0_head0=0.0000000000
manual_vs_module_max_abs_diff=0.000000000000
```
重点看输出的 attention weights 矩阵。它应该是下三角有效：
```text
第 0 行只能看第 0 列；
第 1 行只能看第 0、1 列；
第 2 行只能看第 0、1、2 列；
未来位置权重为 0。
```
### 实验 B：改变 head 配置
运行：
```bash
PYTHONDONTWRITEBYTECODE=1 python course/labs/trace_attention_shapes.py --hidden_size 96 --num_attention_heads 6 --num_key_value_heads 3
```
记录：
```text
head_dim =
q_proj(x).shape =
k_proj(x).shape =
xq.shape =
xk after repeat shape =
scores.shape =
```
你应该看到：
```text
head_dim=16
q_proj(x).shape=(2, 5, 96)
k_proj(x).shape=(2, 5, 48)
xq.shape=(2, 5, 6, 16)
xk after repeat shape=(2, 5, 6, 16)
scores.shape=(2, 6, 5, 5)
```
这说明：
- head 配置改变的是 Attention 内部多头形状。
- 输出仍然要回到 `[batch, seq, hidden_size]`。
## 7. 本节检查
如果你真懂了本节，应该能不看答案说清楚：
1. Attention 要解决的核心问题是什么。
2. 从一个 token 的角度，Attention 是如何从上下文拿信息的。
3. Q/K/V 为什么要分成三种角色。
4. QK scores 的每个元素表示什么。
5. 为什么 softmax 后乘的是 V，不是 K。
6. causal mask 和第 10 课的 next-token loss 是什么关系。
7. `Attention.forward` 从 `x` 到 `output` 的完整执行顺序是什么。
8. 为什么最终必须用 `o_proj` 回到 `[batch, seq, hidden_size]`。
## 8. 下一课
第 12 课进入 `RoPE 与 KV cache`：我们会讲位置编码如何作用在 Q/K 上，以及生成时 past key/value 如何让模型不用每一步重算全部历史。