1. 从 llm.c 的极简哲学说起：为什么第三篇要死磕注意力与 MLP
如果你跟着前两篇一路走过来，应该已经对 llm.c 这个项目的脾气有所了解了。它不追求工程上的花哨，也不搞一堆抽象层，而是用最朴素的 C 代码，把 GPT-2 从零到一搭出来。到了第三篇，我们终于要碰整个 Transformer 架构里最核心、也最容易被讲玄乎的两块骨头：多头自注意力机制和MLP 前馈网络。顺带还要把夹在它们中间的LayerNorm彻底讲透。
很多人学 GPT 类模型，卡就卡在这一步。看论文里的公式觉得懂了，看 PyTorch 的nn.MultiheadAttention觉得懂了，但一旦让你用 C 手写一遍，立刻发现脑子里全是浆糊：Q、K、V 到底怎么切分？多头到底是"多"在哪里？缩放因子为什么是1/sqrt(d_k)？残差连接加在哪个位置？LayerNorm 到底对哪个维度做归一化？这些问题不解决，你永远只是调包侠，而不是真正理解 GPT-2 的人。
llm.c 的价值恰恰在这里。它逼着你把每一个张量形状、每一次内存布局、每一个循环边界都想清楚。这篇我就按照 llm.c 第三篇的节奏，把多头注意力机制原理、MLP 的内部结构、LayerNorm 的归一化维度这三件事，从代码视角掰开揉碎讲一遍。适合已经看过前两篇、对 GPT-2 整体结构有概念，但一到注意力细节就发怵的读者。如果你还没看过前两篇，建议先补一下张量基础和嵌入层那部分，不然这篇里的形状推导会让你有点晕。
我先把这一篇要解决的核心问题列出来，你可以对照自己的理解程度：
- 单头注意力到底在算什么，为什么需要"多头"
- Q、K、V 的投影矩阵在 llm.c 里是怎么存的，形状怎么对齐
- 缩放点积注意力的缩放因子从哪来，不缩放会怎样
- 因果掩码（causal mask）在代码里怎么实现，为什么不能简单置零
- MLP 为什么是两层线性加一个 GELU，中间维度为什么是 4 倍
- LayerNorm 的均值方差到底在哪个维度上算，和 BatchNorm 差在哪
- 残差连接和 LayerNorm 的先后顺序（Pre-LN vs Post-LN）对训练的影响
这些问题，我会一个一个拆。你不需要有深厚的数学背景，但需要愿意跟着形状推导走一遍。下面正式开始。
2. 多头自注意力机制：从单头到多头的形状推演
2.1 单头注意力：一次"加权平均"的本质
先把最朴素的单头注意力讲清楚。假设输入序列长度为T，每个 token 的嵌入维度为C（GPT-2 small 里C=768）。输入张量形状是[T, C]。注意力的目标是：让每个位置的 token 去"看"其他所有位置，然后根据相关性做一次加权平均。
具体怎么做？用三个线性投影把输入分别映射成 Query、Key、Value：
Q = X @ W_q，形状[T, C]K = X @ W_k，形状[T, C]V = X @ W_v，形状[T, C]
然后计算注意力分数scores = Q @ K^T，形状变成[T, T]。这个矩阵的第i行第j列，表示第i个 token 对第j个 token 的"关注程度"。接着做缩放、掩码、softmax，得到注意力权重attn，形状还是[T, T]。最后out = attn @ V，形状回到[T, C]。
用一句话概括：注意力就是让每个 token 根据自己和其他 token 的相似度，对其他 token 的 Value 做加权求和。相似度用 Q 和 K 的点积衡量，权重用 softmax 归一化。
这里有个细节很多人第一次看会忽略：Q @ K^T得到的是[T, T]，而不是[T, C]。也就是说，注意力的计算复杂度和序列长度的平方成正比。这就是为什么长上下文那么吃显存——T翻倍，注意力矩阵大小翻四倍。llm.c 里这个矩阵是实打实分配出来的，你能直观看到它的内存占用。
2.2 多头到底"多"在哪：切分维度的选择
单头注意力的表达能力有限，因为它只学了一种"关注模式"。多头注意力的思路是：把C维的 Q、K、V 切成H份，每份维度是C/H，每一份独立做一次注意力，最后把结果拼回来再投影一次。
GPT-2 small 里H=12，C=768，所以每个头的维度d_head = 768/12 = 64。这个 64 不是随便定的，它要保证H * d_head = C，这样拼接后维度才能对得上。
关键问题来了：多头是"切分 Q、K、V 的维度"，还是"复制 H 份独立的 Q、K、V"？答案是前者。在 llm.c 里，W_q、W_k、W_v的形状都是[C, C]，投影完之后得到[T, C]，然后 reshape 成[T, H, d_head]，再 transpose 成[H, T, d_head]，让每个头独立计算。这样做的原因是：如果每个头都用完整的C维做投影，参数量会变成H倍，而且计算量也爆炸。切分维度是性价比最高的方案。
我用一个表格把形状变化列清楚，这是理解多头注意力的关键：
| 步骤 | 操作 | 形状变化 |
|---|---|---|
| 输入 | X | [T, C] |
| QKV 投影 | X @ W_qkv | [T, 3C] |
| 拆分 | reshape + transpose | [3, H, T, d_head] |
| 注意力分数 | Q @ K^T | [H, T, T] |
| 缩放+掩码+softmax | - | [H, T, T] |
| 加权求和 | attn @ V | [H, T, d_head] |
| 合并头 | transpose + reshape | [T, C] |
| 输出投影 | @ W_o | [T, C] |
注意 llm.c 里为了效率，通常把W_q、W_k、W_v合并成一个[C, 3C]的大矩阵，一次矩阵乘法搞定三次投影。这是个很实用的工程技巧，减少 kernel 启动次数，也提高内存访问的连续性。你在自己实现的时候可以分开写，逻辑更清晰，但生产环境建议合并。
2.3 缩放因子 1/sqrt(d_head)：不是玄学，是方差控制
1/sqrt(d_head)这个缩放因子，很多人背下来了但不知道为什么。我用方差的角度解释一遍。
假设 Q 和 K 的每个元素都是独立同分布，均值 0，方差 1。那么点积q · k = sum(q_i * k_i)的方差是多少？每个q_i * k_i的方差是1 * 1 = 1，d_head个独立项相加，方差就是d_head。所以点积结果的方差是d_head，标准差是sqrt(d_head)。
当d_head = 64时，点积结果的标准差是 8。如果不缩放，softmax 的输入值域会很大，导致 softmax 输出接近 one-hot，梯度几乎消失。除以sqrt(d_head)之后，方差被拉回 1，softmax 的输入分布比较温和，梯度更健康。
llm.c 里这个缩放通常直接乘在 Q 上，或者乘在 scores 上，效果一样。我建议乘在 Q 上，因为这样只需要对[T, C]做一次缩放，而不是对[H, T, T]做，省计算。不过 llm.c 为了代码清晰，可能是在 scores 上乘的，你看具体实现。
提示：如果你自己写的时候忘了缩放，训练初期 loss 会下降得很慢甚至不降，因为 softmax 饱和了。这个坑我踩过，排查了半天才发现是漏了
1/sqrt(d_head)。
2.4 因果掩码：为什么不能简单置零
GPT 是自回归模型，预测第t个 token 时只能看到前t个 token，不能看到未来。所以在注意力分数矩阵[T, T]里，位置(i, j)当j > i时必须被屏蔽掉。
最常见的做法是把这些位置设成一个很大的负数（比如-1e9），这样 softmax 之后这些位置的权重接近 0。为什么不能直接设成 0？因为 softmax 的公式是exp(x_i) / sum(exp(x_j))，如果设成 0，exp(0) = 1，反而会分到权重。必须设成负无穷（实践中用大负数近似），exp(-1e9) ≈ 0，才能真正屏蔽。
llm.c 里通常用一个预计算的掩码矩阵，或者在上三角区域直接写-INFINITY。我建议用预计算的方式，因为掩码是固定的，没必要每次前向都重新算。另外要注意，掩码是在 softmax 之前加的，不是之后。
还有个细节：softmax 的数值稳定性。如果一行里所有值都是-1e9（理论上不会，因为对角线j=i是允许的），或者最大值太大，直接exp会溢出。标准做法是先减去每行的最大值再exp。llm.c 里应该做了这个处理，你写的时候别忘了。
3. MLP 前馈网络：4 倍膨胀背后的设计逻辑
3.1 两层线性 + GELU：结构虽简单，位置很关键
Transformer 每个 block 里，注意力后面跟着一个 MLP（也叫 FFN，前馈网络）。GPT-2 的 MLP 结构是：Linear(C, 4C) -> GELU -> Linear(4C, C)。输入[T, C]，先升维到[T, 4C]，激活后再降回[T, C]。
为什么中间要膨胀 4 倍？这个问题没有标准答案，但有几个合理的解释。第一，注意力层主要负责 token 之间的信息交换，MLP 负责对每个 token 独立做非线性变换，膨胀给模型更大的容量去存储知识。第二，4 倍是个经验值，实验发现比 2 倍或 8 倍效果好。第三，从参数量看，MLP 占了整个 block 参数量的三分之二左右（2 * C * 4C = 8C^2，而注意力是4C^2），所以 MLP 其实是知识存储的主力。
llm.c 里这个 MLP 的实现很直接：两个矩阵乘法加一个 GELU。GELU 的公式是x * Phi(x)，其中Phi是标准正态分布的累积分布函数。实践中常用近似：0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))。llm.c 里应该用的是这个 tanh 近似，因为它比精确的 erf 实现快。
3.2 GELU vs ReLU：为什么 GPT-2 选了 GELU
ReLU 是max(0, x)，简单粗暴，负数直接归零。GELU 是平滑的，负数区域不是硬归零，而是有一个小的负值。这个平滑性带来两个好处：一是梯度更平滑，训练更稳定；二是负值区域保留了一点信息，表达能力更强。
我用一个生活类比：ReLU 像是一个严格的老师，答错了直接零分；GELU 像是一个温和的老师，答错了给个低分但不是零分，学生还能从错误中学到东西。当然这个类比不严谨，但能帮你记住区别。
实测下来，GELU 在 Transformer 类模型上确实比 ReLU 好一点点，但差距不大。如果你自己搭模型，用 ReLU 也能跑，只是可能收敛慢一点。llm.c 忠实还原 GPT-2，所以用 GELU。
3.3 MLP 里的残差和 LayerNorm：顺序问题
MLP 本身不包含残差和 LayerNorm，但它在 block 里的位置和这两个组件紧密相关。GPT-2 用的是 Pre-LN 结构：x = x + Attention(LN(x))，然后x = x + MLP(LN(x))。也就是说，LayerNorm 在注意力和 MLP 之前，残差连接把输入直接加到输出上。
这个顺序很重要。Post-LN（原始 Transformer 论文用的）是x = LN(x + Attention(x))，LayerNorm 在残差之后。Pre-LN 训练更稳定，不需要 warmup 也能收敛，所以 GPT-2 之后的大部分模型都用 Pre-LN。llm.c 里你注意看 LayerNorm 的调用位置，就能确认是哪种。
注意：如果你把 Pre-LN 和 Post-LN 搞混，训练可能会发散。我建议自己实现时先用 Pre-LN，稳。
4. LayerNorm 的归一化维度：最容易搞错的地方
4.1 LayerNorm 到底对哪个维度做归一化
这是初学者最容易晕的地方。输入形状[T, C]，LayerNorm 是对每个 token 的C维特征做归一化，也就是对最后一维做。每个 token 独立计算均值和方差，然后归一化，再乘gamma加beta。gamma和beta的形状都是[C]，对所有 token 共享。
对比一下 BatchNorm：BatchNorm 是对 batch 维度做归一化，每个特征维度独立。在 NLP 里，序列长度可变，batch 里不同样本长度不一样，用 BatchNorm 很麻烦。LayerNorm 不依赖 batch，每个样本独立，所以更适合 NLP。
我用一个表格对比：
| 维度 | LayerNorm | BatchNorm |
|---|---|---|
| 归一化维度 | 特征维 C | Batch 维 |
| 统计量来源 | 单个样本 | 整个 batch |
| 是否依赖 batch | 否 | 是 |
| 训练/推理一致性 | 一致 | 不一致 |
| 适合场景 | NLP、序列 | CV、固定 batch |
llm.c 里 LayerNorm 的实现通常是：对每个 token，遍历C个元素算均值和方差，然后归一化。因为C=768，这个循环不算大，但T可能很长，所以总体计算量不小。优化的话可以用向量化指令，但 llm.c 为了可读性可能用朴素循环。
4.2 均值方差的计算：数值稳定性不能忽视
LayerNorm 的公式是y = (x - mean) / sqrt(var + eps) * gamma + beta。eps是个很小的数（通常1e-5），防止除零。方差的计算有两种方式：一种是E[(x - mean)^2]，另一种是E[x^2] - mean^2。后者只需要遍历一次，但数值稳定性差，当x的均值很大时，两个大数相减会损失精度。llm.c 里我建议用前者，遍历两次但更稳。
还有个细节：gamma和beta是可学习参数，初始化时gamma=1，beta=0。这样初始状态下 LayerNorm 是恒等变换（除了归一化），不会破坏输入分布。
4.3 LayerNorm 在 llm.c 里的内存布局
llm.c 用连续内存存张量，[T, C]的布局是行优先，第t个 token 的C个特征连续存放。LayerNorm 遍历时，外层循环t，内层循环c，这样内存访问是连续的，缓存友好。如果你反过来写，性能会差很多。
gamma和beta是[C]的一维数组，对所有 token 共享。计算时，内层循环里同时读x[t*C + c]、gamma[c]、beta[c]，三个数组都是连续访问，效率高。
5. 把注意力、MLP、LayerNorm 串起来：一个 block 的完整前向
5.1 数据流梳理：从输入到输出
现在把这三块串起来。一个 GPT-2 block 的前向过程是：
- 输入
x，形状[T, C] ln1 = LayerNorm(x)，形状[T, C]attn = MultiHeadAttention(ln1)，形状[T, C]x = x + attn，残差连接ln2 = LayerNorm(x)，形状[T, C]mlp = MLP(ln2)，形状[T, C]x = x + mlp，残差连接- 输出
x，形状[T, C]
这个结构重复N层（GPT-2 small 是 12 层），最后接一个 LayerNorm 和输出投影。
注意残差连接的作用：它让梯度可以直接回传，缓解深层网络的梯度消失。没有残差，12 层堆起来训练会很困难。llm.c 里残差就是简单的逐元素加法，但要注意原地操作还是新分配内存。原地操作省内存，但可能覆盖还需要用的输入，要小心。
5.2 参数量核算：钱花在哪了
我算一下 GPT-2 small 一个 block 的参数量，让你对"钱花在哪"有个概念：
- 注意力 QKV 投影：
C * 3C = 768 * 2304 = 1,769,472 - 注意力输出投影：
C * C = 589,824 - MLP 第一层：
C * 4C = 768 * 3072 = 2,359,296 - MLP 第二层：
4C * C = 3072 * 768 = 2,359,296 - LayerNorm 两个：
2 * 2C = 3072
加起来一个 block 约 708 万参数，其中 MLP 占了 471 万，约三分之二。12 层就是 8500 万左右，加上嵌入层和输出层，总共约 1.24 亿参数。这就是 GPT-2 small 的规模。
这个核算告诉你：如果你想压缩模型，MLP 是重点；如果你想提升容量，加 MLP 宽度比加注意力头数更划算。
5.3 常见实现错误清单
我在实现和调试过程中踩过的坑，列出来帮你避雷：
- 形状对不上：多头 reshape 时
H * d_head必须等于C，否则 transpose 后维度错乱。 - 掩码位置错：掩码必须在 softmax 之前加，且用大负数不是零。
- 缩放漏掉：
1/sqrt(d_head)忘了乘，训练不收敛。 - LayerNorm 维度错：对
T维归一化而不是C维，结果完全不对。 - 残差覆盖：原地加法把还需要用的
x覆盖了，导致后续计算错误。 - GELU 近似公式错：tanh 里的系数写错，激活值偏差大。
- softmax 溢出：没减最大值，
exp溢出成 inf。
这些错误单看都很低级，但在 C 代码里没有形状检查，一旦写错就是静默错误，loss 不降你都不知道哪错了。建议每写一个模块，先用小规模数据（比如T=4, C=8, H=2）手动验证输出。
6. 调试与验证：怎么确认你的实现是对的
6.1 用小规模数据做单元测试
llm.c 这种底层实现，最怕的就是"看起来对但实际错"。我的做法是：构造极小的输入，手动算一遍期望输出，然后对比。比如T=2, C=4, H=2, d_head=2，输入用简单的整数，手动算注意力分数、softmax、加权求和，然后和代码输出对比。
LayerNorm 也可以这样测：输入[1, 2, 3, 4]，均值 2.5，方差 1.25，归一化后是[-1.34, -0.45, 0.45, 1.34]（近似），乘 gamma 加 beta 后对比。
6.2 梯度检查：数值梯度 vs 解析梯度
如果你实现了反向传播，梯度检查是必须的。用数值梯度(f(x+h) - f(x-h)) / 2h和解析梯度对比，误差应该在1e-5量级。llm.c 如果只做前向推理，可以跳过这步，但如果你想训练，这步不能省。
6.3 和参考实现对比
最靠谱的验证是和一个已知正确的实现对比。你可以用 PyTorch 搭一个同样的 GPT-2 block，加载相同的权重，输入相同的数据，对比每一层的输出。如果每一层都对得上，那你的 C 实现就是对的。这个过程中你会定位到具体是哪一层出了问题，比盲猜高效得多。
提示：对比时注意浮点精度，C 的 float 和 PyTorch 的 float32 可能有微小差异，允许
1e-4左右的误差。
7. 一些实战心得与性能考量
7.1 内存布局决定性能
llm.c 用连续内存，矩阵乘法的手写实现里，循环顺序对性能影响巨大。比如C = A @ B，A是[M, K]，B是[K, N]，C是[M, N]。如果按i, j, k循环，内层访问B[k*N + j]是跨行访问，缓存不友好。改成i, k, j，内层访问B[k*N + j]连续，性能能差好几倍。
注意力里的Q @ K^T也是同理。K^T不是真的转置，而是通过索引访问实现，循环顺序要配合。
7.2 多头注意力的并行化思路
llm.c 是单线程 C，但多头天然可以并行，因为每个头独立。如果你用 OpenMP，可以在头维度上并行，#pragma omp parallel for加在h循环上。不过要注意每个头写回输出的位置不能冲突，或者用 reduction。
7.3 什么时候该用融合算子
生产环境里，QKV 投影通常融合成一个矩阵乘法，注意力分数和 softmax 也可能融合。llm.c 为了教学清晰，可能分开写。你自己实现时，先保证正确，再考虑融合。融合能减少内存读写，但调试难度大。
7.4 关于数值精度的取舍
GPT-2 原始是 float32，但推理时可以用 float16 甚至 int8 量化。llm.c 如果做 float32，精度没问题，但内存占用大。如果你想省内存，可以试试混合精度，但要注意 LayerNorm 和 softmax 对精度敏感，最好保持 float32。
我在实际使用中发现，LayerNorm 的eps取值对结果有影响，1e-5是 GPT-2 的默认值，但如果你自己训练，可以试试1e-6或1e-4，看哪个收敛更好。这个没有定论，得实验。
最后再分享一个小技巧：如果你在调试注意力时发现输出全是 NaN，先检查 softmax 有没有减最大值，再检查掩码有没有用大负数，最后检查缩放因子。这三个地方是 NaN 的高发区。按这个顺序排查，基本能定位问题。
原理与实战：从基础到应用)
