# 注意力 · CausalSelfAttention 逐行拆解
> 一行不跳地读完 nanoGPT 的因果自注意力——qkv 合并投影、多头拆分、缩放点积、因果掩码、Flash Attention 与手写版的对照，把 GPT 里最关键也最容易被示意图糊弄过去的部分彻底讲透。
- 作者：David（道雾轩）
- 专栏：nanoGPT 源码解读（https://daiw.org/manual/nanogpt.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/nanogpt/attention
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）
# 注意力 · CausalSelfAttention 逐行拆解
到了。这是整个 GPT 里最关键的一块，也是最容易被「一张示意图」糊弄过去的一块。这一章我们**一行都不跳**，把 `CausalSelfAttention` 从头到尾啃下来。
先用一句话说清它要干的事：
> **在预测「下一个 token」时，当前位置该重点参考前文里的哪些位置？注意力就是算这个「参考权重」、并按权重把前文信息汇总过来的机制。**
举个例子：读到「小明把钥匙忘在了车里，他回去拿」，要理解「他」指谁，模型得把注意力**大部分投给「小明」**。注意力机制，就是让每个位置学会「该回头看谁、看多重」。
## `__init__`：先把零件备齐
```python
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
# key, query, value projections for all heads, but in a batch
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
# output projection
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
# regularization
self.attn_dropout = nn.Dropout(config.dropout)
self.resid_dropout = nn.Dropout(config.dropout)
self.n_head = config.n_head
self.n_embd = config.n_embd
self.dropout = config.dropout
# flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
if not self.flash:
print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
# causal mask to ensure that attention is only applied to the left in the input sequence
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size))
```
- `assert config.n_embd % config.n_head == 0`：嵌入维度必须能被头数整除。因为一会儿要把 `n_embd` 维**平均切成 `n_head` 份**，每头分 `n_embd / n_head` 维，除不尽就没法切。
- `self.c_attn = nn.Linear(n_embd, 3 * n_embd)`：**一个矩阵，一次算出 Q、K、V 三样**。你可能以为要三个投影，这里却把它们**拼成一个宽三倍的线性层**——一次矩阵乘法比三次快，算完再切开。这是很典型的工程优化。
- `self.c_proj = nn.Linear(n_embd, n_embd)`：注意力算完后的**输出投影**，把汇总来的信息再过一层线性变换。
- 两个 `Dropout`：`attn_dropout` 作用在注意力权重上，`resid_dropout` 作用在输出上，都是抗过拟合用（默认 `dropout=0` 时是空操作）。
- `self.flash = hasattr(...)`：检查当前 PyTorch 有没有内置的 `scaled_dot_product_attention`（即 **Flash Attention**,PyTorch ≥ 2.0 才有）。有就用它，快得多。
- **只有没有 Flash 时**，才用 `register_buffer` 造一个**因果掩码** `bias`——这就是让模型「只能看过去」的关键道具，下面细讲。
**Q、K、V 是什么？** 一个常用的类比：每个 token 都发出一个「**查询** Query（我想找什么）」、亮出一个「**钥匙** Key（我是什么）」、并携带一份「**内容** Value（我能提供什么信息）」。某个位置的 Query 去和所有位置的 Key 比对，匹配度高的，就多取一点那个位置的 Value。注意力，本质就是这样一场「按匹配度加权取信息」的过程。
## `forward`：数据真正流动起来
前向传播是这一章的核心，我们切成四步。开头先记住三个维度符号：
```python
def forward(self, x):
B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
```
- `B`：batch，一批有几个序列。
- `T`：time / 序列长度，当前这段有几个 token。
- `C`:channel，就是 `n_embd`，每个 token 向量多长。
后面所有 reshape，盯住这三个字母就不会晕。
### 第一步：一次算出 Q、K、V，并切成多头
```python
# calculate query, key, values for all heads in batch and move head forward to be the batch dim
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
```
- `self.c_attn(x)`：把 `(B, T, C)` 的输入过那个「宽三倍」的线性层，得到 `(B, T, 3C)`。
- `.split(self.n_embd, dim=2)`：沿最后一维**切成三段**，每段 `C` 宽，正好就是 `q`、`k`、`v`，各自 `(B, T, C)`。这就是前面「一次算完再切开」的兑现。
- 接下来三行长得一样，做的是**拆多头**。以 `q` 为例：
- `.view(B, T, n_head, C // n_head)`：把最后那 `C` 维**劈成 `n_head` 份**，每份 `hs = C // n_head` 维(head size)。形状变成 `(B, T, nh, hs)`。
- `.transpose(1, 2)`：把 `T` 和 `nh` 两个维度**对调**，变成 `(B, nh, T, hs)`。
为什么要 `transpose`？这样一来，`nh`（头）就挪到了和 `B`(batch)相邻的前排。后面做批量矩阵乘法时，PyTorch 会把 `(B, nh)` 当成「一批批」来并行处理——**等于让 `n_head` 个头各算各的注意力，互不干扰，还全程并行**。
**为什么要「多头」？** 一个头只能学一种「该看谁」的模式。多头就是让模型**同时开好几路注意力**：也许一头专门盯语法搭配、一头盯指代关系、一头盯远处的主题词……各看各的，最后拼起来。把一个大注意力拆成几个小注意力并行，表达力更丰富，计算量还不变。
### 第二步：算注意力——两条路
nanoGPT 给了两套实现：有 Flash 走快的，没有就手写。**两者数学上等价**，先看手写版最能看清原理：
```python
# causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
if self.flash:
# efficient attention using Flash Attention CUDA kernels
y = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 0, is_causal=True)
else:
# manual implementation of attention
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
```
手写版这四行，正是那个著名公式的代码化：
$$
\text{Attention}(Q,K,V)=\text{softmax}\!\left(\frac{QK^{\top}}{\sqrt{d_k}}+M\right)V
$$
逐行对上：
- **`att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))`** —— 这是 $\frac{QK^\top}{\sqrt{d_k}}$。
- `q @ k^T`：每个位置的 Query 和**所有**位置的 Key 做点积，得到一张 `(T, T)` 的**相似度表**——第 `i` 行第 `j` 列，就是「位置 `i` 该有多关注位置 `j`」。
- 除以 `√d_k`（`d_k` 就是 `hs`）：**缩放**。点积会随维度变大而数值膨胀，不压一压，softmax 会被推到极端、梯度消失。这个 `√d_k` 就是「缩放点积注意力」里的「缩放」。
- **`att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))`** —— 这是加上掩码 $M$，也是「**Causal**（因果）」这个词的来历，下一节专门讲。
- **`att = F.softmax(att, dim=-1)`** —— 把每一行的分数**归一化成概率**（每行加起来为 1）。现在第 `i` 行就是「位置 `i` 把注意力如何分配给各位置」的权重。
- **`att = self.attn_dropout(att)`** —— 训练时随机丢弃一些权重，防过拟合。
- **`y = att @ v`** —— 拿权重去**加权汇总** Value。位置 `i` 的输出，就是所有位置 Value 按第 `i` 行权重的加权和。**至此，每个位置都「按需」把前文信息收集到了自己身上。**
Flash 版那一行 `scaled_dot_product_attention(..., is_causal=True)` 做的是**完全一样的数学**，只是把这几步融进一个高度优化的 CUDA 内核里，不在显存里实际造出那张 `(T, T)` 的大表，又快又省显存。`is_causal=True` 就等价于下面要讲的掩码。
### 第三步：因果掩码——「只能看过去」的那道法术
这是整章的题眼。语言模型生成时是**从左往右**一个个吐字的，预测第 `i` 个 token 时，它**只能看到第 0 到 `i-1` 个**，绝不能偷看第 `i` 个之后的——否则就是「抄答案」，训练出的模型一到真实生成（没有未来可抄）就废。
「不许看未来」这条铁律，在代码里就是那个 `bias` 缓冲：
```python
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
.view(1, 1, config.block_size, config.block_size))
```
- `torch.ones(block_size, block_size)`：一张全是 1 的方阵。
- `torch.tril(...)`：取**下三角**(triangular-lower)，对角线及以下保留 1，上方全变 0。对 `T=4` 来说长这样：
```
看j=0 j=1 j=2 j=3
位置i=0 [ 1 0 0 0 ]
位置i=1 [ 1 1 0 0 ]
位置i=2 [ 1 1 1 0 ]
位置i=3 [ 1 1 1 1 ]
```
第 `i` 行只有第 0..`i` 列是 1——**正好就是「位置 `i` 只允许看位置 0 到 `i`」**。
再回看那行 `masked_fill`：
```python
att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
```
凡是掩码里为 `0` 的位置（即「未来」），把相似度分数强行填成 **负无穷**。这样接下来 `softmax` 一算，`e^{-∞} = 0`，**这些未来位置的权重精确地变成 0**，一点信息都汇总不过来。一行 `masked_fill`，就把「不许看未来」这条物理定律焊死进了模型。
这正是 GPT(decoder-only)和 BERT(encoder)的分水岭：BERT 没有这道下三角掩码，每个位置能同时看左右两边，适合「理解」；GPT 加了它，只能从左看到右，才能一个接一个地「生成」。**一张下三角矩阵，决定了模型是「阅读理解型」还是「写作接龙型」。**
### 第四步：合并多头，收尾输出
```python
y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
# output projection
y = self.resid_dropout(self.c_proj(y))
return y
```
- `y` 此刻是 `(B, nh, T, hs)`——每个头各算出了自己那份结果。
- `.transpose(1, 2)`：把 `nh` 和 `T` 换回来，变 `(B, T, nh, hs)`。
- `.contiguous().view(B, T, C)`：把 `nh` 个头的结果**并排拼回**成一个 `C` 维向量(`nh * hs = C`)。`contiguous()` 是因为 `transpose` 只改了「视图」没真挪内存，`view` 前得先把内存排整齐。**几个头「分头侦查」得到的情报，在这里汇总成一份。**
- `self.c_proj(y)`：过输出投影，把拼起来的多头结果再融合一下。
- `self.resid_dropout(...)`：最后再来一道 dropout，返回。
输出 `y` 的形状，又变回了进来时的 `(B, T, C)`——**尺寸不变，但每个位置的向量，已经揉进了它该关注的前文信息**。这也是为什么 `Block` 里能直接 `x + self.attn(...)`：进出同形，才好做残差相加。
## 把这一章串起来
`CausalSelfAttention` 做的事，一句话：
> **每个位置，拿自己的 Query 去和「自己及之前」所有位置的 Key 比对，算出一组注意力权重（未来位置被掩码强制清零），再按权重把那些位置的 Value 加权汇总到自己身上。多个头并行地做这件事，最后拼合、投影输出。**
回头看代码，它其实只有干净利落的十来行。但就是这十来行，让 GPT 拥有了「联系上下文」的能力——它是「大模型能读懂长文」这件事的**物理基础**。
至此，五个零件里最硬的一块啃完了。剩下的 `LayerNorm`、`MLP`、`Block` 都简单得多——下一章我们一口气把它们收拾掉，你会发现读完注意力再看它们，轻松得像饭后甜点。👉 [其余三个零件：LayerNorm、MLP、Block](https://daiw.org/manual/nanogpt/the-other-pieces)