biscuitの博客

1. Attention 前向传播是一条可追踪的张量流水线#

Attention.forward 接收某个 Decoder Block 中经过 RMSNorm 的隐藏状态：

$T_q$ 表示本次调用送入模型的 Query 长度。训练时通常一次送入整段序列，此时 $T_q=T$；使用 KV Cache 逐 Token 解码时，通常有 $T_q=1$。

本文使用以下符号：

MiniMind 默认使用 $d=768$、$h_q=8$、$h_{kv}=4$、$d_h=96$，所以：

下面的图先给出整条数据流。虚线缓存支路只在自回归生成时发挥作用；QK Norm 和 RoPE 都不处理 Value。

把这条流水线压缩成一组公式，就是：

后面的代码只是在显式构造这些量，同时不断改变它们的视图和形状。

2. Q/K/V 投影把表示空间拆成多个注意力头#

Attention 初始化时创建四个无偏置线性层：

1
self.q_proj = nn.Linear(
2
    config.hidden_size,
3
    config.num_attention_heads * self.head_dim,
4
    bias=False
5
)
6
self.k_proj = nn.Linear(
7
    config.hidden_size,
8
    self.num_key_value_heads * self.head_dim,
9
    bias=False
10
)
11
self.v_proj = nn.Linear(
12
    config.hidden_size,
13
    self.num_key_value_heads * self.head_dim,
14
    bias=False
15
)
16
self.o_proj = nn.Linear(
17
    config.num_attention_heads * self.head_dim,
18
    config.hidden_size,
19
    bias=False
20
)

nn.Linear(in_features, out_features) 保存的权重形状是 [out_features, in_features]，前向计算采用 $XW^\top$。因此：

forward 先完成投影proj，例如q_proj 输出的 768 维（内存里是连续的一条），再用 view 把最后一个总维度拆成“头数 × 头维度”view 不会再学习一次投影，也不会把信息分发给独立的小网络；它只是重新解释连续内存中的维度，768 个连续的数,前 96 个算头 0,接下来 96 个算头 1,以此类推。。真正决定每个头看到什么子空间的是 $W_Q$、$W_K$ 和 $W_V$ 中不同的行，$W_Q$ 的第 095 行产生头 0 的 query,第 96191 行产生头 1 的。：

1
bsz, seq_len, _ = x.shape
2

3
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
4

5
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
6
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
7
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)

对应的形状变化是：

MiniMind 默认 $h_qd_h=d$，所以 Query 投影（q_proj）前后的总维度相同；但这不等于没有发生变换。它仍然用一个 $d\times d$ 参数矩阵把每个 Token 的隐藏表示映射到了新的 Query 空间。K/V 总维度只有：

这正是 GQA 比标准多头注意力节省 K/V 参数和缓存的来源。默认配置下，四个投影的主要参数量为：

如果 K/V 也各有 8 个头，标准 MHA 对应的投影参数量则是 $4d^2$。

1
  标准 MHA:                     GQA (MiniMind):
2

3
  Q:  Q0 Q1 Q2 Q3 Q4 Q5 Q6 Q7   Q:  Q0 Q1 | Q2 Q3 | Q4 Q5 | Q6 Q7
4
       │  │  │  │  │  │  │  │        └─┬─┘   └─┬─┘   └─┬─┘   └─┬─┘
5
  K:  K0 K1 K2 K3 K4 K5 K6 K7   K:    K0      K1      K2      K3
6
  V:  V0 V1 V2 V3 V4 V5 V6 V7   V:    V0      V1      V2      V3
7

8
  8 组 K/V                         4 组 K/V，每组服务 2 个 Q 头

代码里的 n_rep 就是这个”每组服务几个”:

1
self.n_rep = self.n_local_heads // self.n_local_kv_heads   # 8 // 4 = 2

GQA 是 MHA 和 MQA 之间的插值:

3. QK Norm 与 RoPE 分别控制数值尺度和位置信息#

完成多头拆分(view)后，MiniMind 先归一化 Q/K，再施加 RoPE：

1
xq, xk = self.q_norm(xq), self.k_norm(xk)
2

3
cos, sin = position_embeddings
4
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)

这里的 q_norm 和 k_norm 都是 RMSNorm(self.head_dim)。以单个 Query 头向量 $q\in\mathbb{R}^{d_h}$ 为例：

代码在最后一维上计算均方根，因此每个 Token 的每个头独立计算自己的归一化尺度。与此同时，$\gamma_Q\in\mathbb{R}^{d_h}$ 会广播到 $B$、$T_q$ 和 $h_q$ 三个维度，即所有 Query 头共享同一组可训练缩放参数；Key 使用另一组独立的 $\gamma_K$。Value 不参与 QK Norm，因为注意力分数由 Q 与 K 的点积决定，而 V 是被权重汇总的内容。

自定义 RMSNorm 还会先把输入转成 float32 计算，再转回原来的数据类型：

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

6
def forward(self, x):
7
    return (self.weight * self._norm(x.float())).type_as(x)

这一步不改变形状：

随后 apply_rotary_pos_emb 将位置旋转施加到 Q/K：

1
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
2
    def rotate_half(x):
3
        return torch.cat(
4
            (-x[..., x.shape[-1] // 2:],
5
              x[..., :x.shape[-1] // 2]),
6
            dim=-1
7
        )
8

9
    q_embed = (
10
        q * cos.unsqueeze(1)
11
        + rotate_half(q) * sin.unsqueeze(1)
12
    ).to(q.dtype)
13
    k_embed = (
14
        k * cos.unsqueeze(1)
15
        + rotate_half(k) * sin.unsqueeze(1)
16
    ).to(k.dtype)
17
    return q_embed, k_embed

若把一个头向量沿最后一维分成等长两半：

那么位置 $t$ 的旋转可以写为逐元素运算：

cos 和 sin 的形状为 $[T_q,d_h]$，unsqueeze(1) 后变为 $[T_q,1,d_h]$，再通过广播同时作用于 Batch 与所有头。RoPE 不在隐藏状态上直接加一个位置向量，而是在 Q/K 空间中按位置旋转坐标，使后续点积能够表达相对位置信息。

两步的职责可以明确区分：QK Norm 控制参与点积的向量尺度，RoPE 改变参与点积的方向以编码位置。它们都只处理 Q/K，也都保持张量形状不变。

4. KV Cache 与 GQA 在不同维度上减少重复计算#

RoPE 之后，代码先拼接历史缓存，再展开 GQA 的 K/V 头：

1
if past_key_value is not None:
2
    xk = torch.cat([past_key_value[0], xk], dim=1)
3
    xv = torch.cat([past_key_value[1], xv], dim=1)
4

5
past_kv = (xk, xv) if use_cache else None
6

7
xq = xq.transpose(1, 2)
8
xk = repeat_kv(xk, self.n_rep).transpose(1, 2)
9
xv = repeat_kv(xv, self.n_rep).transpose(1, 2)

假设缓存中已有 $T_p$ 个历史 Token，本次又输入 $T_q$ 个 Token，那么拼接后的长度为：

缓存的形状是：

这两个缓存解决的是“时间维度上的重复”：历史 Token 的 K/V 已经算过，下一步生成时直接复用即可，不必把整个前缀重新送过模型。每层缓存的元素数为：

MiniMind 的 $h_{kv}=h_q/2$，因此其 KV Cache 元素数只有同头数 MHA 的一半。

但是 Query 有 $h_q$ 个头，缓存中的 K/V 只有 $h_{kv}$ 个头。点积计算前，repeat_kv 在头维度上把每个 K/V 头逻辑展开 $r$ 次：

1
def repeat_kv(x, n_rep):
2
    bs, slen, num_key_value_heads, head_dim = x.shape
3
    if n_rep == 1:
4
        return x
5
    return (
6
        x[:, :, :, None, :]
7
        .expand(bs, slen, num_key_value_heads, n_rep, head_dim)
8
        .reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
9
    )

默认的对应关系是：

1
Query heads:  Q0  Q1 | Q2  Q3 | Q4  Q5 | Q6  Q7
2
K/V heads:    K0  K0 | K1  K1 | K2  K2 | K3  K3
3
group:        └─ G0 ─┘ └─ G1 ─┘ └─ G2 ─┘ └─ G3 ─┘

用形状表示为：

Query 不需要复制，只需转置：

GQA 解决的是“头维度上的冗余”：一组 Query 头共享同一组 K/V 表示。repeat_kv 是为了让普通批量矩阵乘法能够按相同头数执行，并不表示缓存中真的保存了 $h_q$ 份 K/V；past_kv 在展开之前就已经建立，仍然只保存 $h_{kv}$ 个头。

还要注意，代码默认 $h_q$ 能被 $h_{kv}$ 整除，因为 n_rep 使用整数除法。若自定义配置破坏了：

展开后的头数就无法正确匹配 Query。

5. 缩放点积、掩码与 Softmax 共同得到注意力权重#

经过转置与 GQA 展开后，三个张量的形状为：

手动分支直接实现了缩放点积注意力：

1
scores = (
2
    xq @ xk.transpose(-2, -1)
3
) / math.sqrt(self.head_dim)
4

5
if self.is_causal:
6
    scores[:, :, :, -seq_len:] += torch.full(
7
        (seq_len, seq_len),
8
        float("-inf"),
9
        device=scores.device
10
    ).triu(1)
11

12
if attention_mask is not None:
13
    scores += (
14
        1.0 - attention_mask.unsqueeze(1).unsqueeze(2)
15
    ) * -1e9
16

17
weights = F.softmax(scores.float(), dim=-1).type_as(xq)
18
output = self.attn_dropout(weights) @ xv

首先计算每个 Query 与全部 Key 的相似度：

若 Q/K 各维度具有近似相同的尺度，未经缩放的点积方差会随 $d_h$ 增长。除以 $\sqrt{d_h}$ 可以避免 Softmax 输入随头维度变大而过度极端。

因果掩码把“当前 Query 不应该看到的未来 Key”加为 $-\infty$。没有缓存时，$T_q=T_{kv}=T$，它是熟悉的上三角矩阵：

有缓存时，scores[:, :, :, -seq_len:] 只在本次新增的 $T_q\times T_q$ 区域加入上三角掩码。所有历史列都保持可见，因为历史 Token 对当前 Query 来说都位于过去。逐 Token 解码时 $T_q=1$，这个 $1\times1$ 上三角矩阵为 0，当前 Token 可以关注全部 $T_p+1$ 个 Key。

外部 attention_mask 通常形如 $[B,T_{kv}]$，两次 unsqueeze 后广播为 $[B,1,1,T_{kv}]$。其中 0 对应的 Key 位置被加上一个极小值，从而在 Softmax 后取得近似 0 的权重。这里的核心不是“用 0 乘掉分数”，而是在 Softmax 之前对不合法位置施加加性掩码：

最后沿 Key 维归一化，并对 Value 加权求和：

scores.float() 让 Softmax 在 float32 中计算，降低半精度指数运算溢出或下溢的风险；随后 type_as(xq) 再把权重转回 Q 的数据类型。attn_dropout 只在训练模式下随机丢弃部分注意力权重，不改变形状。

预训练脚本调用模型时没有显式传入 attention_mask。PretrainDataset 采用右侧 Padding，真实 Token 无法越过因果掩码看到未来的 PAD；PAD 位置虽然可以看到前文，但其标签被设为 -100，不会进入语言模型 Loss。因此这一路径仍然成立。若 Batch 使用左侧 Padding 或更复杂的有效区间，就应正确传入 attention_mask。

6. Flash Attention 是同一数学运算的融合实现#

当运行环境支持且配置启用 Flash Attention 时，MiniMind 会优先调用 PyTorch 的融合算子：

1
if (
2
    self.flash
3
    and seq_len > 1
4
    and (not self.is_causal or past_key_value is None)
5
    and (attention_mask is None or torch.all(attention_mask == 1))
6
):
7
    output = F.scaled_dot_product_attention(
8
        xq,
9
        xk,
10
        xv,
11
        dropout_p=self.dropout if self.training else 0.0,
12
        is_causal=self.is_causal
13
    )
14
else:
15
    # 手动计算 scores、mask、softmax 和 output
16
    ...

这不是另一种注意力目标。两条分支都在计算：

差别在工程实现：融合算子可以分块完成点积、Mask、Softmax 与 Value 汇总，避免把完整的 $T_q\times T_{kv}$ 分数矩阵长期写入显存。序列越长，这种中间张量的显存开销越明显。

MiniMind 的条件也揭示了两条分支各自常见的使用场景：

所以阅读手动分支最容易看懂数学过程，但不应据此认为训练一定真的显式保存了 scores；实际是否融合取决于配置、运行环境与输入条件。

7. 合并多头后才回到 Decoder Block 的残差主干#

每个 Query 头都已得到自己的上下文向量后，代码把头维移回末尾并合并：

1
output = output.transpose(1, 2).reshape(
2
    bsz, seq_len, -1
3
)
4
output = self.resid_dropout(self.o_proj(output))
5
return output, past_kv

对应的形状变化为：

多头结果不是求平均，而是先拼接成 $h_qd_h$ 维向量，再由 $W_O$ 混合各头的信息并映射回模型隐藏维度。由于输出仍为 $[B,T_q,d]$，它才能与进入 Attention 前的残差分支相加。

不过，Attention.forward 本身只返回 Attention 输出，没有在内部执行残差加法。加法位于外层 MiniMindBlock.forward：

1
residual = hidden_states
2
hidden_states, present_key_value = self.self_attn(
3
    self.input_layernorm(hidden_states),
4
    position_embeddings,
5
    past_key_value,
6
    use_cache,
7
    attention_mask
8
)
9
hidden_states = residual + hidden_states

因此完整关系是：

resid_dropout 的名字表示它作用于即将进入残差加法的 Attention 输出，而不是说残差连接发生在 Attention 类内部。把模块边界分清，阅读代码调用链时就不容易把两个不同位置的 RMSNorm、Dropout 与 Add 混在一起。

8. 两条形状轨迹连接训练与增量解码#

为了让每个维度都能实际核对，下面使用一个缩小后的配置：

整段输入 $T_q=5$ 时，形状轨迹如下：

注意返回缓存仍然只有 2 个 K/V 头，而不是 GQA 展开后的 4 个头。

现在假设这 5 个 Token 已经写入缓存，下一次只输入 1 个新 Token。此时 $T_p=5$、$T_q=1$、$T_{kv}=6$：

这条轨迹说明了 KV Cache 的核心收益：第 6 个 Token 的 Query 仍需与 6 个 Key 比较，但前 5 个 Token 的 K/V 不再重复投影。生成长度增长时，每一步只为新增 Token 计算新的 Q/K/V，再把 K/V 追加到缓存中。

回到源码，可以把整个 Attention.forward 读成一句话：先将 $X$ 投影为多头 Q/K/V，用 QK Norm 稳定点积、用 RoPE 注入位置，再拼接紧凑的历史 K/V 并按 GQA 展开，经过带掩码的缩放点积注意力汇总上下文，最后合并各头并投影回 $d$ 维。

下一篇

Attention 已经从黑盒变成了一条完整的张量流。下一篇继续阅读同一个 Decoder Block 中的 FeedForward、RMSNorm 与残差连接，并进一步区分普通 SwiGLU FFN 和 MiniMind 可选的 MoE 路径。