注意力 · CausalSelfAttention 逐行拆解 · nanoGPT 源码解读

一行不跳地读完 nanoGPT 的因果自注意力——qkv 合并投影、多头拆分、缩放点积、因果掩码、Flash Attention 与手写版的对照，把 GPT 里最关键也最容易被示意图糊弄过去的部分彻底讲透。

注意力 · CausalSelfAttention 逐行拆解

到了。这是整个 GPT 里最关键的一块，也是最容易被「一张示意图」糊弄过去的一块。这一章我们一行都不跳，把 CausalSelfAttention 从头到尾啃下来。

先用一句话说清它要干的事：

在预测「下一个 token」时，当前位置该重点参考前文里的哪些位置？注意力就是算这个「参考权重」、并按权重把前文信息汇总过来的机制。

举个例子：读到「小明把钥匙忘在了车里，他回去拿」，要理解「他」指谁，模型得把注意力大部分投给「小明」。注意力机制，就是让每个位置学会「该回头看谁、看多重」。

`init`：先把零件备齐

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`：数据真正流动起来

前向传播是这一章的核心，我们切成四步。开头先记住三个维度符号：

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，并切成多头

        # 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 走快的，没有就手写。两者数学上等价，先看手写版最能看清原理：

        # 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}}$ $\frac{Q K ^{⊤}}{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 缓冲：

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：

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 加了它，只能从左看到右，才能一个接一个地「生成」。一张下三角矩阵，决定了模型是「阅读理解型」还是「写作接龙型」。

第四步：合并多头，收尾输出

        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

注意力 · CausalSelfAttention 逐行拆解

本页目录