MiniMind 学习笔记之 04 注意力之外：位置、记忆、省显存、省时间、深加工Flash Attention、RoP

上承〔03 · 自注意力：模型怎么“看懂”上下文？〕。上一篇把自注意力拆完了：Q、K、V 怎么来，注意力权重怎么算，多头怎么并行。但一个完整的 Transformer 块，不是只有注意力。注意力只是“让 token 互相看”的那一半。

03 篇源码走读时，有几处只点了名、没展开。其中四个——Flash Attention、RoPE、KV Cache、GQA——是注意力的配套：位置、记忆、省显存、省时间。它们在源码里出现，但每个都是独立的工程话题，放在源码走读的语境里讲不透。这一篇把这四个逐一展开。

注意力之外，还有三块没讲过：FFN 负责每个 token 的深加工，残差和 RMSNorm 保证深层能训得动，最后这些零件怎么拼成一个完整的块。这一篇一并补上。

六块拼完，一个完整的 Transformer 块就出来了。

一、位置：RoPE

1.1 自注意力不知道谁在前，谁在后

第三篇讲注意力的时候，一直回避了一个问题。

自注意力的计算没有依赖链。每个 token 的 Q、K、V 都是独立算出来的，每个位置的注意力都是独立做的匹配和加权。这种独立性让它可以并行——训练时，整个序列一次性算完；推理时，当前 token 和缓存里所有历史 token 一次性匹配。和 RNN 不同，它不需要等前一个 token 算完才能开始算下一个。

这意味着什么？如果打乱“我”、“爱”、“你”的顺序，注意力机制算出来的结果是一样的。“我爱你”和“你爱我”，在模型眼里暂时没有区别。

但语言是有顺序的。“猫追老鼠”和“老鼠追猫”，字一样，意思相反。模型要能区分，就必须给每个 token 一个信号：你在第几个位置。

这个信号，就是位置编码。

1.2 从正弦编码到 RoPE

大模型刚出来时，位置编码用的是正弦编码（Sinusoidal Positional Encoding），2017 年《Attention Is All You Need》那篇论文提出的。

做法很直接：给每个位置的 token，在它的输入向量上加一个固定的向量。这个向量用正弦和余弦函数算出来，每个位置不同，每个维度不同。

$\text{输入} = \text{embedding} + PE(pos)$

后来换了方案。2021 年，苏剑林等人提出了旋转位置编码（RoPE，Rotary Position Embedding）。目前主流大模型——Llama、Qwen、DeepSeek、MiniMind——用的都是 RoPE。

两者最容易被误解的区别，是“加 vs 乘”。 这个说法不算错，但没说到点上。两者都是把位置信息作用到向量上，真正的区别在下面三处。

第一，作用对象不同。

正弦编码作用在 embedding 上。token 一进来就被加了位置信息，然后才进 Transformer 层。

RoPE 作用在 Q 和 K 上。token 先经过投影变成 q、k，然后才被旋转。位置信息是“注意力计算前”才加的，不是“进入模型时”加的。

第二，绝对 vs 相对。

正弦编码标记的是绝对位置。位置 5 的 token，加的是位置 5 对应的那个固定向量。模型学到的位置关系，依赖绝对位置的取值。

RoPE 编码的是相对位置。两个 token 的 q 和 k 都旋转之后，它们的点积只和旋转角度之差有关——也就是只和它们的距离有关。位置 5 和位置 6 的距离是 1，位置 105 和位置 106 的距离也是 1，RoPE 算出来的点积是一样的。

这是 RoPE 最核心的优势：相对位置编码，自然融进点积里。

第三，数学性质不同。

正弦编码加到 embedding 上之后，注意力分数 $q \cdot k$ 展开会有四项：内容·内容、内容·位置、位置·内容、位置·位置。位置项和内容项纠缠在一起，模型得自己理清。

RoPE 旋转之后， $q' \cdot k'$ 有一个漂亮的数学性质：只和相对角度有关。位置信息不是“混进去的噪音”，而是“自然地变成了相对距离的信号”。

1.3 RoPE 的核心思路

RoPE 不往向量里“加”位置，而是把向量旋转一个角度。

角度由位置决定。位置越靠后，旋转角度越大。

具体怎么做？把 Q 和 K 向量里的数两两分组，每组两个数 $(x_1, x_2)$ ，看作一个二维平面上的点。然后把这个点按一个角度旋转：

\begin{pmatrix} x_1' \\ x_2' \end{pmatrix} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix} \begin{pmatrix} x_1 \\ x_2 \end{pmatrix}

旋转角度 $\theta$ 和 token 的位置 $m$ 有关： $\theta = m \cdot \omega$ 。位置越靠后，旋转越多。

每一对分配一个不同的频率 $\omega$ 。低频对旋转得慢，高频对旋转得快。

这里有一个容易搞混的地方：RoPE 作用在多头的每个头上。

一个 token 的向量经过 q_proj 投影后，变成 768 维，然后reshape成 8 个头，每个头 96 维。RoPE 对每个头单独做——每个头的 96 维拆成 48 对，每对分配一个频率，各自旋转。

【q_proj 这个 768 × 768 的矩阵，内部可以理解为 8 个 768 × 96 的子矩阵拼在一起。每个子矩阵负责一个头。也就是说，输入 x（768 维）和每个子矩阵相乘，得到 96 维的输出。8 个子矩阵的输出拼起来，就是 768 维。然后 reshape 成 [8, 96]，每个头拿到自己的 96 维。】

8 个头共用同一张频率表。位置 5 的 token，在 8 个头里的旋转角度是一样的（都是 $5 \times \omega_i$ ），但每个头的 96 维内容是各自投影出来的，所以旋转后的结果各不相同。

1.3.1 RoPE 改不改 QKV 矩阵的参数？

这个问题值得单独说清楚，因为它容易引起误解。

先说结论：RoPE 不改 QKV 矩阵的参数，但会影响它们被训练成什么样。

第一层：RoPE 不修改 QKV 的权重。

QKV 矩阵是 nn.Linear 层，它们的权重 $W_Q$ 、 $W_K$ 、 $W_V$ 是训练学出来的。RoPE 是一个固定的数学变换，作用在 QKV 投影的输出上，不动 QKV 自己的权重。

流程是：

x → q_proj → q → RoPE 旋转 → q'

q_proj 的参数在第一步就被用掉了。RoPE 只在第二步作用，它没有自己的可学习参数，也不修改 q_proj 的权重。

第二层：RoPE 和 QKV 在训练中协同。

它们在训练中协同，但角色不同。

QKV 矩阵负责“把输入变成 Q、K、V”。它们学的是“什么样的内容该被提取出来”。
RoPE 负责“把位置信息编进去”。它不学，是固定的旋转。

训练时，反向传播的梯度会穿过 RoPE，一路传回 QKV 矩阵。也就是说，虽然 RoPE 自己不学，但它会影响 QKV 矩阵被训练成什么样。

举个例子。模型学到“q 和 k 的点积要能反映距离”。这个学习目标同时约束了两件事：

QKV 矩阵要产出合适的 q、k，让它们旋转后能正确匹配。
RoPE 的旋转角度是固定的，QKV 矩阵得适应这套固定的旋转规则。

所以不是“RoPE 改变 QKV 的参数”，而是“QKV 的参数在 RoPE 固定的旋转规则下被训练”。RoPE 定规矩，QKV 学怎么在这个规矩下把事办好。

第三层：RoPE 有没有可学习的部分？

标准实现里，没有。频率 $\omega_i$ 是公式算出来的固定值，不参与训练。

但有一些变体让部分频率可学习，或者让频率缩放因子可学习。那是改进版，不是原始设计。

一句话总结：RoPE 不参与 QKV 矩阵的旋钮调节。它是一套固定的旋转规则，QKV 矩阵在训练中学会适应这套规则。梯度会穿过 RoPE 传回 QKV，所以它们协同工作，但 RoPE 自己不被训练。

1.4 为什么旋转能编码相对位置？

关键在点积。

两个 token 做注意力匹配时，算的是 $q \cdot k$ 。经过 RoPE 之后，q 和 k 都旋转了各自的角度。

数学上有一个漂亮的结论：两个旋转后的向量做点积，结果只和它们的旋转角度之差有关。

假设“猫”在位置 $m$ ，“坐”在位置 $n$ 。经过 RoPE 后，“猫”的 q 旋转了 $m\omega$ ，“坐”的 k 旋转了 $n\omega$ 。点积的结果，只取决于 $(m - n)\omega$ ，也就是两个 token 的距离。

“猫”在位置 2，“坐”在位置 3，距离是 1。“猫”在位置 100，“坐”在位置 101，距离也是 1。RoPE 算出来的注意力分数，对这两个场景是一样的。

1.5 用一个具体例子走一遍

前面讲了原理，现在拿具体数字走一遍。

沿用第三篇的分词假定：

他 | 把 | 苹果 | 吃了
位置 0 | 1 | 2 | 3

苹果在位置 2，吃了在位置 3。

假设 head_dim = 4（真实是 96，这里只留两对，方便手算）。每一对分配一个不同的频率：

第 1 对： $\omega_1 = 1.0$ （高频，转得快）
第 2 对： $\omega_2 = 0.1$ （低频，转得慢）

“苹果”的 q 和“吃了”的 k，初始都是最简单的向量：

苹果的 q = [1, 0, 1, 0]
吃了的 k = [1, 0, 1, 0]

前两个数是第 1 对，后两个数是第 2 对。

“苹果”在位置 2：

第 1 对旋转角度： $2 \times 1.0 = 2$ 弧度
第 2 对旋转角度： $2 \times 0.1 = 0.2$ 弧度

“吃了”在位置 3：

第 1 对旋转角度： $3 \times 1.0 = 3$ 弧度
第 2 对旋转角度： $3 \times 0.1 = 0.3$ 弧度

第 1 对（高频）：

苹果的 q 第 1 对 (1, 0)，旋转 2 弧度：

q₁' = [cos(2), sin(2)] ≈ [-0.416, 0.909]

吃了的 k 第 1 对 (1, 0)，旋转 3 弧度：

k₁' = [cos(3), sin(3)] ≈ [-0.990, 0.141]

点积：

q₁' · k₁' = (-0.416)(-0.990) + (0.909)(0.141) ≈ 0.540

第 2 对（低频）：

苹果的 q 第 2 对 (1, 0)，旋转 0.2 弧度：

q₂' = [cos(0.2), sin(0.2)] ≈ [0.980, 0.199]

吃了的 k 第 2 对 (1, 0)，旋转 0.3 弧度：

k₂' = [cos(0.3), sin(0.3)] ≈ [0.955, 0.296]

点积：

q₂' · k₂' = (0.980)(0.955) + (0.199)(0.296) ≈ 0.995

总点积：

q' · k' = 0.540 + 0.995 = 1.535

这个 1.535，就是“苹果”和“吃了”的注意力分数，包含位置信息。

换个距离，看分数怎么变。

把“吃了”挪到位置 9。距离从 1 变成 7。

第 1 对（高频）：

吃了的 k 第 1 对 (1, 0)，旋转 9 弧度：

k₁' = [cos(9), sin(9)] ≈ [-0.911, 0.412]

点积：

q₁' · k₁' = (-0.416)(-0.911) + (0.909)(0.412) ≈ 0.379 + 0.375 ≈ 0.754

点积只取决于角度差：k 转了 9 弧度，q 转了 2 弧度，差是 7 弧度。7 弧度绕了一圈多（一圈 ≈ 6.28 弧度），7 − 6.28 = 0.72 弧度——比距离 1 时的 1 弧度还小。高频对在这种情况下已经分不清“距离 7”和“距离 0.72”了。

第 2 对（低频）：

吃了的 k 第 2 对 (1, 0)，旋转 0.9 弧度：

k₂' = [cos(0.9), sin(0.9)] ≈ [0.622, 0.783]

点积：

q₂' · k₂' = (0.980)(0.622) + (0.199)(0.783) ≈ 0.610 + 0.156 ≈ 0.766

总点积：

q' · k' = 0.754 + 0.766 ≈ 1.520

距离 1 时总分是 1.535，距离 7 时总分是 1.520。非常接近。

问题出在哪？高频对绕圈了，把“距离 7”误判成了“距离 0.72”。低频对因为转得慢，角度差只从 0.1 变成 0.9，还在合理范围内，但它对总分的贡献被高频对的误判抵消了一部分。

多频率配合的意义：

如果只有高频对，距离一长就绕圈，分不清距离 1 和距离 7。

如果只有低频对，相邻位置的角度差太小（距离 1 时只有 0.1 弧度），区分不了相邻位置。

两个频率一起，高频负责近处，低频负责远处。单个频率会误判，但两个频率的“误判方式”不同。 模型可以通过多个频率的组合，反推出真实距离。

MiniMind 的 head_dim 是 96，分成 48 对。48 个频率从高到低，覆盖了从“相邻位置”到“整个序列长度”的所有尺度。

为什么是两两分组？

因为二维平面上的旋转有现成的公式。高维空间里的旋转没有这么简洁的公式，但拆成一对一对的二维平面，每一对独立旋转，就很好算。

以 MiniMind 为例：每个头的 96 维向量，拆成 48 对。每对分配一个不同的频率。48 个二维旋转拼起来，就是这个头的旋转；8 个头用同一张频率表各转各的，合起来就是整个 768 维向量的位置编码。

48 对各自旋转，向量会不会被拆得支离破碎？

不会。96 维向量在数学上就是 96 个独立的数。分成 48 对，对与对之间从来没有“绑在一起”过。

旋转有一个性质：它不改变向量的长度。一对向量 (x₁, x₂) 旋转 θ 度之后，长度还是 √(x₁² + x₂²)。旋转只改变方向，不改变大小。

48 对各自旋转之后：

每一对内部的长度不变
整个 96 维向量的总长度也不变
变化的只是每一对的方向

这不是把向量打散，而是让同一个向量在不同尺度上同时记录位置。 48 个不同频率的指针，各转各的，合起来就是完整的位置信息。

为什么只旋转 Q 和 K，不旋转 V？

回到注意力的完整公式：

\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V

第一步， $QK^T$ 算出匹配分数，Softmax 变成权重。这一步决定“谁看谁”，需要位置信息。所以 Q 和 K 必须带位置。

第二步，权重乘以 V，加权求和。这一步只关心“取什么内容”，不关心“这两个 token 隔多远”。位置信息在这里没有用处。

如果硬把位置信息也塞进 V，同一个内容出现在位置 3 和出现在位置 7，取回来的材料就不同了。但内容本身和位置无关——“吃了”这个动作，不管它在句子的哪个位置，它提供的“动作信息”应该是一样的。

位置只影响匹配，不影响内容。 所以只旋转 Q 和 K。

1.6 MiniMind 的 RoPE 实现

打开 MiniMind 的 model/model_minimind.py，RoPE 分两步。

第一步：预计算频率。

def precompute_freqs_cis(dim, end, rope_base, rope_scaling=None):
    freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
    return freqs_cos, freqs_sin

逐一拆解：

dim 是 head_dim，MiniMind 里是 96。
rope_base 是基频，MiniMind 用的是 1,000,000。LLaMA-1 用的是 10,000。基频越大，频率衰减越慢，长距离依赖效果越好。
torch.arange(0, dim, 2) 生成 0, 2, 4, ..., 94，共 48 个数。除以 dim 后，得到 48 个不同的频率。
t 是位置序列，从 0 到 end - 1。MiniMind 的 end 是 32768。
torch.outer(t, freqs) 得到一个 32768 × 48 的矩阵。第 $m$ 行第 $i$ 列，就是位置 $m$ 在第 $i$ 个频率上的旋转角度。
最后把 cos 和 sin 各拼一份，变成 32768 × 96。前 48 列和后 48 列相同，是为了和 rotate_half 配合。

第二步：应用旋转。

def rotate_half(x):
    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
    return torch.cat([-x2, x1], dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin):
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

rotate_half 的作用是：把向量分成前后两半，交换位置，前半取负。这正好对应了旋转矩阵里的 $-\sin$ 那一项。

q_embed = q * cos + rotate_half(q) * sin，就是旋转公式的向量化写法。

逐行拆代码。

第一行：算频率。

freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))

torch.arange(0, dim, 2) 生成 [0, 2, 4, ..., 94]，共 48 个数。除以 dim 后，得到 [0, 0.0208, 0.0417, ..., 0.979]。把上一步的结果作为指数，底数是 rope_base。

当指数是 0 时， $1000000^0 = 1$ ，频率是 $1/1 = 1$ 。当指数是 0.979 时， $1000000^{0.979} \approx 750000$ ，频率是 $1/750000 \approx 1.3 \times 10^{-6}$ 。

取倒数，得到频率。最终 48 个频率，从高到低：

ω₀ ≈ 1.0
ω₁ ≈ 0.75
ω₂ ≈ 0.56
...
ω₄₇ ≈ 1.3 × 10⁻⁶

第 0 对频率最高（1.0），转得最快。第 47 对频率最低（约 1.3e-6），转得最慢。

为什么基频是 1,000,000？ LLaMA-1 用的是 10,000，MiniMind 用的是 1,000,000。基频越大，低频部分的频率越低。低频对负责长距离。频率越低，绕圈越慢，能区分的距离越长。这是 MiniMind 支持 32768 上下文的一个基础。

第二行：位置序列。

t = torch.arange(end, device=freqs.device)

end 是 32768，所以 t = [0, 1, 2, ..., 32767]。

第三行：算角度矩阵。

freqs = torch.outer(t, freqs).float()

外积得到一个 32768 × 48 的矩阵。第 $m$ 行第 $i$ 列的数，是 $m \times \omega_i$ ——位置 $m$ 在第 $i$ 对上的旋转角度。

第四、五行：算 cos 和 sin。

freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)

对上一步的矩阵逐元素取 cos 和 sin，各得到一个 32768 × 48 的矩阵。然后把两个同样的矩阵拼起来，变成 32768 × 96。

为什么要拼一份重复的？

因为 rotate_half 的写法。它把 96 维的向量分成前后两半（各 48 维），交换位置，前半取负。RoPE 的旋转公式是：

q' = q \odot \cos + \text{rotate\_half}(q) \odot \sin

要让这个式子对 96 维向量成立，cos 和 sin 也必须是 96 维的。前半段对应“原向量”，后半段对应“rotate_half 之后的向量”。所以 cos 拼一份重复的，让前后两半都用同一组值。

为什么要取整到 64 的倍数？ 为了让矩阵维度对齐硬件的计算友好度。64 是 GPU 上矩阵乘法的常见对齐粒度。

为什么叫 hidden？

这个名字来自早期神经网络的分层叫法。网络分三层：输入层、隐藏层、输出层。

输入层：直接接收外部数据。
输出层：直接给出最终结果。
隐藏层：夹在中间的那些层，既不直接接触输入，也不直接给出输出。

“hidden”的意思是“不直接面对外部”，不是“藏起来看不见”。隐藏层是模型真正做计算的地方，只是它的内部状态不直接暴露给用户。

hidden_size 就是隐藏层里向量的维度。MiniMind 选了 768，所以整条链路上流动的向量，长度都是 768。

这个维度和头数的关系是：

\text{hidden\_size} = \text{头数} \times \text{head\_dim}

MiniMind 有 8 个 Q 头，每个头分到 768 ÷ 8 = 96 维。这个 96 就是 head_dim。

1.7 外推：训练没见过那么长怎么办

模型一次处理的 token 数量，是固定的。训练时，每条训练数据都切成几百个 token，模型一次看这么多。这个量级就是训练时的序列长度。

但推理的时候，用户可能输入一段很长的文本，比如 4096 个 token。模型要处理比训练时更长的序列。

RoPE 是靠旋转角度标记位置的。位置 0 旋转 0 度，位置 1 旋转 $\omega$ 度，位置 2 旋转 $2\omega$ 度……训练时，模型只见过位置 0 到几百的旋转角度。位置 3000 该旋转多少度，公式能算出来，但模型没见过这个角度对应的模式，它不知道该怎么处理。

所以需要长度外推：想办法让训练时只见过几百的模型，能处理 4096 甚至更长的序列。

MiniMind 支持 YaRN（Yet another RoPE extensioN）做长度外推。

YaRN 的思路是：不改模型，改旋转频率。

具体做法：把旋转频率分档。高频部分（对应短距离位置关系）保持原样，低频部分（对应长距离位置关系）做插值——把训练时没见过的那些“远距离角度”，压缩到训练时见过的范围内。

这样，模型遇到长序列时，旋转角度不会超出它熟悉的模式，效果就保住了。

MiniMind 的配置里，max_position_embeddings=32768，通过 YaRN 可以把有效上下文从几百扩展到 32768。控制开关是 inference_rope_scaling 标志。

二、记忆：KV Cache

2.1 推理时的浪费

模型推理的时候，是一个词一个词往外吐的。

第一次：输入“今天”，算出“天气”。第二次：输入“今天天气”，算出“真”。第三次：输入“今天天气真”，算出“好”。

每次都要重新算一遍前面所有 token 的 k 和 v。但前面 token 的 k 和 v 在第一次就算过了，不会因为后面新增了 token 而改变。

生成第 n 个词时，要重算前面 n-1 个 token 的 k 和 v。生成一个长度为 N 的序列，总计算量是 O(N²)。

2.2 缓存的做法

解决办法：把算过的 k 和 v 缓存起来，下次直接用。

第一次算“今天”，得到“今天”的 k 和 v，存起来。第二次算“天气”，只需要算“天气”的 k 和 v，然后把缓存的“今天”的 k、v 拼在前面。第三次同理。

每次只需要算当前 token 的 k 和 v，前面所有 token 的 k、v 从缓存里取。

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)

有了 KV Cache，每生成一个 token 的计算量从 O(n) 降到 O(1)。生成长度为 N 的序列，总计算量从 O(N²) 降到 O(N)。

2.3 缓存的代价

KV Cache 的大小 = 2 × 层数 × KV 头数 × head_dim × 序列长度 × 字节数。

MiniMind 的配置：8 层，4 个 KV 头，head_dim 96，fp16 存储。序列长度 2048 时：

2 × 8 × 4 × 96 × 2048 × 2 ≈ 25MB

一条序列 25MB。如果同时处理 100 条序列，就是 2.5GB。

2.4 为什么只缓存 K 和 V，不缓存 Q

注意力分数的计算方式是：当前 token 的 q，和所有 token 的 k 做点积。

生成第 1 个 token 时，用“今天”的 q₁，和 k₁ 做点积。生成第 2 个 token 时，用“天气”的 q₂，和 k₁、k₂ 做点积。

每一步，只用当前这个 token 的 q。前面 token 的 q 用完了就不再用了。所以 q 不需要缓存。

K 和 V 不一样。k₁ 和 v₁ 在第 1 步算出来，第 2 步还要用，第 3 步还要用，一直用到序列结束。所以必须缓存。

2.5 K 要参与位置编码，为什么前面算过的 K 不变

RoPE 给 k 加位置信息，方式是旋转。旋转角度由 token 的位置决定。

“今天”在位置 0，它的 k 旋转 0 度。 “天气”在位置 1，它的 k 旋转 1 度。

位置是绝对的。 “今天”永远在位置 0，不会因为后面新增了“天气”而变到位置 1 去。

所以“今天”的 k，经过 RoPE 旋转后，永远是同一个向量。缓存的是“已经应用了 RoPE 的 k”。 取出来直接用，不需要重新加位置编码，因为位置没变。

这就是 KV Cache 能成立的根本原因：位置不变，K 就不变。

三、省显存：GQA

3.1 三种注意力：MHA、MQA、GQA

MHA、MQA、GQA，说的是 Q 头和 KV 头之间的数量关系。

MiniMind 有 8 个 Q 头，每个头 96 维。

MHA：8 个 Q 头，8 个 KV 头。每个 Q 头配一组自己的 K 和 V。质量最好，KV Cache 最大。

MQA：8 个 Q 头，1 个 KV 头。所有 Q 头共用一套 K 和 V。KV Cache 最小，质量有损失。

GQA：8 个 Q 头，4 个 KV 头。两个 Q 头一组，共享一组 K 和 V。这是 MiniMind 的方案。

3.2 K 和 V 被投影降维了

一个 token 的向量 x 是 768 维。它进来之后，分三条路走：

x (768 维)
  │
  ├─ q_proj → q (768 维)
  ├─ k_proj → k (384 维)
  └─ v_proj → v (384 维)

Q 没有降维，K 和 V 降了。

q_proj 的形状是 768 × 768，输出 768 维。8 个 Q 头，每头 96 维。

k_proj 和 v_proj 的形状是 768 × 384，输出只有 384 维。4 个 KV 头，每头 96 维。

降的是总维度，不是每个头的维度。 每个 KV 头还是 96 维，和 Q 头一样。变的是头的数量：Q 有 8 个头，K 和 V 只有 4 个。

不是模型把算出来的 768 维 K、V 压缩了，而是 k_proj 和 v_proj 这两个矩阵从一开始就只输出 384 维。投影矩阵本身就是 768 × 384，不是 768 × 768。

3.3 具体走一遍：“苹果”的 Q 和“吃了”的 K 怎么匹配

拿“他把苹果吃了”举例。假设“苹果”在位置 2，“吃了”在位置 3。

“苹果”的 Q 有 8 个头，每头 96 维。“吃了”的 K 有 4 个头，每头 96 维。

GQA 的匹配方式：

苹果的 Q 头 1 和 Q 头 2，都去和吃了的 K 头 1 匹配。
苹果的 Q 头 3 和 Q 头 4，都去和吃了的 K 头 2 匹配。
苹果的 Q 头 5 和 Q 头 6，都去和吃了的 K 头 3 匹配。
苹果的 Q 头 7 和 Q 头 8，都去和吃了的 K 头 4 匹配。

同一组内的两个 Q 头，看到的是同一套 K 和 V。 它们提问的角度不同（各自的 Q 不同），但被匹配的“名片”是一样的。

3.4 GQA 省了多少

如果 KV 头数等于 Q 头数（MHA），MiniMind 的 KV Cache 是：

2 × 8 × 8 × 96 × 2048 × 2 ≈ 50MB

用 GQA，KV 头数从 8 降到 4，KV Cache 直接减半，变成 25MB。

四、省时间：Flash Attention

4.1 GPU 的内存墙

Flash Attention 是一种针对注意力机制的计算优化技术，由斯坦福大学 Tri Dao 等人在 2022 年提出。

要理解它，得先知道 GPU 的内存是分层的：

SRAM（高速缓存）：容量很小（约 20MB），但读写速度极快（约 19TB/s）。
HBM（高带宽显存）：容量很大（数十 GB），但读写速度相对慢得多（约 1.5TB/s）。

标准的注意力计算会生成一个 N×N 的注意力矩阵（N 是序列长度）。当序列很长时，这个矩阵非常庞大。传统实现会把这个大矩阵在慢速的 HBM 和快速的 SRAM 之间来回搬运，绝大部分时间浪费在数据搬运上。

4.2 分块计算：不让大矩阵完整出现

Flash Attention 的核心思路：不要让那个 N×N 的大矩阵完整地出现在 HBM 里。

序列长度 2048 时，N×N 就是 2048 × 2048 = 419 万个数字。fp16 存储，就是 8MB。

SRAM 只有 20MB，理论上放得下。但 SRAM 还要放 Q、K、V 的中间结果，留给这个矩阵的空间并不多。序列再长一点，比如 4096，矩阵就变成 32MB，SRAM 彻底放不下了。

Flash Attention 的做法：把这个大矩阵切成小块，一块一块地算。

还是 2048 × 2048 的矩阵，切成 64 × 64 的小块。一块只有 4096 个数字，8KB。每次只把一块加载到 SRAM 里，算完这块的输出，扔掉，再加载下一块。

但这里有一个问题：Softmax 需要看到一整行的所有分数，才能算归一化。切成小块之后，每一块只看到了一部分分数，怎么算 Softmax？

Flash Attention 用了一个数学技巧：在线 Softmax。它一边遍历小块，一边维护当前的最大值和累加和，最后再统一归一化。这个技巧保证了结果和标准注意力完全一致，不是近似。

4.3 重计算：以计算换内存

反向传播时，需要用到前向算出的注意力矩阵来算梯度。

标准实现会把前向的注意力矩阵存下来，反向时直接读。但 Flash Attention 没有把这个矩阵完整地算出来，也没存。

答案是：重新算一遍。 反向传播走到注意力这一步时，Flash Attention 拿着已经存下来的 Q、K、V，重新执行一遍分块计算。

这是典型的“以计算换内存”。代价是训练时多算一遍；收益是显存占用大幅降低。

4.4 效果

通过分块计算和重计算，Flash Attention 把注意力对 HBM 的访问次数大幅降低，显存占用从 O(N²) 降为 O(N)。在 GPT-2 上训练速度可提升 3 倍，并能支持长达 64K 的序列长度。

4.5 MiniMind 怎么用 Flash Attention

语言模型有一个额外的需求：因果掩码。

MiniMind 的注意力实现里，核心的注意力计算调用是这样的：

# 手工路径（SDPA 不可用时）：上三角加 -inf
scores[:, :, :, -seq_len:] += torch.full((seq_len, seq_len), float("-inf"), device=scores.device).triu(1)
# Flash 路径（默认）：因果掩码就是 kernel 的一个开关
attn_output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=self.is_causal)

is_causal=True：kernel 内部自动遮住未来位置，不需要外部的掩码矩阵。

分工是：

分块计算、在线 Softmax、重计算：Flash Attention 默认就做，使用者不用管。
因果掩码：通过 is_causal=True 告诉它要不要做。

五、深加工：FFN

5.1 FFN 是什么，为什么叫“前馈”

FFN 的全称是 Feed-Forward Network，中文叫前馈网络。

“前馈”这个词，意思是信息只向前流动，不回头。和它相对的是“循环神经网络”（RNN），RNN 的信息会从后面的步骤回传到前面的步骤。FFN 不循环，数据进去、出来，就完事了。

在 Transformer 的语境里，FFN 特指 Transformer 块里注意力后面的那个子层。它的工作对象是每个 token 的向量，一个一个独立处理——token 之间不发生任何信息交互。

注意力：token 和 token 之间交换信息。横向的。
FFN：每个 token 自己消化信息。纵向的。

一句话：注意力负责开会讨论，FFN 负责会后自己消化。

5.2 为什么需要一个“非线性”的加工

注意力有一个根本性的限制：它整体是线性的。

注意力做的事，是把所有 token 的 V 按权重加起来。加权求和是线性操作。

线性变换有一个致命弱点：叠多少层，整体还是线性的。 假设两层线性变换， $y = W_2(W_1 x)$ ，展开就是 $y = (W_2 W_1)x$ ，等价于一个矩阵。叠 100 层也一样。

纯线性模型没法处理“如果这个数大于 0，就往一个方向走；否则往另一个方向走”这种带条件的判断。而语言里到处是这种判断。

所以注意力之后，必须有一个地方引入非线性。这就是 FFN 的任务。

5.3 激活函数：非线性的来源

引入非线性的关键，是激活函数。

ReLU 的公式是 $\max(0, x)$ 。负数砍成 0，正数原样通过。简单、快、有效，但 0 点不可导，而且“硬切”太粗暴。

SiLU 的公式是 $x \cdot \sigma(x)$ ，其中 $\sigma(x)$ 是 Sigmoid。负数区域缓慢趋近 0，正数区域近似线性增长，0 附近有一个轻微的下凹。它光滑可导，训练更稳定，而且保留了负数区域的微弱信号。

SiLU 是“光滑版”的 ReLU，现代 LLM 更偏爱它。

5.4 FFN 在干什么：升维、非线性、降维

最简单的一层 FFN，是两层线性加一个激活：

\text{FFN}(x) = \text{ReLU}(x W_1) W_2

第一步：升维。 $W_1$ 把 768 维的 x 升到更高维，比如 3072 维。维度越高，能表达的模式越复杂。768 维升到 3072 维，模型有了一块更大的“画布”。

第二步：非线性。 ReLU（或 SiLU）作用在每一个维度上，切割空间。

第三步：降维。 $W_2$ 把 3072 维压回 768 维。

升维 → 非线性切割 → 降维，这就是 FFN 做的全部事情。

5.5 SwiGLU：给 FFN 加一道门

ReLU FFN 有两层。现代 LLM 换成了 SwiGLU，有三层：

\text{FFN}(x) = \big(\text{SiLU}(x W_{\text{gate}}) \odot x W_{\text{up}}\big) W_{\text{down}}

三道工序，各司其职。

up_proj：把材料摊开。 它把 768 维的向量升到 2432 维，相当于把原来压缩在一张纸上的信息，摊开到一张更大的桌子上。768 维里混在一起的“名词信息”“动作信息”“修饰信息”，摊到 2432 维之后，各自有机会落在不同的位置上。

gate_proj：决定放行哪些。 它同样把 768 维升到 2432 维，但它的任务不是“摊开材料”，而是出一份放行名单。它输出一串门控值，每个值对应 2432 维里的一个位置。

举个例子。某个维度上：

gate_proj 输出 2.0，经过 SiLU 约等于 1.76。门开得很大。
up_proj 在同一个维度输出 3.0，这是原始材料。
两者相乘：1.76 × 3.0 = 5.28。信息几乎全量通过。

换个维度：

gate_proj 输出 -2.0，经过 SiLU 约等于 -0.24。门几乎关死，还翻了符号。
up_proj 在同一个维度还是输出 3.0。
相乘：-0.24 × 3.0 = -0.72。信息被压得很小，方向也反了。

同一个 3.0 的材料，因为门控值不同，通过的量和符号都变了。 这就是“门”做的事：不是简单地“开或关”，而是按比例缩放。

down_proj：把结果收回去。 2432 维里筛完、缩放过后的信息，还散在桌子上。down_proj 把它们收拢，压回 768 维。

为什么要三道工序，而不是两道？

原版 ReLU 太粗暴。一个维度要么全通，要么全断，没有中间状态。SwiGLU 加了 gate_proj，把“筛选”这件事从硬切变成了软调节。表达能力强了很多。代价是多了一个矩阵，参数增加约 50%。所以中间维度从 4 倍 hidden_size 降到 8/3 倍左右，把账拉平。

5.6 中间维度为什么是 π 倍

MiniMind 用的不是 8/3。它用的是 π。

self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)

算一下：768 × π / 64 = 37.699，向上取整到 38，再乘回 64，得到 2432。

所以 MiniMind 实际用的是 2432。而且 2432 就是默认公式的结果，不需要任何显式配置。

为什么取整到 64 的倍数？ 为了让矩阵维度对齐硬件的计算友好度。64 是 GPU 上矩阵乘法的常见对齐粒度。

5.7 MiniMind 的 FFN

class FeedForward(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

三个矩阵：

gate_proj：768 → 2432，负责门控。
up_proj：768 → 2432，负责提供信息。
down_proj：2432 → 768，负责压回原维度。

bias=False：不加偏置项，省参数。

5.8 FFN 占了多少参数

把 FFN 和注意力对比一下。

每层注意力：

矩阵	形状	参数量
q_proj	768 × 768	589,824
k_proj	768 × 384	294,912
v_proj	768 × 384	294,912
o_proj	768 × 768	589,824
合计		1,769,472

每层 FFN：

矩阵	形状	参数量
gate_proj	768 × 2432	1,867,776
up_proj	768 × 2432	1,867,776
down_proj	2432 × 768	1,867,776
合计		5,603,328

FFN 的参数量是注意力模块的 3.17 倍。

为什么差这么多？看中间维度就明白了。注意力模块里，每个投影的输出维度不超过 768。FFN 里，gate_proj 和 up_proj 的输出是 2432 维，是 768 的 3 倍多。中间维度一大，矩阵就大，参数就多。

放到整个模型看：8 层 FFN 合计约 4480 万，8 层注意力合计约 1420 万，加上 embedding 层的 490 万。FFN 占了总参数的约 70%，注意力只占约 22%。

这个比例不是 MiniMind 的特例。所有主流 Transformer 都这样。有研究认为，FFN 的高维中间层像一个 key-value 记忆库，大部分“知识”存储在这里。所以参数量的大头落在 FFN，是有道理的。

注意力：找关系，决定 token 之间怎么互动。
FFN：记知识，决定每个 token 自己该被加工成什么样。

六、训得动：残差 + RMSNorm

6.1 深层网络的困境

8 层积木，每层都做一次注意力、一次 FFN。如果每层都把向量彻底改写，会发生什么？

第一层的输入是 embedding 向量，经过注意力、FFN，变成一个新向量。这个新向量进入第二层，又被彻底改写。到第八层出来，原始的 embedding 信息可能已经被磨得差不多了。

更严重的是，训练时梯度要从最后一层反向传到第一层。每经过一层，梯度都可能被放大或缩小。8 层下来，要么梯度爆炸，要么梯度消失。

6.2 残差连接

残差连接的做法极其简单：把输入直接加到输出上。

\text{输出} = \text{层}(x) + x

不是“彻底改写”，而是“在原向量上加一个修正量”。

这样做的效果是：原始信息永远有一条直通路。 不管中间的层把 x 变换成什么样，x 本身始终保留在输出里。每一层只需要学“在这个基础上，还应该修正什么”，而不是“从头重建整个表示”。

梯度也受益。反向传播时，残差连接提供了一个恒等路径，梯度可以沿着这条路直接传回去，不受中间层的影响。

6.3 RMSNorm：把数值拉回可控范围

神经网络逐层计算时，矩阵乘法会不断放大数值尺度。一层一层叠加，数值可能变得极大或极小。太大就溢出，太小就变成零，梯度也跟着消失。

归一化的作用，是在每个计算模块的入口，把数值拉回一个标准范围。

传统的 LayerNorm 做两件事：减去均值，除以标准差。

RMSNorm 只做一件事：除以均方根。

\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{d}\sum_{i=1}^{d} x_i^2 + \epsilon}} \cdot \gamma

不减去均值。就这一个区别。

为什么可以省掉？因为实践发现，减去均值对最终效果影响很小，但计算量少了。大模型里，少一步计算，乘以几百层，就是可观的节省。

6.4 MiniMind 的实现

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return self.weight * self.norm(x.float()).type_as(x)

x.pow(2).mean(-1) 算的是每个 token 向量的均方值。
torch.rsqrt 是平方根的倒数。
self.weight 是一个可学习的缩放因子，初始为 1。

归一化只改变向量的尺度，不改变方向。语义信息主要编码在方向里，所以不会因为归一化而丢失。

6.5 残差和归一化的组合

residual = hidden_states
hidden_states = self.self_attn(self.input_layernorm(hidden_states))
hidden_states += residual   # 第一次残差：注意力之后

residual = hidden_states
hidden_states = self.feed_forward(self.post_attention_layernorm(hidden_states))
hidden_states += residual   # 第二次残差：FFN 之后

两次残差，保证了信息在注意力模块和 FFN 模块之间流动时，原始内容始终有一份保留。

七、组装：一个完整的 Transformer 块

现在把六块拼起来。

输入 hidden_states
  │
  ├─ 保存 residual = hidden_states
  │
  ├─ hidden_states = RMSNorm(hidden_states)          ← 归一化
  │
  ├─ hidden_states = Attention(hidden_states)         ← 注意力
  │     ├─ Q、K、V 投影
  │     ├─ RoPE 旋转 Q 和 K
  │     ├─ KV Cache 拼接
  │     ├─ 注意力计算
  │     └─ o_proj 输出
  │
  ├─ hidden_states = hidden_states + residual         ← 残差连接
  │
  ├─ 保存 residual = hidden_states
  │
  ├─ hidden_states = RMSNorm(hidden_states)          ← 归一化
  │
  ├─ hidden_states = FeedForward(hidden_states)       ← FFN
  │
  └─ hidden_states = hidden_states + residual         ← 残差连接
  │
输出 hidden_states

这个块重复 8 次，就是 MiniMind 的全部积木。

组件	配置
hidden_size	768
层数	8
Q 头数	8
KV 头数	4
head_dim	96
FFN 中间维度	2432
归一化	RMSNorm
激活函数	SiLU
位置编码	RoPE（支持 YaRN）

注意力让 token 互相看，位置告诉模型谁在前谁在后，KV Cache 让推理不用重算，GQA 省显存，Flash Attention 省时间，FFN 让每个 token 自己消化，残差和 RMSNorm 保证深层能训得动。

这六块合起来，才是一个完整的 Transformer 块。

八、动手实验

实验一：看 RoPE 的旋转效果

import torch

def precompute_freqs(dim, end, rope_base=1000000):
    freqs = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(end)
    freqs = torch.outer(t, freqs).float()
    return freqs

freqs = precompute_freqs(96, 100)
print("位置 0 的前 5 个频率:", freqs[0, :5])
print("位置 1 的前 5 个频率:", freqs[1, :5])
print("位置 50 的前 5 个频率:", freqs[50, :5])

实验二：对比有无残差连接

import torch

x = torch.randn(1, 10, 768)

# 无残差：每层彻底改写
h = x
for _ in range(8):
    h = torch.randn(1, 10, 768) * 0.1
print("无残差，最终与原始输入的相似度:",
      torch.cosine_similarity(x.flatten(), h.flatten(), dim=0).item())

# 有残差：每层加修正量
h = x
for _ in range(8):
    h = h + torch.randn(1, 10, 768) * 0.1
print("有残差，最终与原始输入的相似度:",
      torch.cosine_similarity(x.flatten(), h.flatten(), dim=0).item())

实验三：KV Cache 的效果

seq_len = 1000
no_cache_size = sum(range(seq_len))

print(f"无缓存的总计算量: {no_cache_size}")
print(f"有缓存的总计算量: {seq_len}")
print(f"节省比例: {1 - seq_len / no_cache_size:.1%}")

九、本篇概念清单

概念	本篇交代到什么程度
RoPE（旋转位置编码）	讲透：原理、公式、MiniMind 实现、多频率设计、可学习性
相对位置编码	讲清：为什么旋转后点积只和距离有关
YaRN	讲清：外推的思路，MiniMind 的支持
KV Cache	讲透：为什么需要、怎么做、代价是什么
MHA / MQA / GQA	讲透：三种注意力的关系和取舍
Flash Attention	讲清：内存墙、分块计算、重计算
FFN / SwiGLU	讲透：注意力负责什么、FFN 负责什么、SwiGLU 的三层结构
残差连接	讲透：为什么需要、怎么实现
RMSNorm	讲透：和 LayerNorm 的区别、公式、实现
Transformer 块	讲透：完整组装顺序

本篇要牢记的只有三个词：RoPE、KV Cache、SwiGLU。

十、回到开头的问题

第三篇结尾说，注意力让 token 互相看。这一篇补上了剩下的部分。

模型怎么知道词序？RoPE 把位置信息变成旋转角度，融进 Q 和 K 的点积里。

怎么记住前文？KV Cache 把算过的 K、V 存起来，推理时不用重算。

显存不够怎么办？GQA 让多个 Q 头共享一组 K 和 V，KV Cache 减半。

算得太慢怎么办？Flash Attention 分块计算，不把 N×N 矩阵物化到 HBM 里。

每个 token 自己怎么消化？FFN 用 SwiGLU 做通道维度的非线性变换。

深层怎么训得动？残差连接保证信息直通，RMSNorm 把数值拉回可控范围。

六块拼起来，就是一个完整的 Transformer 块。8 个块叠起来，就是 MiniMind 的全部积木。

十一、思考题

如果把 RoPE 的 rope_base 从 1,000,000 改回 LLaMA-1 的 10,000，会发生什么？提示：想想频率衰减速度。
KV Cache 在训练时用不用？为什么？
残差连接是 输出 = 层(x) + x。如果改成 输出 = 层(x) + 0.1 * x，会有什么问题？
SwiGLU 的中间维度在 MiniMind 里用的是 π 倍率，不是 8/3 倍。为什么要取整到 64 的倍数？

备注：本篇的源码分析基于 MiniMind 仓库的 model/model_minimind.py。RoPE 的公式推导参考了苏剑林等人的原始论文。YaRN 的细节可以参考其论文《YaRN: Efficient Context Window Extension of Large Language Models》。Flash Attention 的原理参考 Tri Dao 等人的论文《FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness》。