拆开 LLM（二）：手撕注意力——后来的一切优化都在救这一条公式

2026-07-17

llm-anatomy 是一个"把 LLM 拆开看"的教学连载：不调库、不背黑话，每一章用 numpy 把大模型的一个器官亲手造一遍，所有结论都配可复现的实验数字。
上一篇《拆开 LLM（一）：mini-autograd》我们手写了自动求导，解决了"梯度从哪来"。这一篇进入正题中的正题：注意力。

先看两个数字

原理放后面，先给你两个本文实测的数字：

一个 64 维的单层注意力，逐 token 生成 1024 个 token：每步全量重算，全程 1024 步累计要 1566.4 ms；挂上 KV cache 全程只要 19.2 ms——81.5 倍。
把 √d 缩放去掉，维度加到 1024：softmax 输出的平均熵从 3.688 塌到 0.133，最大概率飙到 0.947——注意力还没开始训练，就已经"死认一个 token"了。

这两个数字背后是同一条公式：

$$\text{Attention}(Q, K, V) = \text{softmax}!\left(\frac{QK^\top}{\sqrt{d}}\right)V$$

softmax 注意力这条公式本身，2017 年定稿之后一个字都没改过。但 2026 年你能叫得出名字的推理优化——KV cache、GQA、MLA、FlashAttention——几乎全部是在这条公式的不同环节上做手术。这一章我们把公式亲手写出来，再把每台手术钉回它动刀的位置。

2014 年：注意力不是为大模型发明的

先考个古，因为注意力的出身经常被讲反。

2014 年的机器翻译主流是 Seq2Seq：一个编码器 RNN 把整句源语言读完，压成一个固定长度的向量，再由解码器 RNN 从这个向量里把译文一个词一个词吐出来。

问题很直白：一句 50 个词的句子，凭什么能无损地塞进一个几百维的向量？翻到译文第 40 个词的时候，你需要的源句细节早就被挤没了。当时的实验也确实显示：句子越长，翻译质量掉得越狠。这就是著名的信息瓶颈。

Bahdanau 等人的解法（也就是后来所谓的"注意力"）思路极其朴素：别压缩了，留着。编码器把每个源词的隐状态全部保留；解码器每吐一个词之前，先拿自己当前的状态（query）去和所有源词的隐状态（keys）逐一打分，softmax 归一化成权重，再对源词信息（values）加权求和——等于每一步都现场决定"我现在该重点看源句的哪几个词"。

所以注意力最初是台对齐机器：它的产出物就是一张"译文第 i 个词 ↔ 源文第 j 个词"的软对齐表。论文里那张热力图（法语英语词序对调被清晰地对齐出来）比 BLEU 分数更有说服力。它当时的身份是翻译系统的一个配件，离"大模型的心脏"还差着三年。

手撕三种打分函数

“拿 q 去和每个 k 打分"这件事，2014–2017 年间出现过三种主流写法。全部代码在仓库 ch02/attention.py，统一接口：q 是 (d_q,) 的单条查询，K 是 (n, d_k) 的 n 条键，返回 (n,) 的分数。

先把地基打好——softmax：

1
2
3
4
5
def softmax(x, axis=-1):
    """数值稳定版：先减最大值再 exp，防止 exp 溢出"""
    x = x - x.max(axis=axis, keepdims=True)  # 平移不改变 softmax 结果
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)

那行减最大值不是洁癖：np.exp(800) 直接溢出成 inf，softmax 输出变 nan，整个模型当场去世。softmax 对整体平移免疫，减掉最大值后指数最大是 exp(0)=1，怎么都炸不了。这是本文第一个"不做就会死"的细节，后面还有一个。

加性 / Bahdanau（2014）

1
2
3
4
5
6
7
def additive_score(q, K, W_q, W_k, v):
    """加性 / Bahdanau (2014): score = v · tanh(W_q q + W_k k)

    q: (d_q,)  K: (n, d_k)  W_q: (h, d_q)  W_k: (h, d_k)  v: (h,)
    """
    h = np.tanh(W_q @ q + K @ W_k.T)  # (n, h) 广播：W_q q 加到每一行
    return h @ v                      # (n,)

本质是把 q 和 k 拼起来喂给一个单隐层小神经网络，让网络自己学"什么叫相关”。表达力最强，还天然支持 q、k 维度不同。

乘性 / Luong general（2015）

1
2
3
4
5
6
def multiplicative_score(q, K, W):
    """乘性 / Luong general (2015): score = q W k

    q: (d_q,)  K: (n, d_k)  W: (d_q, d_k)
    """
    return K @ (W.T @ q)              # (n,) 先算 W^T q，避免 K@W 大矩阵

把小神经网络换成一个双线性矩阵 W——“相关性"由一次矩阵变换定义。参数比加性少，算起来就是两次矩阵乘。

缩放点积 / Vaswani（2017）

1
2
3
4
5
6
def scaled_dot_score(q, K):
    """缩放点积 / Vaswani (2017): score = q·k / sqrt(d)，要求 d_q == d_k

    q: (d,)  K: (n, d)
    """
    return K @ q / np.sqrt(q.shape[-1])  # (n,)

连 W 都不要了：直接点积。“相关"被定义为"方向相近”。（Transformer 里可学的部分挪到了生成 Q/K/V 的投影矩阵上，打分本身零参数。）

点积为什么赢

论《表达力》，加性最强；Transformer 论文自己也承认，d 大的时候不带缩放的点积效果不如加性。那点积凭什么通吃至今？

因为它不是数学上最优，而是硬件上最优。 把 n 条 query 摞成矩阵，缩放点积对所有 (q, k) 对的打分就是一句 Q @ K.T——一次标准矩阵乘。而矩阵乘是整个计算机工业被优化得最狠的运算：BLAS 库、GPU、后来的 Tensor Core，全是围着它建的。加性打分每个 (q, k) 对都要过一遍 tanh 小网络，中间张量是 (n_q, n_k, h) 三维的，吃显存又难以榨干矩阵乘单元。

同样的分数量级，点积能把硬件跑到接近峰值，加性只能干瞪眼。**深度学习的历史一再重演这个剧本：赢的不是最聪明的公式，是最贴合硬件的公式。**这个视角记住，本文结尾讲 FlashAttention 时它会再次出现。

√d 不是装饰，是保命符

公式里那个不起眼的 √d，是第二个"不做就会死"的细节。

推一下：设 q、k 的每个分量独立同分布、均值 0 方差 1，则 $q \cdot k = \sum_{i=1}^{d} q_i k_i$ 是 d 个独立随机变量之和——方差恰好是 d。维度翻倍，logits 的离散程度就翻倍。而 softmax 里的差距是指数放大的：logits 只要拉开五六个标准差，最大那项就几乎独占所有概率。

空口无凭，跑 ch02/scale_demo.py：随机 q 和 64 条随机 k，d 从 16 加到 1024，比较缩放前后的 logits 方差、softmax 熵和最大概率（256 次重复取均值）：

$ uv run python ch02/scale_demo.py
每个 softmax 覆盖 64 个 key，均匀分布的熵 = ln(64) = 4.159 nats，重复 256 次取均值

    d |   var(qk) | var(qk/sqrt d) |   H_raw | H_scaled | maxP_raw | maxP_scaled
--------------------------------------------------------------------------------
   16 |     16.04 |          1.003 |   1.362 |    3.681 |    0.583 |       0.109
   64 |     63.89 |          0.998 |   0.575 |    3.675 |    0.792 |       0.112
  256 |    251.79 |          0.984 |   0.265 |    3.697 |    0.897 |       0.104
 1024 |   1013.87 |          0.990 |   0.133 |    3.688 |    0.947 |       0.108

三个观察：

var(qk) 一列就是理论值 d 本身：16.04、63.89、251.79、1013.87——推导没骗你。
不缩放：d=1024 时熵只剩 0.133 nats，最大概率 0.947——64 个候选里 softmax 几乎把全部注押给一个。这叫饱和：注意力权重还没学过一步，就已经退化成硬指针了。
除以 √d 之后：方差钉死在 1.0 左右，熵稳在 3.68，最大概率 0.11——维度随便加，softmax 的"锐度"不动。

饱和为什么致命？softmax 的梯度大小由 $p_i(1 - p_i)$ 这类项主导：p 越接近 0 或 1，梯度越接近 0。饱和的 softmax 是一堵梯度墙——上游的 Q、K 投影矩阵收不到任何有效梯度，学不动。Transformer 论文只用一个脚注交代这件事，但你现在有了表格：√d 不是美学选择，是让注意力在初始化时就处在"可学习区间"的保命符。

从对齐机器到 self-attention

到这里注意力还是台"对齐机器”：q 来自解码器，k、v 来自编码器，两边是两拨人。2017 年《Attention Is All You Need》干的事，一句话概括：让 Q、K、V 全部来自同一个序列——每个 token 既提问、也应答、也供货。这就是 self-attention。（考据一句：self-attention 的概念 2016–2017 年已在多篇工作里出现，AIAYN 的贡献不是发明它，而是把它从配件扶正为整个架构的主角。）

身份从此反转：注意力不再是 RNN 的外挂配件，而是把 RNN 整个踢出局，自己当了架构主角。序列内任意两个位置直接一步交互，不用像 RNN 那样隔 100 个词传 100 次话。

完整版实现（ch02/attention.py）：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
def scaled_dot_product_attention(Q, K, V, mask=None):
    """完整缩放点积注意力：softmax(Q K^T / sqrt(d)) V

    Q: (n_q, d)  K: (n_k, d)  V: (n_k, d_v)  mask: (n_q, n_k) True=可见
    返回: (输出 (n_q, d_v), 注意力权重 (n_q, n_k))
    """
    d = Q.shape[-1]
    logits = Q @ K.T / np.sqrt(d)                 # (n_q, n_k)
    if mask is not None:
        logits = np.where(mask, logits, -np.inf)  # 不可见处 = -inf，softmax 后严格为 0
    w = softmax(logits, axis=-1)                  # (n_q, n_k) 每行和为 1
    return w @ V, w                               # (n_q, d_v)

mask 参数就是 GPT 系模型"只能看过去"的全部实现——一个下三角布尔阵：

1
2
3
def causal_mask(n):
    """(n, n) 下三角布尔阵，True = 可见（只看得到自己和过去）"""
    return np.tril(np.ones((n, n), dtype=bool))

这十几行代码里埋着几个我见人踩过无数次的坑，逐条编号：

坑 1：mask 用 0 而不是 -inf。 把不可见位置的 logit 置 0？0 是个正常分数，softmax 照样分给它概率。必须是 -np.inf，exp 之后才严格为 0。

坑 2：mask 加在 softmax 之后。 先 softmax 再把未来位置抹零，每行的概率和就不是 1 了，输出整体被缩小，而且缩小的比例还随位置变化。mask 必须动在 logits 上。

坑 3：softmax 忘了减最大值。 前面说过，logits 一大直接 nan。尤其是配合坑 1 的 -inf，减最大值的写法恰好也能正确处理它。

坑 4：忘了 √d。 不会报错，不会 nan，训练也能跑——只是学得又慢又差，而且维度越大越差。最贵的 bug 就是这种不报错的 bug。

自检输出（uv run python ch02/attention.py）——下三角有值、上三角严格为 0、每行和为 1：

causal 注意力权重（保留 2 位）:
 [[1.   0.   0.   0.   0.   0.  ]
 [0.48 0.52 0.   0.   0.   0.  ]
 [0.66 0.14 0.2  0.   0.   0.  ]
 [0.23 0.3  0.13 0.34 0.   0.  ]
 [0.08 0.08 0.45 0.26 0.13 0.  ]
 [0.24 0.23 0.16 0.12 0.06 0.2 ]]

2026 年的推理优化全家桶，都在救同一条公式

现在到本章高潮。请你带着这条公式往下读：

$$\text{softmax}!\left(\frac{QK^\top}{\sqrt{d}}\right)V$$

KV cache：省掉"重算历史"这一步

自回归生成是一次加一个 token。关键观察：因为 causal mask 的存在，历史 token 的 K 和 V 在后续每一步里一个数都不会变——第 t 步唯一的新东西是新 token 自己的那一行 q、k、v。

那"每步把全序列的 Q/K/V 和整张注意力矩阵重算一遍"就是纯粹的浪费。把算过的 K、V 存下来（这就是 KV cache），每步只算新 token 的 q 对全体历史 K/V：

1
2
3
4
5
        x = xs[t - 1]                            # 只处理新 token
        q = x @ Wq                               # (D,)，O(D²)
        K[t - 1], V[t - 1] = x @ Wk, x @ Wv      # 追加进 cache
        s = K[:t] @ q / np.sqrt(D)               # (t,)，O(t·D)
        o = (softmax(s) @ V[:t]) @ Wo            # 新 token 天然只见过去，无需 mask

每步的注意力计算从 O(t²·d) 降到 O(t·d)——一整个数量级的 n 被消掉了。ch02/kv_cache_demo.py 用一个 64 维单层注意力做逐 token 生成，两条路径的输出经 np.allclose 验证逐位一致，计时如下：

$ uv run python ch02/kv_cache_demo.py
    n |   no-cache 总耗时 |   cache 总耗时 |    加速比 |  KV cache 显存
------------------------------------------------------------------
  128 |         4.7 ms |      1.6 ms |   2.9x |    65,536 B
  512 |       144.3 ms |      7.7 ms |  18.8x |   262,144 B
 1024 |      1566.4 ms |     19.2 ms |  81.5x |   524,288 B

最后一步（第 n 个 token）理论 FLOPs，O(n²) vs O(n):
    n |     no-cache | with cache |      比值
----------------------------------------------
  128 |    7,348,224 |     65,536 |    112x
  512 |   79,699,968 |    163,840 |    486x
 1024 |  293,609,472 |    294,912 |    996x

注意加速比的走势：2.9x → 18.8x → 81.5x，序列越长赚得越多——因为省掉的是平方项。每 token 耗时画出来更直观：无 cache 的蓝线是一条越来越陡的二次曲线，有 cache 的绿线贴地飞行：

Per-token latency: no cache vs KV cache

没有 KV cache，就没有今天"按 token 计费还能盈利"的推理服务。它是所有 LLM 推理引擎的第一块地基，没有之一。

代价：显存里躺着一头大象

天下没有白吃的午餐。KV cache 是拿空间换时间——上面表里最后一列，玩具模型的 cache 也随 n 线性长（1024 个 token 已占 524,288 字节）。放大到真实尺寸就吓人了。kv_cache_demo.py 顺手算了一个 7B 级模型（32 层、head_dim=128、fp16）的投影：

方案                       |    每token | n=  1024 | n=  4096 | n= 32768 | n=131072
--------------------------------------------------------------------------------------------
MHA (32 kv头)             |     512KB |    0.50G |    2.00G |   16.00G |   64.00G
GQA ( 8 kv头)             |     128KB |    0.12G |    0.50G |    4.00G |   16.00G
MLA (每层压成576维潜向量)        |      36KB |    0.04G |    0.14G |    1.12G |    4.50G

看 MHA 那行：上下文拉到 128k，光 KV cache 就要 64GB——比 7B 模型权重本身（约 14GB）大好几倍。你的显卡不是装不下模型，是装不下模型的"记忆"。于是刀落在公式的 K 和 V 上：

GQA（Grouped-Query Attention）：query 保留 32 个头，但 K/V 只留 8 个头，每 4 个 query 头共享一组 K/V。动的是K、V 矩阵的"份数"——cache 直接砍到 1/4（表里 64G → 16G）。LLaMA-2-70B 之后几乎成了开源标配；推到极致就是 MQA（全体共享 1 个 KV 头）。
MLA（Multi-head Latent Attention，DeepSeek 系）：不砍头数，砍K、V 本身的存储形态——每层把 token 压进一个约 576 维的低秩潜向量（严格说是 512 维压缩 KV 潜向量 + 64 维解耦 RoPE 分量），cache 里只存这个压缩包，用时再展开（数学上可以吸收进投影矩阵，不用真的解压）。表里同样 128k 上下文只要 4.5G，比 MHA 小一个数量级还多。

一句话总结这两位：**GQA 少存几份 K/V，MLA 把 K/V 压扁了存。**动机全在上面那张表里，不用背。

FlashAttention：公式一个字没改，改的是取数的路线

最后一位主角最反直觉：FlashAttention 对 $\text{softmax}(QK^\top/\sqrt{d})V$ 的数学零修改，输出和朴素实现精确等价（不是近似！），却能快好几倍、省一个数量级的显存。它改的是什么？

改的是内存访问顺序。朴素实现要把 n×n 的注意力矩阵完整写进 GPU 显存（HBM）再读回来——n=32k 时这是一个 10 亿元素的中间产物，而 GPU 的算力增速远快于显存带宽增速，瓶颈根本不在算，在搬。FlashAttention 的两板斧：

Tiling（分块）：把 Q、K、V 切成小块，装进 GPU 片上的 SRAM（快但小，类比 CPU 的 L1 cache）里算完一块再换下一块，n×n 矩阵从头到尾不落 HBM。
Online softmax：softmax 要"看全一行才能归一化"，看似和分块矛盾。解法是流式计算——边看边维护"目前为止的最大值和指数和"，每来一块就把之前的部分结果重新缩放一次修正。一行看完，归一化恰好也完成了。

顺带更新版本号：FlashAttention 已经迭代到第 4 代（2026 年 3 月的论文，面向新一代 GPU，宣称利用率显著提升），但原理没变——还是 tiling + online softmax，改访存，不改数学。

还记得"点积为什么赢"那节的结论吗？**FlashAttention 是同一个剧本的续集：第一次是公式迁就硬件，这一次是实现迁就硬件。**数学从头到尾没动过。

一张表收束

公式环节	优化技术	救的是什么
$QK^\top$ 里历史 K/V 的重复计算	KV cache	计算：每步 O(n²·d) → O(n·d)
cache 里 K/V 的份数（头数）	GQA / MQA	显存：K/V 头 32 → 8 → 1
cache 里 K/V 的存储形态	MLA	显存：K/V 压成低秩潜向量
$\text{softmax}(QK^\top)$ 的 n×n 中间矩阵	FlashAttention	带宽：tiling + online softmax，不落显存
softmax 输入的方差	√d 缩放	梯度：防饱和，让注意力学得动

一条 2017 年定稿的公式，本身一个字没改，养活了此后近十年的优化工业。理解了每个优化钉在公式的哪个位置，你再看新论文的"XX Attention"，第一个问题就该是：它动的是哪个环节？

不救了，直接换

上面这些全是"救"。但 2025 下半年起，出现了第一批不救、直接绕开手术台的方案：DeepSeek V3.2 的 DSA 稀疏注意力用一个轻量 indexer 先挑 token 再算注意力——在"和谁算"上直接做减法，API 价格借此再砍一半；Qwen3.5 干脆把约四分之三的层换成 Gated DeltaNet 线性注意力。但想看懂它们换掉的是什么，你还是得先把这条公式亲手推一遍——这正是本章存在的意义。

试一试（包括把它弄坏）

git clone https://github.com/KerroKapple/llm-anatomy && cd llm-anatomy
uv run python ch02/attention.py
uv run python ch02/scale_demo.py
uv run python ch02/kv_cache_demo.py

跑通只是及格，把它弄坏才算学会：

拆掉保命符：把 scale_demo.py 里的 scaled = raw / np.sqrt(d) 改成 scaled = raw，或直接在 scaled_dot_product_attention 里删掉 / np.sqrt(d)，看 d=1024 时熵怎么塌、maxP 怎么冲上 0.9+。再想一层：如果初始化时把 W_q 的方差调小，能不能部分挽救？（这其实就是一些真实模型 init 技巧的出发点。）
打开天眼：把 attention.py 自检里的 mask=causal_mask(n) 换成 mask=None，看权重矩阵上三角出现非零值——这就是"未来泄漏"。训练语言模型时犯这个错，loss 会低得离谱（照着答案抄当然低），一到真实生成就现原形。
给 cache 装滑窗：改 kv_cache_demo.py，让 cache 只保留最近 256 个 token（K[:t] 换成 K[max(0, t-256):t]），对比生成输出从哪个位置开始和全量 cache 分道扬镳、每 token 耗时曲线怎么从线性增长变成一条平线。恭喜，你亲手做出了滑动窗口注意力的雏形——顺着这条路走下去就是 Mistral 的滑窗和 StreamingLLM。
数一数自己的显存：照着 7B 投影表的算法，算一算你最常用的那个开源模型（层数、KV 头数、head_dim 翻它的 config.json）在你的显卡上最多能撑多长的上下文。算完你就明白服务商的"长上下文"贵在哪了。

下一章预告

打分、缩放、mask、cache——注意力这台发动机我们已经完全拆开了。但发动机不等于整车：残差连接、LayerNorm、FFN、位置编码，一个都还没登场。

第 3 章把它们全部焊上：手写一个完整的 Transformer block（PyTorch 从下一章登场，第 0 章约定过的交棒点），然后逐个"换件"——把 LayerNorm 换成 RMSNorm、把绝对位置编码换成 RoPE、把 FFN 换成 SwiGLU——换完你会发现，手里这台车已经叫 LLaMA 了。

仓库在 github.com/KerroKapple/llm-anatomy，每章一个目录，所有实验数字可复现。有问题欢迎开 issue 讨论。