Lecture 2 · Transformer：从词向量到注意力，以及它的计算账本 · Transformers: From Word Vectors to Attention, and Their Arithmetic

对应 Maps to：Stanford CS349E / EE339 Efficient ML Inference at Scale（Autumn 2026），Lecture 2 · Thu Sep 24 · Transformers（Fredrik Kjolstad）。阅读材料：T. Mikolov, K. Chen, G. Corrado, J. Dean, Efficient Estimation of Word Representations in Vector Space（2013，word2vec）；A. Vaswani et al., Attention Is All You Need（NeurIPS 2017）。原课作业 1（LLM Inference，9/24–10/1）的一部分主题落在本讲。 依据说明：课程网站只公开了课表和阅读材料，没有 slides。本讲以两篇阅读材料为主线，另外用到这些公开文献：Mikolov et al. 2013b（Distributed Representations of Words and Phrases，负采样的出处）；Su et al. 2021（RoFormer，RoPE 的出处）；Shazeer 2019（Fast Transformer Decoding，MQA）；Ainslie et al. 2023（GQA）；DeepSeek-V3 技术报告 2024（MLA）；How to Scale Your Model 第 4 章（Transformer 的 FLOPs 和字节账）。这些补充是我选的，不代表课堂上实际讲了什么。 配套代码 Code：code/lecture02/（numpy；1 个共享模块 + 5 个脚本，每个都在 15 秒内跑完） 预计阅读时间：3–4 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 解释”词向量”是怎么从”共现统计”里学出来的（word2vec 的 CBOW / skip-gram、负采样），以及 word2vec 为什么是一个”先算账、再设计”的典型例子。
2. 从”软的字典查找”出发，推导缩放点积注意力（scaled dot-product attention），并说清楚每个设计的理由：为什么除以 $\sqrt{d}$、为什么要因果掩码、为什么要多头、为什么要位置编码（以及 RoPE 为什么只依赖相对位置）。
3. 画出一个 Llama 风格 decoder-only Transformer 的一层：pre-norm、RMSNorm、GQA 注意力、SwiGLU MLP、残差连接；能用 einsum 记号写出每一步。
4. 逐个矩阵地数出参数、FLOPs、权重字节、KV 字节，并和第 1 讲的公式对账（本讲的代码证明它们分毫不差）。
5. 解释 KV cache 为什么能成立、省了多少、代价是什么；写一个带 KV cache 的 decode 循环，并知道怎样验证它是对的。
6. 比较 MHA、GQA、MQA、MLA 四种注意力在推理时的账本：每 token 的 KV、读 KV 的时间、注意力核心的算术强度。

0.2 先修与本讲范围 Prerequisites & scope

本讲不讲训练（反向传播、优化器、学习率），只在 word2vec 那里手推了一个简单的梯度。我们关心的是”推理时这个模型要算什么、搬什么”。

0.3 本讲地图 Map

§1  从词到向量：word2vec 的代价账、CBOW 与 skip-gram、负采样、类比、嵌入 = 查表           ← 代码 01
§2  注意力：RNN 的瓶颈；软查找；Q/K/V；√d；因果掩码；多头；位置编码与 RoPE               ← 代码 02
§3  组装 decoder-only Transformer：一层的结构；RMSNorm；SwiGLU；输出；einsum 记号
§4  账本：参数、FLOPs、字节，和第 1 讲对账；注意力何时重要；prefill / decode 的计算图      ← 代码 03
§5  推理：自回归生成、KV cache 的正确性、省了多少、代价、批量 decode                     ← 代码 04
§6  为推理而改的注意力：MQA、GQA、均值池化转换、MLA                                     ← 代码 05
§7  现状（2025–2026）
§8–12  小结、常见误区、练习（含答案）、延伸阅读、下一讲预告

0.4 记号 Notation

沿用第 1 讲的记号，并加上 Transformer 内部的维度。单个大写字母的写法来自 How to Scale Your Model 第 4 章（§3.5 的 einsum 会用到）：

注意 How to Scale Your Model 里 $K$ 是 KV 头数、$G$ 是组大小、$H$ 是每头维度，和很多论文的习惯不一样。本讲正文用 $G$ 表示组数、$R$ 表示组大小（和代码 transformer_lab.py 一致），只有在 §3.5 写 einsum 字符串时用书里的字母。

1. 从词到向量：word2vec · From words to vectors

代码：01_word2vec.py

1.1 问题：计算机眼里的词

神经网络只会处理数字。最朴素的做法是给词表里每个词一个编号，再写成独热向量（one-hot）：词表有 $V$ 个词，第 $i$ 个词就是一个长度 $V$、只有第 $i$ 位是 1 的向量。

独热向量有两个毛病：

• 太长：$V$ 可能是几十万到上百万。
• 没有相似度：任何两个不同词的独热向量都互相正交。“猫”和”狗”的距离，与”猫”和”民主”的距离完全一样。模型在”猫”上学到的东西，一点也迁移不到”狗”上。

我们想要的是每个词一个短的、稠密的向量（几十到几千维），而且意思相近的词，向量也相近。这就是分布式表示（distributed representation）。

1.2 分布式假设：一个词的意思由它的上下文决定

“你会通过一个词的伙伴认识它”（语言学家 J. R. Firth 的名言）。“猫”和”狗”经常出现在相似的句子里（“喂___”、“___在睡觉”），“民主”则不会。所以，只要让经常出现在相似上下文里的词有相似的向量，就能从大量文本里自动学出语义——不需要任何人工标注。

在 word2vec 之前，神经网络语言模型（NNLM，Bengio et al. 2003）已经在做这件事：它用前 $N$ 个词预测下一个词，词向量是副产品。问题是太贵。

1.3 先算账：NNLM 为什么训练不动

word2vec 论文的第一件事是定义代价。训练代价 $O = E \times T \times Q$，$E$ 是遍数，$T$ 是训练词数，$Q$ 是每个训练样本要”碰”的参数个数。论文给出了几种模型的 $Q$（$N$ 个输入词、投影维 $D$、隐藏层 $H$、词表 $V$、上下文距离 $C$）：

代码 01 演示 2 把论文给的典型值代进去（$N = 10$、$D = H = 500$、$C = 10$、$V = 10^6$）：

读这张表要分两步：

1. 完整 softmax 的 $H \cdot V$ 一项是 5 亿：输出层要给 100 万个词每个都打一次分。先用层次 softmax（把词表组织成一棵 Huffman 树，只需走 $\log_2 V \approx 20$ 步）或负采样（§1.5）把它降下来。
2. 降完之后，NNLM 的瓶颈变成 $N \cdot D \cdot H$——那个非线性隐藏层（250 万）。word2vec 的关键决定是整个去掉隐藏层：模型变成”对数线性”（log-linear）的，表达能力弱了，但 $Q$ 小了两个数量级，同样的算力能多看两个数量级的数据。

这是全课第一次出现的主题：先算账，再设计。论文的赌注是”简单模型 + 更多数据”胜过”复杂模型 + 较少数据”，结果赌对了：论文表 6 里，在 60 亿词上训练的 skip-gram（1000 维）类比准确率 65.6%，NNLM（100 维）是 50.8%；训练时间分别是 2.5 天 × 125 个 CPU 核和 14 天 × 180 个 CPU 核。

Transformer 会走一条看起来相反、其实同样的路：每个 token 的运算量很大，但全部是能并行的大矩阵乘法，硬件吃得下（§2.1）。“算账”的对象从 FLOPs 变成了”FLOPs 能不能被硬件高效执行”。

1.4 CBOW 与 skip-gram

word2vec 有两种结构：

• CBOW（continuous bag-of-words）：用周围的词（前后各 4 个）预测中间的词。周围词的向量直接取平均（所以叫”词袋”，不管顺序）。
• skip-gram：反过来，用中间的词预测周围的每个词。离得远的词关系弱，所以随机少采一些（每次在 $1 \ldots C$ 之间随机取一个窗口大小 $R$）。

两者都只有两个矩阵：输入向量 $v_w$（最后当作词向量用）和输出向量 $u_w$。skip-gram 对一对 (中心词 $c$, 上下文词 $o$) 的概率是

分母要对全部 $V$ 个词求和，这就是 $D \cdot V$ 那一项的来源。论文报告 skip-gram 在语义类比上明显更好，CBOW 在句法类比上略好、训练也更快。

1.5 输出层太贵：负采样

后续论文（Mikolov et al. 2013b）提出了比层次 softmax 更简单的负采样（negative sampling）：不再计算完整的概率，而是把问题变成二分类——“这一对是真实的共现吗？“对每个真实对，再从词表里随机抽 $k$ 个”负样本”，目标是

直觉：真实上下文的点积要大，随机词的点积要小。每个训练对只需算 $k + 1$ 个点积，而不是 $V$ 个。负样本按”词频的 0.75 次方”抽，这是经验值，比按词频抽更常抽到少见词。代码 01 用 numpy 手写了它的梯度（推导在 train_sgns 的注释里，只有三行）。

1.6 实验：类比从共现里长出来

真实语料太大，代码 01 造了一个”规则已知”的合成语料：

• 性别对（king/queen、man/woman、father/mother……）：上下文 = 这个”角色”的话题词（palace、crown……）+ 这个”性别”的词（he/his 或 she/her）。
• 国家 / 首都（france/paris……）：国家和首都共享该国的专属词，国家另配”国家类”词，首都另配”城市类”词。
• 形容词 / 比较级（big/bigger……）：共享话题词，原级配 very/quite，比较级配 than/even。

skip-gram 只看到”谁和谁一起出现”。2 万句、15.5 万个词、30 维、负采样 $k = 5$ 训练 3 遍后（代码 01 演示 3）：

类比的算法就是论文 §4 的方法：计算 $x = v(a^*) - v(a) + v(b)$，在全部词里找余弦相似度最大的（排除三个问题词）。10 个”男 → 女”差向量两两之间的平均余弦相似度是 0.98——性别被学成了一个几乎恒定的方向，所以 $v(\text{queen}) \approx v(\text{king}) - v(\text{man}) + v(\text{woman})$。

两个值得注意的地方：

• 最近邻看不出这些关系。king 的最近邻是 throne、palace、crown、kingdom（都是它的话题词），而不是 queen 或 man。关系藏在方向里，不在距离里。这正是论文要专门设计类比测试集（8,869 道语义题 + 10,675 道句法题）的原因。
• 我们的语料是刻意造得很干净的，所以准确率是 100%；论文在真实语料上的最好结果约 65%。

1.7 维度和数据要一起加

代码 01 演示 4 扫描了语料大小和向量维度（三组关系的平均类比准确率）：

• 两维空间放不下几十个词、好几种互相独立的关系，数据再多也没用。
• 数据太少时加维度也没用。$d = 5$ 时数据加 5 倍只从 25.5% 涨到 38.7%，$d = 10$ 时却从 60.9% 涨到 100%。

这和论文表 2 的结论一样：维度和数据要一起加。这个”维度”就是后来 Transformer 的 $d$（d_model），它直接进入第 1 讲的每一条公式。

1.8 系统视角：嵌入 = 查表

独热向量乘以嵌入矩阵 $E[V, d]$，结果就是 $E$ 的第 $i$ 行。数学上是矩阵乘法（$2Vd$ 次运算、读整个矩阵），实现上只是按下标取一行（gather，0 次运算、只读 $d$ 个数）。代码 01 演示 1 验证两者完全相等。

这解释了第 1 讲公式里的一个细节：“参与矩阵乘法的参数”不含词嵌入表。但输出端正好相反：

lm_head 要给词表里每个词打分，必须整张读、整张乘——它占 8B 模型全部权重的 6.5%，decode 每一步都要读一遍。小模型常让两者共享（tied）一个矩阵，省的是参数和显存；decode 时 lm_head 还是要整张读，带宽一点没省。

2. 注意力：软的字典查找 · Attention as a soft lookup

代码：02_attention_from_scratch.py

2.1 动机：RNN 的瓶颈是”顺序”

word2vec 给了每个词一个固定的向量，但词义依赖上下文（“bank”是银行还是河岸）。2017 年之前，处理句子的主流是循环神经网络（RNN）：

第 $t$ 步必须等第 $t-1$ 步算完。Attention Is All You Need 表 1 用三个指标比较了几种层：

按运算量算，$n < d$ 时自注意力更省；论文当时的句子（几十个词）远比 $d = 512$ 短。但更重要的是后两列：

• 顺序步数决定了能不能并行。代码 02 演示 6 在 numpy 上量了一下（float32，$d = 256$，墙钟时间，编写时的机器）：同样的运算量，RNN 一步一步算比”把 $n$ 个向量摞成一个矩阵、一次乘完”慢约 10 倍（几次运行在 9–11 倍之间；$n = 256$ 时一次是 675 µs vs 74 µs）。每一步只是向量乘矩阵，每个权重只用一次，还要付每步的调用开销。在有上万个计算单元的 GPU 上，差距更大。
• 路径长度决定了学长距离依赖有多难：RNN 里第 1 个词的信息要经过 $n$ 次变换才能到第 $n$ 个词，注意力里只要一步。

所以 Transformer 胜出的系统原因是：训练时所有位置一起算，全是大矩阵乘法。记住这一点，因为推理的 decode 阶段又退回了”一步一个 token”（§5）。

2.2 软查找

Python 字典是”硬”查找：给一个 key，完全相等才取出 value。注意力把它变”软”：

query $q$ 和哪个 key $k_i$ 的点积大，就主要取哪个 value $v_i$。代码 02 演示 1：5 个 key、value 是 10–50，query 是”有点像 key 2”的向量：

$\tau$ 小时接近硬查找，$\tau$ 大时接近平均。和硬查找最大的不同是：软查找对 $q$、$k$、$v$ 处处可导，模型可以用梯度下降学会”查什么”、“用什么当索引”、“取出什么”。

2.3 Q、K、V 从哪来：自注意力

自注意力（self-attention）里，三者都来自同一个序列：对每个 token 的向量 $x$，

每个 token 都拿自己的 $q$ 去”查”所有 token 的 $k$，取回它们 $v$ 的加权和。于是每个 token 的新向量都混进了上下文的信息——“bank”旁边是”river”还是”money”，它的输出向量就会不同。把所有 token 摞成矩阵，整个过程是两次矩阵乘法：

论文 §3.2.1 解释了为什么选点积而不是”加性注意力”（用一个小网络算相似度）：两者理论复杂度差不多，但点积能直接用高度优化的矩阵乘法代码。又是一个硬件友好的选择。

2.4 为什么除以 $\sqrt{d}$

论文脚注 4：如果 $q$、$k$ 的每个分量独立、均值 0、方差 1，那么点积 $q \cdot k = \sum_{i=1}^{d} q_i k_i$ 的方差是 $d$。$d$ 大时，分数之间的差距很大，softmax 的输出接近 one-hot，梯度几乎为 0。代码 02 演示 2（1 个 query、64 个 key，200 次平均）：

不缩放时，$d = 1024$ 的最大权重已经是 0.96（几乎 one-hot），雅可比范数（控制梯度能传回多少）从 0.30 掉到 0.06；除以 $\sqrt{d}$ 之后两列都不随 $d$ 变。缩放只是乘一个常数，可以并进 $W_q$，推理时零代价。它还让低精度推理受益：分数的范围小，不容易在 FP16 / FP8 里溢出（前置讲 P3 §5.4）。

2.5 因果掩码：不能偷看答案

语言模型逐个生成 token，训练时也必须保证第 $t$ 个位置只能看到位置 $\le t$ 的 token。做法：在 softmax 之前，把 $j > t$ 的分数设成 $-\infty$，权重就变成 0。

代码 02 演示 3 在一个 12 个 token 的序列里只改第 6 个：

有掩码时，前面位置的输出完全不受后面 token 的影响。这条性质就是 KV cache 的基础：位置 $j$ 的 $k$、$v$ 只依赖 token $0 \ldots j$，以后不会变，可以存起来复用（§5.2）。双向模型（BERT、原论文的编码器）做不到这一点。

2.6 多头：同样的运算量，几个地方一起看

一个头的 softmax 只能产生一个分布：如果一个 token 需要同时关注”前一个词”和”句子的主语”，一个头就要在两者之间妥协。多头注意力把 $d$ 维切成 $h$ 份，每份 $d/h$ 维，各自独立做注意力，最后拼起来再乘 $W_o$：

代码 02 演示 4（$d = 512$，$T = 1024$）：

• 头数不改变运算量（论文 §3.2.2：“由于每个头的维度降低了，总运算量和全维度的单头注意力差不多”）。多头几乎是免费的。
• 头太多、每头太窄时，每个头的矩阵乘法太”瘦”，硬件效率下降。真实模型的 $d_h$ 通常是 64–128。
• 最后两行预告了 §6：GQA / MQA 让几个 query 头共用一组 K/V，省的是 K/V 的投影和每 token 的 KV，注意力核心的运算不变。它们省的是字节，不是 FLOPs。

2.7 位置编码：注意力本身不知道顺序

注意力只是”加权求和”，对顺序视而不见。代码 02 演示 5(a)：没有位置编码、没有掩码时，打乱 10 个 token 的顺序，输出只是跟着打乱（差 $5 \times 10^{-15}$）——“猫追狗”和”狗追猫”得到同一组向量。所以必须注入位置信息。（有趣的是，因果掩码本身也泄露了位置：只能看前面的 token。有些模型干脆不用位置编码，就靠这一点。）

正弦位置编码（原论文）：给第 pos 个位置加上一个向量，

每一对维度是一个转速不同的”钟表指针”：前几对转得快，能分辨相邻位置；后几对转得慢，能分辨远距离。论文选它的理由是：对任意固定偏移 $k$，$PE_{\text{pos}+k}$ 是 $PE_{\text{pos}}$ 的线性函数（每一对旋转 $k\omega_i$），模型容易学会”看相对位置”。代码 02 演示 5(b) 验证了这个线性关系（误差 $10^{-15}$）。论文还试过可学习的位置向量，结果几乎一样；选正弦是因为”可能外推到训练时没见过的长度”。

RoPE（旋转位置编码，Su et al. 2021）更进一步：不是把位置加到输入上，而是把 $q$、$k$ 按位置旋转。第 $i$ 对数在位置 $m$ 旋转角度 $m\theta_i$：

两个旋转后的向量做点积，旋转角度相减，结果只依赖 $m - n$。代码 02 演示 5(c)：同一对 $q$、$k$ 放在位置 (5, 2)、(105, 102)、(10005, 10002)，点积都是 3.6432231984——只看相对位置。Llama、Qwen、DeepSeek 等主流模型都用 RoPE。

演示 5(d) 用 RoPE 手工搭了一个”总是看前一个 token”的头：所有 token 内容相同，只靠旋转区分位置；在位置 3、7、11，它都把 0.98 以上的权重放在前一个位置上。真实模型里确实能找到这样的头（previous-token head），它是复制、归纳等能力的基础部件。

【系统联系】RoPE 是在 $k$ 写进 KV cache 之前加的，所以缓存里的 $k$ 已经”带着位置”。这一点在两个地方很重要：§5.3 的”位置编号写错”bug，以及 §6.4 里 MLA 为什么要单独存一份带 RoPE 的 key。

3. 组装一个 decoder-only Transformer · Assembling the model

代码：transformer_lab.py（一个 Llama 风格、能真正运行、会记账的 numpy 实现）

3.1 从编码器–解码器到 decoder-only

原论文是为机器翻译设计的编码器–解码器：编码器用双向注意力读源语言句子，解码器用因果注意力逐词生成目标语言，并通过”交叉注意力”看编码器的输出。base 模型 6 + 6 层、$d = 512$、8 头、$d_\text{ff} = 2048$，6500 万参数，在 8 块 P100 上训练 12 小时；big 模型 2.13 亿参数、训练 3.5 天，在 WMT 2014 英德翻译上达到 28.4 BLEU、英法 41.8 BLEU。

今天的 LLM（GPT、Llama、Qwen、DeepSeek……）几乎都只用解码器（decoder-only）：一个因果注意力的堆叠，输入是”到目前为止的所有 token”，输出是”下一个 token 的分布”。没有编码器，也没有交叉注意力。下面只讲这一种。

3.2 一层的结构

transformer_lab.forward 里的一层（Llama 3 的结构）：

h = RMSNorm(x)
q, k, v = h·Wq, h·Wk, h·Wv          # q: H 个头；k, v: G 组头（GQA）
q, k = RoPE(q), RoPE(k)
a = 因果注意力(q, k, v)
x = x + a·Wo                         # 残差连接
h = RMSNorm(x)
x = x + (SiLU(h·W_gate) ⊙ (h·W_up))·W_down    # SwiGLU MLP + 残差

几个部件的理由：

• 残差连接（$x = x + f(x)$）：每一层只学”在原来基础上改一点”。几十层叠起来时，梯度可以沿着"$+x$"那条路直接传回去。从推理的角度看，残差流（residual stream）就是那个贯穿所有层、长度为 $d$ 的向量，所有层都在往它上面读写。
• 归一化：原论文用 LayerNorm（减均值、除标准差），而且放在残差加法之后（post-norm）。今天几乎都用 pre-norm（先归一化、再进子层、再加回残差），训练更稳定。Llama 用 RMSNorm：只除以均方根、不减均值，少一次规约（reduction），效果差不多。
• MLP：原论文是 $\max(0, xW_1)W_2$（ReLU），$d_\text{ff} = 4d$。Llama 用 SwiGLU：两个上投影，一个过 SiLU 当”门”，逐元素相乘，再下投影——三个矩阵，所以 $d_\text{ff}$ 取约 $\frac{8}{3}d$（再取整），总参数和两个矩阵、$4d$ 的版本差不多。
• GQA（§6.2）：K、V 只有 $G$ 组头。

3.3 输出：从向量到下一个 token

最后一层之后再做一次 RMSNorm，乘以 lm_head $[d, V]$，得到每个词的分数（logits），softmax 之后就是下一个 token 的概率分布。生成时从这个分布里选一个 token：贪心（取最大）、温度采样（分数除以温度再采样）、top-$k$ / top-$p$（只在最可能的几个里采样）。本讲代码用贪心，因为它是确定性的，方便比较”带 cache”和”不带 cache”的结果。

3.4 原始 Transformer vs Llama 3

这张表里只有一项是纯粹为推理做的改动：GQA。§6 会专门讲它和它的”亲戚”。

3.5 einsum 记号

How to Scale Your Model 用 einsum 描述 Transformer 的每一步：每个维度一个字母，出现在两个输入里、但不在输出里的字母被求和。代码 03 演示 3 把 Llama-3.1-8B 的一层写成（书里的字母：B 批、T / S 长度、D 模型维、F MLP 维、N query 头、K KV 头、G 组大小、H 每头维度）：

FLOPs 规则：一次缩并的 FLOPs = 2 × 所有出现过的维度之积（同一个字母只算一次）。比如 btd,df->btf 是 $2BTDF$。这个规则不用画图，只看字母就够了。一层合计 962.1 G，和 analytic_ops 逐项数出来的完全相同。

这个记号有三个好处：

• $QK^\top$ 和 $P \cdot V$ 里 B、K 是”批维”，T、S 同时出现，所以运算量 $\propto T \cdot S$，平方项一目了然。
• GQA 就是把 N 拆成 K·G：K、V 的形状里只有 K 没有 G，这就是它省 KV cache 的原因。
• 第 4 讲的切分记号就是在这些字母上加下标（比如 $D$ 被切到 Y 轴上写成 $D_Y$）。

4. 账本：逐个矩阵数 · The bookkeeping, matrix by matrix

代码：03_transformer_accounting.py。transformer_lab.Counter 在每次矩阵乘法、每次读 KV 时记一笔账；analytic_ops 是同样约定下的纯公式版本。

4.1 为什么要”数”

第 1 讲的公式是写出来的，没有被检验过。公式最容易出错的地方是漏项（比如漏了 lm_head）、重复计数（比如 tied 时把 lm_head 算两次）和差一（decode 时上下文是 ctx 还是 ctx + 1）。最可靠的检验办法是：写一个真的会跑的模型，在它运行时数一数，再和公式比。

记账的约定（transformer_lab.py 开头有完整说明）：

• 矩阵乘法 $x[M, K] \cdot W[K, N]$：FLOPs = $2MKN$；权重字节 = $KN \times$ 每个数的字节；激活字节 = $(MK + MN) \times$ 每个数的字节。
• 注意力核心按融合内核记账（分数矩阵不写回显存）：执行的 FLOPs = $4BTSHd_h$；有用的 FLOPs 只算因果掩码允许的部分；KV 字节 = 读一遍 K 和 V。
• 字节一律按 BF16（2 字节）计，和 numpy 实际用的 float64 无关。

4.2 参数：公式和数组一一对应

代码 03 演示 1，小模型（4 层、$d = 128$、8 个 query 头、2 组 KV、$d_h = 16$、$f = 344$、$V = 512$）逐个数组数出来的参数，和第 1 讲 llm_math 的公式完全相等：每层注意力 40,960、每层 MLP 132,096、全部 824,448、参与矩阵乘法的 757,760。

用同样的公式看 Llama-3.1 的参数花在哪：

MLP 是大头；注意力只占约 17%（GQA 把 K、V 投影压到 Q 投影的 1/4–1/16）；模型越小，词表占比越大。

4.3 一次前向的账，和第 1 讲对账

代码 03 演示 2 在小模型上跑了一次 prefill（$B = 1$、$T = 128$）和一次 decode（$B = 8$、每个序列已有 127 个 token），逐项和 analytic_ops 比较——完全相等。再和第 1 讲的 llm_math 对账：

结论：第 1 讲的账本分毫不差，可以放心用在跑不动的大模型上。几个细节：

• decode 时公式里的 ctx 要取”包括新 token 在内”的长度（新 token 也要和自己做注意力）。
• 朴素实现”执行的”注意力 FLOPs 约是”有用的”两倍：先把 $T \times T$ 全算出来再掩掉一半。FlashAttention 一类的内核会跳过整块被掩掉的部分（第 11 讲）。
• 公式忽略了激活字节。prefill 时它很大：每个运算都读写一遍 $[T, d]$ 级别的张量（小模型 prefill 时，逐元素运算的激活字节 2.76 MB，比 MLP 的权重 1.06 MB 还多）。这部分要靠算子融合省掉（第 11、12 讲）。

4.4 注意力的运算什么时候不能忽略

每个 token 的矩阵乘法 FLOPs 是常数 $2 \times$ matmul_params，注意力核心是 $4 L H d_h \cdot \text{ctx}$，随上下文线性增长。两者相等时：

代码 03 演示 4：

（表中百分比是注意力核心占每 token 总 FLOPs 的比例。）模型越大，矩阵乘法越重，注意力越晚成为运算的主要部分。How to Scale Your Model 对 MHA、$F = 4D$ 的模型给出”$T > 8D$ 时注意力占主导”的经验式，是同一个道理。

但FLOPs 不是 decode 的瓶颈。哪怕注意力的 FLOPs 很少，读 KV cache 的字节也可能远超读权重（§4.5）。

4.5 prefill 和 decode 的计算图

把每一类运算单独套 roofline（Llama-3.1-8B，H100，BF16，屋脊点 295 FLOP/byte），代码 03 演示 5：

prefill，提示 2048 个 token：

decode，$B = 1$，上下文 2048：所有运算的强度都是 1–4，整张图都是 memory-bound，时间几乎全是读权重（MLP 3.37 ms、注意力投影 0.80 ms、lm_head 0.31 ms）。

decode，$B = 64$，上下文 2048：

读这三张表：

1. prefill 的矩阵乘法是 compute-bound（强度上千），但逐元素运算单独执行时是 memory-bound 的，一共 3.7 ms，相当于矩阵乘法时间（约 30 ms）的 12%。逐项相加比第 1 讲”整体一个 roofline”的下界多 6%（34.1 vs 32.2 ms）——差得没有 3.7 ms 那么多，是因为第 1 讲的 prefill_time 把 lm_head 算在了每个提示 token 上（多算约 2.2 ms），而生成只需要最后一个 token 的 logits。融合逐元素运算是内核和编译器的工作。
2. decode 加大批，投影和 MLP 的强度随 $B$ 上涨（约 60），注意力核心的强度却停在 4.0：恰好等于组大小 $R = H/G = 4$。每个序列读自己的 KV，只被自己的 4 个 query 头用，批再大也不变（How to Scale Your Model 第 4 章问题 4 的结论：decode 时注意力的强度 $\to G$，这里的 $G$ 是书里的组大小）。
3. 上下文才 2048、批 64，读 KV 的注意力核心（5.1 ms）就已经比整个 MLP（3.4 ms）还慢了。

5. 推理：prefill、decode 与 KV cache · Inference with a KV cache

代码：04_kv_cache_decode.py

5.1 自回归生成

生成一段文字的过程：

1. prefill：把整段提示一次送进模型，得到最后一个位置的 logits，选出第一个新 token。
2. decode：把新 token 接到序列末尾，再算下一个位置的 logits，选出下一个 token……重复，直到生成结束符或达到长度上限。

最笨的做法是每一步都把”到目前为止的整个序列”重新算一遍。第 $t$ 步要处理 $T_0 + t$ 个 token，生成 $n$ 个 token 一共处理约 $n T_0 + n^2/2$ 个 token：随长度平方增长。

5.2 为什么能缓存

§2.5 已经看到：有因果掩码时，位置 $j$ 的输出只依赖 token $0 \ldots j$。一层一层往上推，位置 $j$ 在每一层的 $k_j$、$v_j$ 也只依赖 token $0 \ldots j$。所以生成新 token 时，旧 token 的 $k$、$v$ 不会变，可以存起来：

• prefill 时，把提示里每个 token 在每一层的 $k$、$v$ 写进 KV cache；
• decode 时，每一步只算 1 个新 token 的 $q$、$k$、$v$，把新的 $k$、$v$ 追加进 cache，再用新 $q$ 和 cache 里的全部 $k$、$v$ 做注意力。

为什么只缓存 K、V，不缓存 Q？因为旧 token 的 $q$ 以后再也用不到：新 token 只需要拿自己的 $q$ 去查所有人的 $k$。

5.3 正确性，以及两个常见 bug

原课作业 1 要求写一个带 KV cache 的 decode 循环。代码 04 演示 1 把它和”每步重算整个序列”的参照逐步比较（4 个序列，提示 16 个 token，贪心生成 24 个）：

正确的实现和参照只差 float64 的舍入误差：两者在数学上是同一个计算，只是换了计算顺序。两个 bug 都不会报错，第 0 步（由 prefill 直接算出）也是对的，从第 1 步起输出就跑偏了。第二个 bug 特别隐蔽：新 token 以为自己在位置 0，RoPE 算出的相对位置全错，模型照样能生成”像样”的文字，只是质量下降。所以测试要和参照逐步比较 logits，而不只是看生成的文字。

5.4 省了多少

代码 04 演示 2，一个稍大的模型（$d = 256$，4 层，float32），提示 64 个 token：

（时间是编写时的墙钟时间，只看趋势；两种做法生成的 token 完全相同。）

两个观察：

• FLOPs 从平方变成线性，生成越长差距越大。用公式推到 Llama-3.1-8B：提示 1000、生成 8000 个 token，不用 cache 要多花约 4,200 倍的运算。没有 KV cache，长输出的推理模型根本跑不起来。
• 墙钟时间的比值比 FLOPs 的比值小得多（139× vs 45×）。用 cache 时每一步只处理 1 个 token：每次调用都有固定开销，矩阵乘法也退化成向量乘矩阵（§2.1 的 RNN 问题又回来了）。FLOPs 省了，但每个 FLOP 变贵了——这就是 decode 的本质，也是第 4–7 讲各种技巧要对付的问题。

5.5 代价：显存

每个 token 每层存一份 $k$ 和 $v$，每份 $G \times d_h$ 个数：

代码 04 演示 3：

8B 模型只要约 12 万个 token 的 KV 就和全部权重一样大；128K 上下文时，8 个序列的 KV 是权重的 8 倍多。KV cache 是用显存换运算，而显存决定了能同时服务几个请求（第 1 讲 §4.4）。

5.6 批量 decode：权重读一次，KV 各读各的

代码 04 演示 4 在小模型上实际运行并记账（上下文 128）：

读权重的字节不随 $B$ 变化（所有序列共享），读 KV 的字节和 FLOPs 都和 $B$ 成正比。强度 $= B f / (W + B k)$ 随 $B$ 增大，但有上限 $f / k$。Llama-3.1-8B 在 4096 上下文时，这个上限约 32 FLOP/byte，远低于 H100 的 295：$B = 1024$ 时强度 31.1，一步下界 169 ms，每芯片 6,074 tok/s，比 $B = 256$ 只多 8%。这就是第 1 讲 §4.3”批处理摊不薄 KV”的结论，在这里由实际运行的模型逐字节数了出来。

要突破这个上限，只能让每个序列的 KV 变小或少读——下一节。

6. 为推理而改的注意力：MQA、GQA、MLA · Attention variants built for inference

代码：05_attention_variants.py

6.1 MQA：所有头共用一组 K/V

Shazeer（2019）在 Fast Transformer Decoding: One Write-Head is All You Need 里分析了增量解码（incremental decoding）：多头注意力每一步要读的 K、V 字节和运算量之比是 $\Theta(n/d + 1/b)$（$n$ 是序列长度，$b$ 是批大小）。$1/b$ 一项可以靠加大批来压，$n/d$ 一项却和批无关——这正是 §5.6 的”上限”。他的办法是多查询注意力（MQA）：保留多个 query 头，但所有头共用一组 K、V。KV cache 缩小 $H$ 倍，上面那一项变成 $n/(dh)$。

代价是质量：MQA 可能掉点、训练不稳定。

6.2 GQA：折中

GQA（Ainslie et al. 2023）介于两者之间：$H$ 个 query 头分成 $G$ 组，每组共用一组 K、V。$G = H$ 就是 MHA，$G = 1$ 就是 MQA。Llama 3 全系用 $G = 8$。

代码 05 演示 1 说明了 GQA 的本质：把 GQA 的 K/V 投影复制给组里每个头，就得到一个输出相同的 MHA（差 $6 \times 10^{-15}$）。GQA 是”强制同组的头共用 K/V”的 MHA，表达能力是 MHA 的子集，换来的是不必存那些重复的 K/V。在小模型上实际运行并记账（$H = 8$，decode $B = 16$，上下文 128）：

注意力核心的强度恰好等于组大小：一组 K/V 读进来，被组里每个 query 头用一次。参数量几乎不变——GQA 省的是推理时的显存和带宽，不是模型大小。

6.3 从 MHA 转换：均值池化 + 补训练

已经用 MHA 训练好的模型，能不能改成 GQA？GQA 论文的办法：把同一组里各头的 K、V 投影矩阵取平均（他们比较过，比”只保留一个头”或”随机初始化”都好），再用原预训练算力的 5% 继续训练（uptraining）。论文报告 T5-XXL 的结果：MHA 推理时间 1.51 ms、平均分 47.2；转成 GQA-8 后 0.28 ms、47.1；转成 MQA 后 0.24 ms、46.6——质量接近 MHA，速度接近 MQA。

为什么需要补训练？代码 05 演示 2 构造了组内各头 K/V 投影相关系数为 $\rho$ 的 MHA，直接均值池化（不训练）：

（每格：与原模型输出分布的平均 KL 散度 / top-1 一致率。每一列是一个组结构和目标一致的不同随机模型，所以列之间有小的波动。）头完全相同时转换是精确的；头越不像、组越大，丢的信息越多。真实模型的头之间有冗余但不完全相同，所以要补训练把损失补回来。

6.4 MLA：只缓存一个潜在向量

DeepSeek-V2/V3 的多头潜在注意力（multi-head latent attention, MLA）换了一个思路：不减少头数，而是把 K、V 联合压缩成一个低维的潜在向量 $c$：

只缓存 $c$（DeepSeek-V3：$d_c = 512$），需要时再”解压”出每个头的 $k$、$v$。但每一步都解压整个上下文太贵了。关键的一步是用矩阵结合律”吸收”上投影：

先把 $W_{UK}$ 乘进 query，在潜在空间里直接和 $c$ 做点积；加权求和也在潜在空间里做，最后再上投影。代码 05 演示 3 实现了两种写法，输出只差 $10^{-15}$。

为什么 RoPE 要单独留一份：RoPE 按位置旋转，$(R_m q)^\top (R_n W c) = q^\top R_m^\top R_n W c$，中间的 $R_m^\top R_n$ 随位置变化，不能并成一个固定矩阵。所以 DeepSeek 把 key 拆成两部分：不带 RoPE 的部分走潜在向量（可以吸收），带 RoPE 的 $d_R = 64$ 维单独计算、所有头共享一份、单独缓存。每 token 每层缓存 $512 + 64 = 576$ 个数。

三种实现的账（DeepSeek-V3 的尺寸：128 头、$d_h = 128$；一个 token、一层、上下文 32768）：

“吸收”两头都占：字节是 MHA 的 1/71，运算量约是 3.4 倍，强度约 240，接近 H100 的屋脊点 295。MLA 用 FLOPs 换字节，把 decode 注意力从”严重 memory-bound”推到了接近平衡。

6.5 四种变体放在一起

代码 05 演示 4：同一个注意力尺寸（DeepSeek-V3：61 层、128 头、$d_h = 128$），BF16 KV，上下文 32K，批大小 32，一块 H100 上一步 decode 的注意力部分：

• 这么宽的注意力用 MHA，一个 32K 序列的 KV 就有 131 GB，一块 H100 都放不下。
• MLA 的 KV 介于 GQA-8 和 MQA 之间：576 个数 / token / 层，相当于 $576 / 256 = 2.25$ 组 KV 头的 GQA。DeepSeek-V3 技术报告说它的质量”和标准 MHA 相当”。
• MLA 是四种里唯一让读 KV 和计算几乎相等的：代价是计算时间是其他三种的 4 倍多。
• 一个对第 4 讲的预告：MLA 的潜在向量是所有头共享的，按头切分的张量并行切不开它，每块卡都要存一整份。这让 MLA 更适合按批切分注意力（DeepSeek 的部署就是这样做的）。

【系统联系】GQA、MQA、MLA 都是模型结构为硬件让步的例子：它们不是为了更好的质量，而是为了 decode 的带宽和显存。第 3 讲的 KV 量化（FP8 / FP4）能在此基础上再省一半到四分之三。

7. 现状（2025–2026）· State of the field

• GQA 是默认：Llama 3、Qwen、Mistral 等开放模型几乎都用 GQA（8 组上下）；MLA 被 DeepSeek-V2/V3/R1 和一些后来的模型采用。
• 注意力变得稀疏或混合：长上下文模型越来越多地混用滑动窗口注意力、线性注意力 / 状态空间模型（SSM）和少量全注意力层，或用学习出来的稀疏注意力只读一部分 KV（例如 DeepSeek-V3.2 的稀疏注意力）。它们攻的都是 §5.6 的那个上限。本仓库 CS349D 第 10 讲专门讲这一类。
• RoPE 外推：为了让模型处理比训练时更长的上下文，出现了一系列 RoPE 缩放方法（位置插值、NTK、YaRN 等），本质上是调整各对数的旋转频率。
• 词表变大：Llama 3 的词表是 128K（Llama 2 是 32K）。更大的词表让同样的文字切成更少的 token（推理更快），代价是 lm_head 更大——对小模型来说是不小的一笔（§1.8）。

8. 小结 · Summary

1. 词向量来自共现：word2vec 用”中心词预测上下文”的简单任务学出稠密向量，关系（性别、首都、比较级）被学成近似恒定的方向。它的设计哲学是先算账：去掉隐藏层、用负采样替代完整 softmax，让 $Q$ 小两个数量级，换来更多数据。嵌入在实现上是查表，lm_head 却是整张矩阵乘。
2. 注意力是可导的软查找：$\text{softmax}(QK^\top/\sqrt{d_k})V$。$\sqrt{d}$ 防止 softmax 饱和；因果掩码让位置 $j$ 只依赖前面的 token（KV cache 的基础）；多头不增加运算量；位置要单独注入，RoPE 让分数只依赖相对位置。
3. Transformer 胜出的系统原因是并行：训练时所有位置一起算，全是大矩阵乘法。decode 却退回了”一步一个 token”。
4. 账本可以逐项验证：真实运行的模型数出来的参数、FLOPs、权重字节、KV 字节和第 1 讲的公式完全相等。prefill 的矩阵乘法 compute-bound，逐元素运算 memory-bound；decode 整张图 memory-bound，注意力核心的强度停在组大小 $H/G$。
5. KV cache 用显存换运算：把 FLOPs 从平方变成线性，但每个 token 每层都要存一份 K、V，几万到几十万个 token 就和权重一样大；批量 decode 的强度有上限 $f/k$。
6. 注意力变体是为推理而生：MQA、GQA 让 query 头共用 K/V，强度 = 组大小；MLA 只缓存一个潜在向量并用结合律吸收上投影，用 FLOPs 换字节，强度接近屋脊点。

9. 常见误区 · Common pitfalls

1. “word2vec 的向量里，king 的最近邻应该是 queen。” 最近邻通常是话题词（throne、palace）。关系藏在方向里，所以要用类比来评价（§1.6）。
2. “多头注意力比单头贵 $h$ 倍。” 每个头的维度是 $d/h$，总运算量不变（§2.6）。
3. “除以 $\sqrt{d}$ 是为了数值稳定，推理时可以去掉。” 它改变了 softmax 的输出，是模型定义的一部分；它的代价是零，因为可以并进 $W_q$。
4. “KV cache 是一种近似，会损失精度。” 它和重算整个序列在数学上是同一个计算，只差浮点舍入（$10^{-15}$），前提是因果掩码（§5.3）。
5. “KV cache 也应该缓存 Q。” 旧 token 的 $q$ 以后再也用不到，新 token 只用自己的 $q$ 去查所有人的 $k$。
6. “GQA 减少了注意力的运算量。” 它减少的是 K/V 投影和 KV cache 的字节，注意力核心的 FLOPs 不变（query 头数没变）。
7. “批大小足够大，decode 就会 compute-bound。” 读 KV 的字节和批大小同比例增长，强度有上限 $f/k$；长上下文时这个上限远低于屋脊点（§5.6）。
8. “注意力的 FLOPs 占比小，所以长上下文不贵。” decode 的瓶颈是字节，不是 FLOPs：上下文 2048、批 64 时，注意力核心的 FLOPs 只占约 7%，时间却是最长的一项（§4.5）。
9. “MLA 就是把 KV 压缩存储、用时再解压。” 如果每一步都解压，运算量会大几百倍；MLA 的关键是吸收，直接在潜在空间里算（§6.4）。
10. “tied embedding 能让 decode 少读 lm_head。” 共享省的是参数和显存；lm_head 还是要整张读、整张乘（§1.8）。

10. 练习 · Exercises

带 💻 的题可以用本讲的模块验证。参考答案在每题后面。

练习 1（参数账）💻。一个 MHA 模型：$L = 64$、$d = 4096$、32 个头、$d_h = 128$、$f = 16384$（两矩阵或三矩阵你自己选一个并说明）、$V = 32000$、不共享 lm_head。按 SwiGLU（三个 MLP 矩阵）算总参数，注意力参数占多少？每 token 的 KV cache 用 int8 存是多少？（这是 How to Scale Your Model 第 4 章的问题 1。）

ModelConfig("q1", 64, 4096, 32, 32, 128, 16384, 32000).n_params() = 17.44 B（书里用 $4 \times 10^3$、$16 \times 10^3$ 这样的整数近似，得到约 16 B）。注意力占 24.6%，接近书里的 1/4：$4d^2 / (4d^2 + 3 \cdot 4d \cdot d) = 1/4$。KV / token（int8）$= 2 \times 64 \times 32 \times 128 = 524{,}288$ 字节 = 512 KiB。

练习 2（word2vec 的代价）。skip-gram + 负采样，$C = 5$、$D = 300$、$k = 15$。每个训练词的 $Q$ 是多少？和同样设置下的完整 softmax（$V = 10^6$）比呢？

负采样：$Q = C \cdot (D + D \cdot (k+1)) = 5 \times (300 + 300 \times 16) =$ 25,500。完整 softmax：$Q = C \cdot (D + D \cdot V) \approx 5 \times 3 \times 10^8 = 1.5 \times 10^9$，大约 6 万倍。

练习 3（einsum FLOPs）💻。$A[I, J, K, L] \cdot B[I, J, M, N, O] \to C[K, L, M, N, O]$ 要多少 FLOPs？取 $I..O = 2, 3, 4, 5, 6, 7, 8$ 用 einsum_flops 验证。

$I$、$J$ 是被求和的维度，没有批维，FLOPs $= 2 \cdot IJKLMNO$。einsum_flops("ijkl,ijmno->klmno", (2,3,4,5), (2,3,6,7,8)) = 80,640 = $2 \times 2 \times 3 \times 4 \times 5 \times 6 \times 7 \times 8$。

练习 4（GQA 的强度上限）。How to Scale Your Model 给出 decode 时 GQA 注意力的强度是 $S R / (R + S)$（$S$ 是上下文长度，$R$ 是组大小，书里写作 $G$）。Llama-3.1-70B（$R = 8$）在 $S = 128$、$1024$、$32768$ 时分别是多少？加大批能提高它吗？

分别约 7.53、7.94、8.00，迅速趋于 $R = 8$。批大小不出现在公式里：每个序列读自己的 KV，加大批对注意力核心的强度没有帮助（§4.5、§6.2）。要提高它，只能增大组大小（更少的 KV 头）或换成 MLA。

练习 5（KV 容量）💻。Llama-3.1-70B，上下文 128K，4 个序列，KV 用 FP8（每个数 1 字节）。KV cache 一共多大？BF16 呢？和 FP8 权重（70.6 GB）比一比。

每 token $2 \times 80 \times 8 \times 128 = 163{,}840$ 字节（FP8）。$4 \times 131072 \times 163840 =$ 85.9 GB；BF16 是 171.8 GB。即使 KV 用 FP8，4 个 128K 的序列也比整个 FP8 模型还大。

练习 6（KV cache 省了多少）💻。Llama-3.1-8B，提示 2000 个 token，生成 2000 个。用 cache 和不用 cache（每步一次完整前向）的总 FLOPs 分别是多少？仿照 04_kv_cache_decode.py 演示 2 的写法计算。

用 cache：prefill + 1999 次 decode ≈ 64.2 TFLOP；不用 cache：2000 次逐渐变长的完整前向 ≈ 94.9 PFLOP，约 1,480 倍。

练习 7（MLA 的 KV）💻。DeepSeek-V3（61 层，$d_c = 512$，$d_R = 64$）每 token 的 KV 用 BF16 是多少？FP8 呢？如果它用 128 头的 MHA（$d_h = 128$），BF16 是多少？

MLA：$61 \times 576 \times 2$ 字节 = 68.6 KiB（BF16），34.3 KiB（FP8）。MHA：$2 \times 61 \times 128 \times 128 \times 2$ = 3.81 MiB，是 MLA 的约 57 倍。

练习 8（RoPE 的相对性）。证明：若把二维向量 $q$ 在位置 $m$ 旋转 $m\theta$、$k$ 在位置 $n$ 旋转 $n\theta$，则旋转后的点积只依赖 $m - n$。为什么这对 KV cache 很方便？

旋转矩阵 $R_\phi$ 满足 $R_a^\top R_b = R_{b-a}$。$(R_{m\theta} q)^\top (R_{n\theta} k) = q^\top R_{m\theta}^\top R_{n\theta} k = q^\top R_{(n-m)\theta} k$，只依赖 $n - m$。对 KV cache 的好处：$k$ 旋转一次之后写进 cache 就不用再改，新 token 不管在哪个位置，只要把自己的 $q$ 按自己的位置旋转即可。（如果位置是”加”到输入上的，效果也一样；但如果某种位置编码要求”所有 key 随当前长度重新计算”，KV cache 就失效了。）

练习 9（双向模型）。BERT 这样的双向编码器能不能用 KV cache 逐个生成 token？为什么？

不能。双向注意力里，每个位置的表示依赖所有位置（包括后面的），新加一个 token，前面所有 token 在每一层的 $k$、$v$ 都会变（代码 02 演示 3：改一个 token，双向模型里前面的位置也变了 0.80）。缓存的旧值失效，只能整段重算。这也是生成式模型都用因果注意力的原因之一。

练习 10（思考题：先融合哪个）。§4.5 的 prefill 表里，逐元素运算一共 12.35 GB、3.69 ms。如果只能融合一类，你会先把 RMSNorm 融合进后面的矩阵乘法，还是先把 SwiGLU 的 SiLU 和逐元素乘法融合进 gate/up 矩阵乘法？用 analytic_ops 查出各自的字节数再回答。

用 analytic_ops(LLAMA31_8B, 1, 2048, 0, last_only=True) 查各类逐元素运算的激活字节：swiglu 5.64 GB（$3 \times 2048 \times 14336 \times 2 \times 32$），residual 3.22 GB，norm 2.15 GB，rope 1.34 GB（softmax 在融合注意力里不计字节），合计 12.35 GB。SwiGLU 的中间张量最大（$f = 3.5d$），先融合它最划算；第二大的其实是残差加法，它可以和前面的 o_proj / down 矩阵乘法的”写回”融合在一起。真实系统（以及第 12 讲的编译器）这些都会做。

11. 延伸阅读 · Further reading

本讲阅读材料

• T. Mikolov, K. Chen, G. Corrado, J. Dean, Efficient Estimation of Word Representations in Vector Space, ICLR Workshop 2013（arXiv:1301.3781）。
• A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. Gomez, Ł. Kaiser, I. Polosukhin, Attention Is All You Need, NeurIPS 2017（arXiv:1706.03762）。

本讲补充材料

• T. Mikolov et al., Distributed Representations of Words and Phrases and their Compositionality, NeurIPS 2013（arXiv:1310.4546）。—— 负采样、高频词下采样。
• J. Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021（arXiv:2104.09864）。
• N. Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019（arXiv:1911.02150）。—— MQA 和增量解码的内存分析。
• J. Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023（arXiv:2305.13245）。
• DeepSeek-AI, DeepSeek-V3 Technical Report, 2024（arXiv:2412.19437）。—— §2.1.1 的 MLA。
• J. Austin et al., How to Scale Your Model, 第 4 章 All the Transformer Math You Need to Know。—— einsum 记号、FLOPs 规则、KV cache 的账。

本仓库的相关讲义

• CS349D 前置讲 P2：Transformer 推理的账本与采样，和本讲 §4–§5 互补。
• CS349D 第 10 讲：混合架构与稀疏注意力。

12. 下一讲预告 · Next lecture

第 3 讲（Priyanka Raina）讲精度、量化与稀疏，阅读材料是 LLM.int8() 和 Switch Transformers。本讲的结论是”decode 的时间花在搬权重和 KV 上”；第 3 讲的问题是：每个数能不能少用几个比特？能不能少读一些参数？ 我们会在本讲的 transformer_lab 模型里插入量化的矩阵乘法（linear_fn 这个钩子就是为它留的），亲眼看到”离群特征让 int8 崩溃、LLM.int8() 的混合精度分解把它救回来”；再看 2:4 结构化稀疏和混合专家（MoE）——“每个 token 只用一部分参数”在 roofline 上是什么样子。