Level 04 | Attention MHA GQA

Mission 1 TODO 1 / TODO 3 / TODO 4

注意力到底在算什么：打分表 → 概率 → 混合 value

🎯 先猜一猜（下注解锁）

注意力公式里，Q·Kᵀ 算完不直接 softmax，非要先除以一个 √head_dim。这个除法看着多余，去掉会怎样？

先押一个假设：如果去掉 √head_dim 这个缩放，最可能出什么问题？

Q 和 K 的 shape 对不上，根本没法相乘注意力分数会变成小数，不再是整数维度越大点积方差越大，softmax 会几乎全压到一个 token 上，梯度消失、学不动

先押一个假设，下面的讲解就会解锁。猜错完全没关系——带着疑问读，记得最牢。

先补的知识

先记一个生活类比：注意力就像“查资料”——你心里的问题是 Query，每份资料的标题是 Key，资料的正文是 Value。标题越对得上你的问题，就越多地去读它的正文。
Q/K/V 都是同一个输入 x 各乘一个不同的 Linear 矩阵得到的。同一句话，从“我想问什么(Q)/我能回答什么(K)/我有什么内容(V)”三个角度各投影一次。
“多头”= 把每个 token 的长向量切成 H 段，每段单独做一遍上面的查资料，让不同的“头”各盯一种规律（有的盯语法、有的盯指代）。计算时摆成 [B, H, S, D]：批次、头、序列长度、每头维度。
softmax：把一排任意大小的分数压成“加起来等于 1 的概率”，分数越高拿到的概率越大。

图解原理

一句话版：每个 token 拿自己的 Query 去和所有 token 的 Key 比一比“我跟你多相关”，得到一张打分表；softmax 把每一行变成一组概率；再按这组概率去混合所有 token 的 Value，混出来的就是这个 token 的新表示。下面把这 4 步拆开、每步配图，并在关键处停下来让你先想一想再看解释。

0. 先建立画面：一句话里，每个词都在“偷看”别的词

读“它躺在桌上”时，你要搞清楚“它”指什么，就得回头看前面的词。注意力做的就是这件事：让每个 token 都能按需去看序列里的其他 token，再把看到的信息揉进自己的新表示里。“按需”——就是下面那张打分表决定的。

想一想：为什么 Q、K、V 要用三个不同的矩阵，而不是直接拿 x 本身去两两点积？

如果直接用 x 和 x 点积，那“我作为提问方”和“我作为被查方”用的是同一个向量，模型没法分别学习“我想找什么”和“我能提供什么匹配特征”这两件不同的事。

拆成 Q、K 两个独立投影后，同一个词可以“以一种姿态发问、以另一种姿态被检索”，匹配关系才学得灵活。V 再单独一个矩阵，是因为“被选中后该交出什么内容”又是第三件事。三个角色、三个矩阵，各司其职。

第一步：把每个 token 的长向量切成 H 个头（TODO 1）

1. 切多头：一根宽向量拆成 H 个小专家

投影后 Q 是 [B, S, H·D]，先 view 成 [B, S, H, D]（把最后一根长向量切成 H 段），再 transpose(1,2) 成 [B, H, S, D]。为什么要转置？因为下一步要让每个头（H）各自独立地在它自己的 [S, D] 小表上做矩阵乘法，把 H 提到前面、和 batch 并排当“批量维”最顺手。

[B, S, H·D]view → [B, S, H, D]transpose(1,2) → [B, H, S, D]

直觉：H 个头像 H 个分工不同的读者，有的盯语法、有的盯指代，各读各的，最后再汇总。

每个头里：query 逐一去问每个 key（TODO 3 第一步）

2. Q @ Kᵀ：打一张 token×token 的相关性分数表

Q 是 [.., S, D]，K 转置成 [.., D, S]，相乘得到 [.., S, S]。第 i 行第 j 列 = 第 i 个 query 和第 j 个 key 的点积 = “token i 该多关注 token j”。

除以 √D，再 softmax（TODO 3 第二、三步）

3. 为什么要除以 √D？防止 softmax “一家独大”

点积 = 把 D 个数乘起来再相加。D 越大，这个和的“波动范围”（方差）就越大，分数越容易出现某个特别大的极端值。直接喂给 softmax，它会几乎把全部概率压到那一个 token 上（接近非 0 即 1 的 one-hot），别的位置概率趋近 0、梯度也趋近 0，模型就学不动了。除以 √D 正好把方差拉回到大约 1 的稳定区间，softmax 才平滑、可训练。

scores = Q·Kᵀ / √Dprobs = softmax(scores)每行加起来 = 1

想一想：为什么偏偏是除以“√D”，而不是除以 D 或别的数？

假设 Q、K 的每个分量都是均值 0、方差 1 的独立随机数。点积是 D 个“乘积项”相加，每项方差约为 1，独立相加后总方差 ≈ D，于是标准差 ≈ √D。

我们想把分数的“典型尺度”拉回到 1 左右，所以除以标准差 √D 最自然——除以 D 会矫枉过正、把分数压得太小，反而让 softmax 太平、谁都差不多。这就是原论文选 √D 的来历。

Causal Mask：生成任务里给打分表的“上三角”填 −∞（代表“还没出现的未来 token”），softmax 后这些位置概率自然变 0，保证每个 token 只能看自己和左边。

用概率加权 value，再把多头拼回去（TODO 3 末 + TODO 4）

4. probs @ V → 合并多头 → 输出投影

probs [.., S, S] 乘 V [.., S, D] 得 [.., S, D]：每个 token 的新向量 = 按注意力概率混合所有 token 的 value。最后 transpose(1,2) 把 H 放回去、contiguous().view 合并成 [B, S, H·D]，再过 o_proj 回到 hidden_dim。

probs @ V → [B,H,S,D]transpose+view → [B,S,H·D]o_proj → [B,S,hidden]

易错点：view 前必须 .contiguous()，因为 transpose 只是改了“怎么读这块内存”的步长、并没真正搬动数据，直接 view 会因内存不连续报错。

行业视角：这张 S×S 打分表，正是大模型的成本中心

注意力之所以打败了 RNN，关键在于它一步就能让任意两个 token 直接相连（RNN 要一格一格传，远距离信息会衰减），而且整张打分表能在 GPU 上并行算完。这是 Transformer 能堆大、能记长上下文的根本。

但代价也在这张表：它的大小是 S×S，序列翻倍、显存和计算就翻 4 倍。这就是为什么“长上下文”这么贵，也催生了 FlashAttention（不把整张表落地到显存，分块边算边累加）这类工程优化——本仓库后面的 FlashAttention 关就是专门拆这个的。今天你手写的这版是“教科书版”，理解它才能看懂工业版在优化什么。

语法热身：批量矩阵乘法只动最后两维

import torch.nn.functional as F
# [B, H, S, D] 想象成 B*H 个独立的 [S, D] 小矩阵
q = torch.randn(2, 4, 5, 8)   # 5 个 query，每个 8 维
k = torch.randn(2, 4, 7, 8)   # 7 个 key

# K 最后两维转置 [.., D, S]，再和 Q 相乘 → 打分表 [.., S, S']
scores = q @ k.transpose(-2, -1)      # [2, 4, 5, 7]
probs  = F.softmax(scores / 8**0.5, dim=-1)  # 沿 key 维归一化
# 前面的 batch、head 维（2,4）原样保留，只有最后两维参与乘法

读法：@ 对 4D 张量只在最后两维做矩阵乘法，前面的维度当“批量”广播。softmax(dim=-1) 表示“对每一行的 key 打分做归一化”。

巩固一下

巩固（动手算）：q=[2,4,5,8]、k=[2,4,7,8]，scores = q @ k.transpose(-2,-1) 的 shape 是？

[2, 5, 4, 7] [2, 4, 5, 7] [2, 4, 8, 8]

回到 notebook 的作业

这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。

TODO 1：投影结果 view 成 [B,S,H,D] 再 transpose(1,2) → [B,H,S,D]。 TODO 3：算 Q@Kᵀ/√D，加 mask，softmax，再 @ V。 TODO 4：transpose 回来 + contiguous().view 合并多头 → o_proj。

Mission 2 TODO 2 / repeat_kv

KV Cache 与 GQA：推理为什么慢，又怎么省显存

🎯 先猜一猜（下注解锁）

GQA 省显存，靠的是只存少量“窄”KV 头。但矩阵乘法又要求 KV 头数和 Query 头数一致——于是代码里有个 repeat_kv 把 KV 复制宽。问题来了：这步该放在缓存之前还是之后？

先押一个假设：repeat_kv（把 KV 扩充到 Query 头数）应该放在哪里？

放在更新缓存之前——先扩充好再存，省得每步重复扩充放在更新缓存之后——缓存里只存窄 KV，每步前向时临时扩充放在 softmax 之后——等算完注意力再扩充

先押一个假设，下面的讲解就会解锁。猜错完全没关系——带着疑问读，记得最牢。

先补的知识

先理解“自回归生成”：模型像打字一样一次只吐 1 个字，然后把这个字接到句子末尾，再用更长的句子算下一个字，循环往复。
算新字时，注意力要让它去看前面所有字的 Key/Value。但前面那些字的 K/V 上一步就算过了、而且不会再变——重算就是纯浪费。
三种省法的名字：MHA（每个 Query 头配一个专属 KV 头，最费）、MQA（所有 Query 头挤用 1 个 KV 头，最省但效果掉）、GQA（分组共享，折中，LLaMA-2/3 采用）。

图解原理

这一关是两个独立但配合使用的优化：① KV Cache——把算过的 K/V 存起来，每步只算新字那一个，把“每步重算整段历史”的浪费去掉；② GQA——让缓存里只存少量共享的 KV 头来省显存，等真正算注意力前再临时把它们复制成和 Query 一样多。最关键的一个细节：复制（repeat_kv）必须放在“存进缓存之后”，先存窄的、后扩宽的——下面会讲为什么反过来就前功尽弃。

1. 痛点：不缓存就要每步重算整段历史（O(N²)）

生成第 100 个 token 时，注意力要用到前 99 个 token 的 K/V。可这 99 个 K/V 在前几步早就算过了，而且不会变。每步都从头算 → 总开销 ∝ N²，序列越长越爆炸。

解法：把历史 K/V 存进 Cache，每步只算新增的那一个（TODO 2）

2. KV Cache：沿 seq 维 cat 一下就行

缓存里存着历史 K/V，形状 [B, H_kv, 旧长度, D]。新 token 算出自己的 xk/xv（长度 1），用 torch.cat(..., dim=2) 拼到时间轴末尾，长度 +1。下一步再接着拼。

想一想：KV Cache 把时间省下来了，那它代价是什么？

天下没有免费的午餐：省了重算的时间，代价是显存。每生成一个 token，缓存就长大一截，要一直占着显存直到这句话生成完。

粗略感受一下规模：层数 L、KV 头数 H_kv、每头维度 D、序列长 S、再算上 K 和 V 两份，缓存大小 ≈ 2 · L · H_kv · D · S。序列越长占用线性增长——长上下文推理时，KV Cache 往往比模型权重还吃显存。这正是下面要用 GQA 来砍它的原因。

另一条线：怎么让 Cache 本身更小？→ GQA

3. MHA → MQA → GQA：在效果和显存间取平衡

MHA 每个 Query 头配一个独立 KV 头，Cache 最大、效果最好；MQA 所有 Query 头共用 1 个 KV 头，Cache 最小但效果掉；GQA 折中——分组共享。下图 4 个 Q 头、2 个 KV 头，每 2 个 Q 共享 1 个 KV，Cache 直接砍半。

想一想：MQA 把 KV 砍到只剩 1 个头、最省显存，为什么大家不直接全用它？

KV 头可以理解成“资料库的不同检索视角”。MHA 有 H 个独立视角，能从很多角度去匹配；MQA 把它压成 1 个，所有 Query 头只能共用同一套 Key/Value，表达力明显变窄，效果通常会掉。

GQA 是中庸之道：保留几组（比如 8 组）视角，让每组被若干 Query 头共享。显存接近 MQA，效果接近 MHA——这就是 LLaMA-2/3 选它的原因。工程里很多决策都是这种“在两个极端之间找平衡点”。

关键顺序：先存“瘦” KV，再临时扩充（repeat_kv）

4. 延迟扩充：缓存窄的，算注意力前才复制宽的

矩阵乘法要求 Q 头数和 K/V 头数一致。GQA 的解法不是把 KV 存成宽的（那就退化回 MHA、白省了），而是只缓存窄的 num_kv_heads 个头，每次前向用 repeat_kv 临时把每个 KV 头复制 num_queries_per_kv 份，对齐到 Query 头数。

cat 更新 cache（窄）new_kv_cache = 窄 KVrepeat_kv → 宽，喂给注意力

易错点：repeat_kv 必须在“存进 new_kv_cache 之后”才做。先扩充再缓存 = 缓存里存了宽 KV = 显存退回 MHA，GQA 等于白做。

行业视角：GQA 为什么是 70B 模型能跑起来的关键之一

LLaMA-2 70B 有 64 个 Query 头，但只用 8 个 KV 头（每 8 个 Q 共享 1 个 KV），KV Cache 直接降到 MHA 的 1/8——在长上下文推理时能省下几十 GB 显存，这往往决定了一张卡到底装不装得下、能开多长上下文、能同时服务多少用户。

而“先存窄、用时再扩”的延迟扩充，正是 vLLM、TensorRT-LLM 这些推理框架的标准做法。它们之所以敢这么干，是因为注意力是 Memory-bound（瓶颈在搬数据而非算数）——临时复制几下的计算量可以忽略，省下的显存带宽才是大头。本仓库后面的 vLLM PagedAttention 关会接着讲怎么把这些 KV 块管得更省。

语法热身：沿时间维拼接 + 复制头

# ① KV Cache：把新 K 拼到历史末尾（dim=2 是 seq 维）
cache = torch.randn(2, 2, 5, 8)   # [B, H_kv, 旧长度5, D]
new_k = torch.randn(2, 2, 1, 8)   # 新 token 的 K，长度1
updated = torch.cat([cache, new_k], dim=2)  # [2, 2, 6, 8]

# ② repeat_kv：每个 KV 头复制 n_rep 份，对齐 Query 头数
#   先插一个新轴 expand，再 reshape 把它并进 head 维
n_rep = 2                          # 4 个 Q 头 / 2 个 KV 头
x = updated[:, :, None, :, :].expand(2, 2, n_rep, 6, 8)
x = x.reshape(2, 2 * n_rep, 6, 8)  # [2, 4, 6, 8] 对齐 Q

读法：cat(dim=2) 在 seq 维接长；expand 不复制内存只“假装”有 n_rep 份，reshape 时才真正铺开，把 KV 头数从 2 撑到 4。

巩固一下

巩固（动手算）：num_heads=8、num_kv_heads=2 时，每个 KV 头要被复制几份去服务 Query 头？

8 2 4

回到 notebook 的作业

这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。

TODO 2：用 torch.cat([k_cache, xk], dim=2) 在 seq 维拼接历史与新 KV。确认 new_kv_cache 存的是拼接后、repeat_kv 之前的“窄”KV。理解 repeat_kv：num_queries_per_kv = num_heads // num_kv_heads，扩充后头数对齐 Query。