搞大模型推理和训练的，迟早会在代码里撞见两样东西：一个是多头注意力机制（Multi-Head Attention），一个是 KV cache。我在做推理性能优化那段时间，几乎每天都要跟这两个概念打交道。刚开始的时候，我以为多头注意力就是多算几个注意力头，KV cache只是一个简单的缓存；真正上手之后才发现，如果只停留在“知道名字”的程度，后面看模型切开、显存估算、推理加速都会很吃力。这篇文章就把这两个东西一次讲透，从原理到工程实现再到常见坑，适合正在学大模型原理的开发者，也适合已经跑过模型、但还没系统梳理过这些细节的从业者。
1. 多头注意力机制：从“一个注意力”到“多个注意力”
1.1 先搞清楚自注意力到底在做什么
多头注意力不是凭空出来的，它建立在自注意力（Self-Attention）之上。自注意力的目标很简单：让序列里的每一个 token，根据自己和序列里其他所有 token 的相似程度，去汇总信息。这个“汇总”不是平均，而是有侧重点的信息抽取。
具体来说，每个 token 会生成三个向量：Query、Key、Value，也就是常说的 Q、K、V。这三个向量来自同一个输入向量乘以三组不同的权重矩阵。你可以用一个小类比来记：Query 是你想搜的内容关键词，Key 是图书馆里每本书的标签，Value 是书本身的内容。先拿 Query 去匹配所有 Key，算出相似度，再用相似度作为权重去加权求和对应的 Value。相似度高的书，贡献的内容就多；相似度低的，基本不影响最终结果。
在自注意力里，Q、K、V 都来自同一个输入序列，所以叫“自”注意力。它的作用就是让模型在处理某一个 token 时，可以自由决定“该多看哪个位置的 token、该少看哪个位置”。比如处理“猫在垫子上睡觉”这句话，模型在处理“猫”的时候，可以学会把较高的注意力权重分配给“睡觉”和“垫子”，而不是机械地只看前一个词。
这个机制在机器翻译、文本生成、分类等任务里被验证过很多次，核心能力就是建模长距离依赖。传统循环神经网络要一步步把信息传过去，距离一远信息就衰减；自注意力则是直接计算任意两个 token 之间的关系，路径长度恒为 1，从结构上解决了长距离依赖问题。理解了这一点，再看多头注意力就会顺很多。
1.2 拆成多头之后，模型到底得到了什么
一个注意力头的问题在于，它只能在一种“关系空间”里计算 token 的重要度。可语言里的关系非常复杂：有相邻词的修饰关系，有跨越多句的指代关系，有语法上的主谓关系，还有语义上的转折关系。如果只用一个注意力头，模型就得在同一个打分空间里把这些不同模式全部压进去，表达力很容易不够。
多头注意力的做法，是把原来的高维空间拆成多个低维子空间，每个子空间独立做一次注意力计算。假设模型的隐藏维度是 d_model，注意力头数是 h，那么通常的做法是先把 Q、K、V 从 d_model 维映射到 d_model 维，然后把它们切成 h 块，每块维度是 d_k = d_model / h。每个头在自己的 d_k 维子空间里做缩放点积注意力，得到自己的输出向量，最后把所有头的输出拼接回 d_model 维，再经过一个输出矩阵。
这样做有三个直接好处。第一，表达能力更强。每个头相当于一套独立的参数，模型可以学习到不同侧重的注意力模式，有些头抓近距离关系，有些头抓远距离关系，整体表达能力比单头强很多。第二，计算效率高。h 个头在数学上是并行计算的，实际代码里通常通过一次矩阵乘法和 reshape 完成，不会真的串行算 h 遍。第三，数值上更稳定。单个头的维度 d_k 往往比 d_model 小很多，比如 d_model=4096、h=32 时，每个头只有 128 维，在低维空间做 softmax 不容易出现过平滑，梯度往回传也更健康。
从直觉上理解，多头有点像让多个不同专长的分析师同时看同一份材料，然后综合他们的判断。每个分析师各有偏好的关注点，有的人盯前后邻词，有的人盯句法成分，有的人盯指代关系。最终模型拿到的不是某个单一视角的判断，而是多个视角的综合结果。
1.3 训练好的多头，观察下来各学各的
很多人在读完官方文档后都会问一个问题：这 h 个头，到底分别学什么？能不能人为指定某几个头去抓语法、某几个头去抓语义？
先说结论：不能指定，但训练完成后确实会出现分化。我在实际跑实验时观察过 attention map 的可视化结果，有些头的注意力权重几乎集中在相邻几个 token 上，学的是局部平滑关系；有些头则会关注句子里间隔很远的某个特定成分，像是在捕捉长距离指代；还有一些头的分布相对均匀，可能承担的是信息兜底的角色。这种现象在很多论文里也被反复验证过，说明多头机制确实给了模型选择的空间，至于每个头最终选择什么模式，是由训练数据分布和损失函数一起决定的。
这里有个容易踩的认知误区：不要以为头数越多一定越好。头数翻倍意味着参数量增加、训练时间拉长，而且头数过多会导致每个子空间维度太小，单个头能表达的信息反而有限。我在实际调参中的经验是，头数一般和 d_model 保持比较常规的比例，比如 d_model=512 时配 8 个头，d_model=1024 时配 16 个头，直接沿用成熟配置通常比自定义极端配置更稳。
2. 多头注意力的工程实现细节
2.1 完整计算流程和维度变换
多头注意力的工程实现，本质上就是一堆 reshape 和矩阵乘法的组合。我用一个具体例子来说，假设输入序列长度为 S，每个 token 的向量维度是 d_model=512，头数 h=8，那么每个头的维度 d_k=64。
第一步，输入 X（形状是 [S, 512]）分别乘上三个权重矩阵 W_Q、W_K、W_V，得到 Q、K、V，形状都是 [S, 512]。第二步，把 Q、K、V 从 [S, 512] reshape 成 [S, 8, 64]，再转置成 [8, S, 64]，相当于把 512 维向量切成 8 份，每份 64 维。第三步，对每一个头独立计算缩放点积注意力。以第一个头为例，Q 是 [S, 64]，K 的转置是 [64, S]，矩阵乘得到 [S, S] 的相似度矩阵，然后除以 sqrt(64)=8，再做 softmax。第四步，用这个 [S, S] 的权重矩阵去加权 V，得到当前头的输出 [S, 64]。第五步，把 8 个头的输出在最后一维上拼起来，得到 [S, 512]，再乘上输出权重矩阵 W_O，得到最终输出。
写成伪代码大概是这样的：
# 输入 x: [S, d_model] Q = x @ W_Q # [S, d_model] K = x @ W_K # [S, d_model] V = x @ W_V # [S, d_model] # 拆头: [S, h, d_k] Q = Q.view(S, h, d_k).transpose(0, 1) # [h, S, d_k] K = K.view(S, h, d_k).transpose(0, 1) V = V.view(S, h, d_k).transpose(0, 1) # 缩放点积注意力 scores = Q @ K.transpose(-2, -1) / sqrt(d_k) # [h, S, S] weights = softmax(scores, dim=-1) output = weights @ V # [h, S, d_k] # 合并多头 output = output.transpose(0, 1).contiguous().view(S, d_model) final = output @ W_O # [S, d_model]这个流程看起来不复杂，但真正容易出错的是维度顺序。PyTorch 里如果用[B, S, H*D]存储，拆头时需要先 view 成[B, S, H, D]，再转置成[B, H, S, D]。如果顺序错了，比如先转置再 view，数据会完全错乱，而且这种 bug 特别难发现，因为模型还能跑，只是效果明显变差。
2.2 为什么 softmax 之前要除以根号 d_k
很多人注意到注意力公式里有个1 / sqrt(d_k)，但没细想过为什么。这个缩放不是随便加的，它有明确的数学意义。
当我们计算 Q 和 K 的点积时，如果 Q、K 的每个元素都是均值 0、方差 1 的分布，那么点积结果的方差大约是 d_k。当 d_k 比较大的时候，比如 64 甚至 128，点积结果的数值会比较大。大的数值进入 softmax 之后，会让 softmax 的输出分布变得非常尖锐，接近 one-hot，也就是大部分概率都集中在一个位置上。这有两个坏处：一是梯度会非常小，因为 softmax 饱和区的梯度接近 0；二是模型会失去“合理分配注意力”的能力，基本变成“只看最相似的那一个 token”。
除以 sqrt(d_k) 之后，点积结果的方差被拉回到 1 左右，softmax 的输入在一个比较温和的量级，梯度能正常回传，注意力分布也更平滑。在工程上，这一步对 FP16 混合精度训练特别重要。如果不做缩放，FP16 下点积结果很容易溢出到很大的数值，softmax 之后直接出现 NaN，那个排查经历真的不想再来第二次。
有人会问，为什么不除以 d_k 而是除以 sqrt(d_k)？因为我们的目标是让方差回到 1，而方差是线性增长的，标准差则按 sqrt 增长，所以用 sqrt(d_k) 才能正好抵消方差增长。
2.3 工程里的常见坑
第一个常见坑是 mask 加错位置。注意力计算里的 mask 分成 padding mask 和 causal mask 两种。padding mask 是把填充 token 的位置加上负无穷，让 softmax 之后权重变成 0；causal mask 是生成任务里禁止当前 token 看到未来 token。两者作用不同，但在多头注意力的代码里都要在 softmax 之前加到 scores 上。加的位置错了，比如加到 softmax 之后，效果就会完全失效。我见过有人把 mask 加到权重矩阵上而不是 scores 上，结果模型训练的时候直接崩溃。
第二个坑是 head 维度的切分顺序。前面提到过，拆头的时候先 view 再 transpose，顺序不能反。实际写代码时，最好用torch.nn.functional.scaled_dot_product_attention这类融合算子，它会帮你处理维度细节，而且底层有优化的 kernel，速度通常比自己手写的实现快好几倍。但注意融合算子的 mask 参数形式跟手写版本不一样，用的时候还是要查一下文档。
第三个坑是输出投影的输出要不要再接残差。Transformer 结构里，多头注意力的输出在进入残差连接之前，通常还要过一个 LayerNorm。如果改变了注意力的内部维度设计，比如调整了头数但忘记同步调整 W_O 的输出维度，维度就对不上，一跑就报错。这种问题报错还算友好，最怕的是维度能对上但语义错乱，排查起来更花时间。
3. KV cache：推理加速里最扎实的一课
3.1 自回归生成为什么慢
大模型做文本生成时，是逐 token 预测的。每生成一个新 token，它要把新 token 拼到已有序列后面，然后再次输入模型，预测下一个 token。如果没有 KV cache，那么每预测一个新 token，模型都要对整个序列从头算一遍注意力。
我举个例子。假设现在要生成第 100 个 token，没有缓存的情况下，模型要重新计算前 99 个 token 对应的 Q、K、V，再算一遍完整的注意力矩阵。生成第 101 个 token 时，又要重新计算前 100 个 token 的 Q、K、V。随着序列越来越长，重复计算的量越来越大，复杂度大致是 O(n^3)，这里的 n 是序列长度。这个效率在推理场景是完全不能接受的，因为用户每问一句话，模型可能生成几百上千个 token，每次都要把之前的内容重算一遍，GPU 再强也扛不住这种浪费。
KV cache 的思路很简洁：既然前面的 token 在生成过程中已经被算过 K 和 V，那就把它们缓存起来，下次生成新 token 的时候直接复用，不用再算。这样每生成一个新 token，真正要重新计算的只有当前 token 的 K 和 V，整体复杂度从 O(n^3) 降到了 O(n^2)，如果再配合一些并行优化，解码阶段的耗时可以从平方级变成近似线性。
3.2 KV cache 到底缓存了什么，怎么边算边用
KV cache 缓存的是每一层注意力计算里的 K 和 V，不是缓存整个隐藏状态。这一点特别重要，因为每一个 Transformer 解码层都有自己独立的 K 和 V，所以 KV cache 不是一个单独的大数组，而是每一层一份缓存。
在常见的生成流程里，推理过程被分成两个阶段：预填充阶段和解码阶段。预填充阶段是第一次处理用户的提示词，这时候输入序列比较长，需要一次性算出所有 token 的注意力，并且把每一层的 K、V 保存下来。解码阶段是逐 token 生成，每来一个新 token，只需要计算它的 Q、K、V，然后把新算出来的 K、V 追加到对应层的缓存里，之后拿 Q 和缓存里的全部 K、V 做注意力计算。
Q 为什么不用缓存？因为 Q 只负责“查询”，它只服务于当前这个 token，用来和历史 K 计算相似度。历史 token 的 Q 在后续生成中用不到，因为当前 token 不会再去和历史 token 的 Q 做计算。所以 KV cache 只存 K 和 V，不存 Q，这也是它名字叫 KV cache 而不是 QKV cache 的原因。
从数据结构来看，每一层的 KV cache 可以看成形状为[batch_size, num_heads, seq_len, head_dim]的张量，K 一份、V 一份。实际使用时，框架会按最大序列长度预分配内存，比如max_seq_len=2048，那么每个张量维度里 seq_len 就按 2048 预留，实际用多少就往里面填充多少。
3.3 一个实际的显存账本
KV cache 最大的代价是显存。很多人对模型参数显存有概念，但对 KV cache 的显存开销完全没有概念，直到部署时才发现显存爆了。
我来算一笔账。假设当前模型是一个 7B 规模的模型，具体参数为：隐藏层维度 d_model=4096，注意力头数 h=32，每个头的维度 head_dim=128，Transformer 层数 L=32。我们使用 FP16 推理，dtype 每个元素占 2 字节。再假设当前只服务一个用户，batch_size=1，最大生成长度 seq_len=2048。
先看单层的 KV cache 大小。每一层的 K 和 V 各一份，每份大小是seq_len * num_heads * head_dim * 2 bytes，也就是2048 * 32 * 128 * 2 = 16MB。K 和 V 两份就是 32MB。模型有 32 层，所以总 KV cache 是32MB * 32 = 1GB。注意这只是 batch_size=1 的情况。
如果在线服务要同时处理 8 个用户，batch_size=8，KV cache 就是 8GB。如果把 max_seq_len 开到 8192，batch_size=32，那 KV cache 就有8192 * 32 * 128 * 2 * 2 * 32 * 32，算出来大约是 64GB。这个数字已经超过了很多单卡显卡的全部显存。这也解释了为什么大模型在线服务经常出现“人数一多就 OOM”的现象，瓶颈往往不在模型权重，而在 KV cache。
对比一下模型权重本身。一个 7B 模型用 FP16 存储，权重显存大约是 14GB。当 batch 增大、序列变长时，KV cache 的显存是可以轻松超过模型权重的。理解了这笔账，再看后面那些 KV cache 优化技术，就会明白它们为什么被那么看重。
4. KV cache 的进阶优化：从分页到架构级改造
4.1 PagedAttention：像操作系统管内存一样管 KV
传统 KV cache 实现里，框架会为每条请求预留一整块连续显存，大小等于max_seq_len * num_layers * num_heads * head_dim * 2。问题在于，实际生成的长度通常远远达不到 max_seq_len，预留的显存大部分被浪费了。而且不同请求生成的 token 数差别很大，有的只生成几十个，有的生成上千个，连续内存分配方式会导致大量的碎片和内部浪费。
PagedAttention 的思路是借鉴操作系统的分页机制，把 KV cache 切成固定大小的块，每个块可以独立分配和释放。生成的时候按需申请新的块，用完的块可以释放给其他请求使用。这样既避免了长时间预留整块显存，又把碎片化程度降到最低。块大小一般由框架自己配置，比如每个块存 16 个 token 的 KV，具体值会影响内存利用率和分配粒度。
这个优化对实际服务的吞吐提升非常明显。我见过同一个模型在同样 batch 下，开启分页式 KV cache 之后显存占用降了接近一半，服务的并发数跟着涨了一截。现在的成熟推理框架基本都支持这项特性，部署时可以直接打开，收益几乎零成本。
4.2 GQA 和 MLA：在源头减少 KV cache
如果说 PagedAttention 是在“管理”KV cache，那 GQA（分组查询注意力）和 MLA（多头潜在注意力）则是从架构层面直接减少 KV cache 的总量。
标准多头注意力 MHA 里，每个查询头都有自己独立的 K、V 头。如果模型有 32 个查询头，那就要存 32 份 K、V。GQA 的做法是让多个查询头共享同一组 K、V 头，比如 32 个查询头只对应 4 个 KV 头，KV cache 容量直接降到原来的 1/8。实际效果上，GQA 在多数任务里的精度损失很小，但推理显存节省非常可观，所以现在很多新模型都默认使用 GQA。
MLA 更进一步，它不直接缓存 K、V，而是缓存一个低秩压缩后的潜在向量。在注意力计算时，从这个潜在向量实时解压出 K、V。这样 KV cache 存储的是压缩后的数据，显存开销比 GQA 还要低。代价是推理时需要多做一步解压计算，但通常在显存瓶颈更严重的场景里，这笔计算换来的显存收益非常值。
如果你正在做部署选型，我建议直接优先选带 GQA 或 MLA 的模型架构，这是解决 KV cache 显存问题最治本的手段。后端的分页、量化都只是尽量用好现有空间，架构层省掉一部分才是真正的减法。
4.3 做推理优化时怎么从 KV cache 入手
我自己的习惯是，拿到一个推理卡顿或者显存溢出的问题，不急着调一堆参数，先把 KV cache 的显存账目算清楚，再按下面的顺序排查和处理。
第一步，确认当前 KV cache 的配置。看 max_seq_len 设置、batch 大小、是否开启缓存复用，先算出理论显存。第二步，打开分页式 KV cache，观察显存变化。这一步几乎无害，如果框架支持，直接开。第三步，考虑 KV cache 量化。KV cache 和权重一样可以做 int8 甚至 int4，但量化后可能存在精度损失，需要在验证集或实际样例上评估一下。第四步，如果显存还紧张，考虑把部分 KV cache 临时 offload 到 CPU 侧，或者使用统一内存管理，让 GPU 只保留必要的数据。第五步，如果以上都不够，那就只能换模型架构，切换到 GQA 或 MLA。
这些步骤里我踩过最多的坑是顺序搞反。有人一上来就量化 KV cache，精度降了不少，但明明只是 max_seq_len 设置过大导致的浪费。先算账，再调整配置，很多时候根本不需要上量化，改一改 batch 和序列长度就能解决大部分问题。
5. 高频问题排查与验证速查
5.1 KV cache 显存爆了的处理顺序
显存爆掉是推理部署里最常见的问题。我自己遇到过几次之后，总结了一个比较固定的处理清单，按优先级排下来是这样：
第一步，调低 max_seq_len。很多场景实际用不到 8192 的长度，改成 4096 甚至 2048，KV cache 显存直接减半。这是成本最低、效果最明显的一步。第二步，限制并发数或者 batch 大小。KV cache 和 batch 成正比，减少并发能立刻缓解显存压力。第三步，打开分页式 KV cache，减少碎片浪费。第四步，做 KV cache int8 量化，显存再降一半，但要注意精度变化。第五步，考虑 CPU offload。把一些层的 KV cache 挪到 CPU 侧，GPU 侧只保留必要的数据，缺点是会增加通信开销。第六步，换模型。如果服务对延迟和显存都敏感，直接换用 GQA 或 MLA 架构的模型，从源头解决问题。
这个顺序的核心原则是：先做无损伤的调整，再做有代价的优化。不要第一步就上量化，因为量化是不可逆的精度折衷，而前两步基本对模型效果没有影响。
5.2 容易误会的三个知识点
关于 KV cache，我见过不少人在理解上出现偏差，这里集中说一下。
第一个误区是“KV cache 缓存的是每一层的隐藏状态”。不对，缓存的是每一层注意力机制里的 K 和 V，不是整个隐藏状态。隐藏状态是注意力输出经过残差、归一化、前馈网络之后的结果，它每轮都在变，不能被简单缓存。KV cache 的粒度是每一层一组 K、V，和隐藏状态完全不是一回事。
第二个误区是“KV cache 只在推理时需要，训练时也需要”。训练阶段确实不需要，因为训练的时候整个序列是完整可见的，所有 token 的 Q、K、V 可以一次算出来，不存在逐步生成的问题。KV cache 的核心动机来自自回归生成，也就是推理阶段逐 token 输出的过程。
第三个误区是“有了 KV cache，解码阶段计算量就变成零了”。不对，KV cache 把重复计算变成了复用，但每生成一个新 token，仍然要算当前 token 的 Q、K、V，要算一次完整的矩阵乘法和注意力加权，还要算后面的前馈网络。它省掉的只是历史 token 的重复计算，不是所有计算。解码阶段真正的瓶颈往往不是算力，而是显存带宽，因为每一步要读取越来越长的 KV 缓存。
5.3 自己写个最小验证
如果你读了上面的内容还觉得有点飘，我建议你花 10 分钟自己写一个带 KV cache 的最小验证。不需要用大模型，一个小规模的 numpy 实现就够了。
核心逻辑是这样的：先定义 Q、K、V 三个矩阵，模拟生成第一个 token，把它的 K、V 存下来；然后生成第二个 token，使用缓存 K 和缓存 V 进行注意力计算。对比两种做法：第一种是把所有 token 的 K、V 都重新算一遍再做注意力；第二种是只算当前 token 的 K、V，然后拼接到缓存上再做注意力。如果实现正确，两种做法输出的结果应该完全一致。
实际验证中最容易出错的地方是拼接维度的顺序。缓存的维度通常设为[heads, seq_len, head_dim]，新 token 计算出来的 K 是[heads, 1, head_dim]，拼接时要在 seq_len 这一维上做 cat，不是在 batch 维或者 head 维。我第一次写的时候在这里翻过车，拼完之后注意力结果怎么都对不上，后来打印了每一步的 shape 才定位到问题。做这个最小验证的过程中踩一遍坑，对 KV cache 的理解会比读十篇文章都深。
另外一个建议是顺便测一下耗时。先生成几百个 token，对比开缓存和不开缓存的耗时差异。实测下来差距会很夸张，这个直观感受会帮你理解为什么 KV cache 是现代推理系统不可或缺的一部分。
最后再说说我的体会。多头注意力和 KV cache 是两个看起来基础、实际上却决定了模型训练和推理能跑多远的概念。如果你手头正好有部署任务，建议先不要急着上各种花哨优化，把自己模型的 KV cache 大小完整算一遍，很多瓶颈就清楚了大半。遇到问题也别慌，从显存账本、缓存策略、架构选择三个方向逐层排查，基本能覆盖掉绝大多数实际场景里的坑。这套思路我用了很长时间，每次换模型、换框架，都能靠它快速定位方向。
