2. 自注意力机制回顾：QKV矩阵的生成与计算过程详解
聊KV缓存之前，咱们得先把自注意力机制这老伙计请出来。说实话，我见过不少同学一上来就啃KV缓存源码，结果卡在QKV怎么来的、怎么算的上面。嗯，这就像盖楼不打地基——迟早要塌。
今天咱们就把这块地基夯实了。我会结合自己踩过的坑，把QKV的生成和计算过程掰开揉碎了讲清楚。
2.1 从输入到QKV：三个矩阵的诞生
假设我们有一段文本，比如"我今天很开心"。首先，它会被切分成token，然后映射成向量。这一步大家应该都熟悉，我就不啰嗦了。
关键来了：每个token的向量，会分别乘以三个不同的权重矩阵——Wq、Wk、Wv。结果就是Query、Key、Value三个向量。
核心公式：
Q = X · Wq
K = X · Wk
V = X · Wv其中X是输入矩阵，形状为 [seq_len, d_model]。
我个人习惯把这三个矩阵想象成三个不同的"视角"：Query负责问问题，Key负责回答问题，Value负责提供具体内容。你想想看，这不就是检索的过程吗？
我的经验：刚开始做Transformer时，我总搞混Wq和Wk的维度。后来我记住一个口诀——"Q和K的最后一维必须相等，因为要做点积"。这个细节在后续的KV缓存中特别重要。
2.2 注意力分数的计算：Q与K的舞蹈
有了Q、K、V，接下来就是自注意力的重头戏——计算注意力分数。说白了，就是让每个token去"关注"其他所有token。
具体分三步走：
- 计算相似度：Q和K做点积，得到注意力分数矩阵。形状为 [seq_len, seq_len]。
- 缩放：除以 √d_k，防止点积结果过大导致softmax梯度消失。d_k是每个头的维度。
- Softmax归一化：把分数变成概率分布，和为1。
完整公式：
Attention(Q, K, V) = softmax(Q · K^T / √d_k) · V这里有个坑，我曾经踩过——mask操作。在解码器里，我们不能让当前token看到未来的token。所以要在softmax之前，把未来位置的分数设成负无穷。嗯，这个细节在KV缓存里会反复出现。
2.3 多头注意力：并行计算的智慧
单头注意力虽然能工作，但表达能力有限。于是就有了多头注意力——把Q、K、V切成h个头，每个头独立计算注意力，最后再拼起来。
我画了一张图，帮你理解这个过程：
你看，每个头其实就是在做同样的事情，只不过参数不同。这就像让多个专家同时分析同一个问题，最后汇总意见——效果往往比单个专家好。
2.4 计算复杂度分析：为什么需要KV缓存？
现在咱们来算一笔账。假设序列长度是n，每个头的维度是d_k，头的数量是h。那么一次自注意力计算的复杂度是：
| 操作 | 计算复杂度 | 说明 |
|---|---|---|
| QKV生成 | O(n · d_model²) | 矩阵乘法，与序列长度线性相关 |
| Q·K^T | O(n² · d_k) | 平方复杂度，序列越长越慢 |
| Softmax | O(n²) | 逐行操作 |
| 分数·V | O(n² · d_k) | 又是平方复杂度 |
看到了吗？Q·K^T这一步是O(n²)的。这就是为什么长序列推理时，自注意力会成为瓶颈。我曾在一次项目中处理2048长度的序列，单次推理就花了300ms——大部分时间都耗在了这里。
注意：在自回归生成中，每次生成一个新token，都要重新计算整个序列的注意力。这太浪费了！因为前面token的K和V其实没变过。这就是KV缓存要解决的问题——把算过的K和V存起来，避免重复计算。
2.5 一个具体的计算示例
光说理论太抽象，咱们来走一遍实际计算。假设有两个token："我"和"爱"。
第一步：生成QKV
# 假设d_model=4，每个token的向量如下
X = [[1, 0, 1, 0], # "我"
[0, 1, 0, 1]] # "爱"
# 随机初始化Wq, Wk, Wv（实际训练得到）
Wq = [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6], [0.7, 0.8]]
Wk = [[0.2, 0.1], [0.4, 0.3], [0.6, 0.5], [0.8, 0.7]]
Wv = [[0.3, 0.3], [0.5, 0.5], [0.7, 0.7], [0.9, 0.9]]
# 计算Q, K, V
Q = X · Wq # 形状 [2, 2]
K = X · Wk # 形状 [2, 2]
V = X · Wv # 形状 [2, 2]第二步：计算注意力分数
# Q · K^T
scores = Q · K^T # 形状 [2, 2]
# 缩放
scores = scores / sqrt(2) # d_k=2
# Softmax
attention_weights = softmax(scores, dim=-1)第三步：加权求和
output = attention_weights · V # 形状 [2, 2]你看，整个过程就是矩阵乘法加softmax。但正是这个简单的过程，支撑起了整个Transformer的辉煌。
我的建议：如果你刚开始接触自注意力，强烈建议手算一遍这个小例子。我当年就是这么干的——虽然慢，但算完之后，整个机制就刻在脑子里了。
2.6 小结：为KV缓存铺路
回顾一下今天的重点：
- QKV是通过输入乘以三个权重矩阵生成的
- 注意力分数 = softmax(Q·K^T / √d_k) · V
- 多头注意力把QKV切成h份，并行计算再合并
- 自注意力的计算复杂度是O(n²)，这是KV缓存存在的根本原因
说白了，KV缓存就是利用了自注意力计算中的一个"漏洞"——在自回归生成时，K和V其实是可以复用的。下一章咱们就深入这个"漏洞"，看看怎么把重复计算省掉，把显存省下来。
嗯，今天就到这儿。记住我说的——手算一遍小例子，比看十遍文章都管用。
公众号:蓝海资料掘金营，微信deep3321