首页 / 资讯中心 / 文章详情

大模型注意力机制详解:MHA、MQA、GQA、MLA原理与KV Cache优化

大模型注意力机制详解:MHA、MQA、GQA、MLA原理与KV Cache优化 ★ FEATURED ARTICLE
如果你在准备大模型方向的面试注意力机制这一段基本绕不开。MHA、MQA、GQA、MLA 这四个缩写几乎每轮都会出现至少一个而且面试官往往不是简单问你知不知道而是让你从原理讲到推理优化再追问几个为什么。我整理这篇文章本质上是把自己手写的面试笔记做了一次系统化梳理用大白话把四个机制讲透再补上那些容易被追问到、但资料里很少讲清楚的细节。适合两类人看一类是准备大模型算法岗面试的人另一类是平时做应用集成、想搞懂模型内部结构但没系统推过公式的同学。1. 先回到原点MHA 多头注意力到底在解决什么问题不管后面怎么变MHAMulti-Head Attention是所有变体的基线。面试的时候如果直接跳过 MHA 讲 MQA、GQA、MLA面试官大概率会觉得你基础不牢。所以先花时间把 MHA 讲清楚后面那些变体也都是围绕着它的痛点做文章的。1.1 缩放点积注意力的数学本质核心公式就一个Attention(Q, K, V) softmax(Q K^T / sqrt(d_k)) V这里 Q、K、V 都是从输入向量经过线性投影得到的。你可以把 Q 理解成“我在找什么”K 是“我身上有什么标签”V 是“我真正提供的内容”。注意力分数就是让每个 Q 去和所有 K 做点积看看匹配程度有多高然后通过 softmax 变成权重最后用权重去加权 V 的内容。这里有个很多人忽略但面试官很喜欢问的点为什么要除以 sqrt(d_k)因为当向量维度 d_k 变大时点积的数值方差也会变大。两个 d_k 维向量点积如果每个元素都是均值 0、方差 1 的分布那点积结果的方差大约是 d_k。方差太大softmax 输入值会落在梯度很小的区域反向传播时梯度会消失训练就很难继续。除以 sqrt(d_k) 之后方差被拉回到 1 左右softmax 的梯度就能保持稳定。这个细节看似简单却是回答“为什么注意力公式长这样”的必备答案。从计算角度看序列长度为 L、隐藏维度为 D 时Q K^T 这一步的复杂度是 O(L^2 D)因为要算 L 个 query 和 L 个 key 之间两两点积每个点积要算 D 维。如果序列从 1K 涨到 128K这一步就是平方级增长这也是后面所有长文本优化都绕不开的问题。1.2 多头机制的价值与代价MHA 和单头注意力的核心区别在于把 D 维空间切成了 H 个 d 维子空间每个头独立做注意力最后拼起来再过一层输出投影。H 是注意力头数通常 d D / H。为什么要切头你可以这样理解如果只有一个头模型每个位置只能产生一种注意力分布相当于“只能从一个角度找信息”。多头相当于让模型同时从多个子空间去检索信息——一个头可能更关注语法关系另一个头可能更关注指代关系另一个头可能更关注距离上的依赖。最后这些头的结果拼接起来表达能力会比单头强很多。代价也很明显每一层要维护 H 组 Q、K、V 投影。假设隐藏维度 D那么 Q、K、V、O 四个投影矩阵的参数量就是 4D²。头数越多可以表达的关系越丰富但对应的训练和推理成本也越高。面试时如果被问“头数和维度的关系”记得强调 D 固定时切头不会增加 Q、K、V 的总计算量因为每个头的维度变小了加起来的计算量仍然是 D 级别真正增加的是投影矩阵带来的参数以及注意力分布计算的灵活度。1.3 MHA 的隐形瓶颈推理时的 KV cacheMHA 在训练阶段的问题还不那么致命但到了推理阶段就会暴露一个巨大的麻烦KV cache。自回归生成时模型每生成一个 token都要重新计算之前所有 token 的注意力。为了不重复计算工程上会把已经算过的 K 和 V 缓存下来这就是 KV cache。MHA 下每个头都有自己的 K 和 V所以每生成一个 token就要往缓存里追加 H 个 key 向量和 H 个 value 向量一共 2D 个浮点数。我们可以算一笔账一个隐藏维度 D4096、层数 L32 的模型生成序列长度 T2048用 fp16 存储。KV cache 大小大约是2 × 4096 × 2048 × 32 × 2 字节 ≈ 1GB这个数字非常惊人。随着序列长度进一步增长KV cache 会线性膨胀长上下文场景下甚至比模型权重还占显存。更麻烦的是推理时每次生成一个 token 都要把全部 KV cache 读一遍显存带宽被大量消耗生成速度也会被拉低。所以MQA、GQA、MLA 本质上都是同一个思路的产物把 KV cache 的存储和读取成本降下来。理解了这一点后面几个变体的设计逻辑就都能串起来了。2. MQA 与 GQA让 KV 头共享起来MQA 和 GQA 的出发点非常直接既然 KV cache 的膨胀是因为每个 query 头都配了一套独立的 K、V那把 KV 头共享出去不就行了2.1 MQA 的原理与数值收益MQAMulti-Query Attention的核心做法是所有 query 头共享一组 K 和 V。也就是说原来有 H 个 key 头和 H 个 value 头现在只剩下 1 个 key 头和 1 个 value 头所有 query 头都用这一组 K、V 去算注意力。这样做的好处非常直观KV cache 的大小从每 token 2D 降到了每 token 2d其中 d 是单个头的维度。如果 H32那 KV cache 直接降到原来的 1/32。还是上面那个例子D4096、L32、T2048MQA 的 KV cache 大约是 32MB和 MHA 的 1GB 相比是数量级的差异。但代价也很明显表达能力变弱了。所有 query 头被迫共享同一份 K、V相当于每个头只能在相同的 key 和 value 里做不同的加权组合这会限制注意力分布的多样性。实践中 MQA 在部分任务上会和 MHA 有一些效果差距尤其是对细节信息敏感的任务。这里要说一个面试高频坑MQA 的“计算量降低”主要发生在 KV cache 的读写和 KV 投影的计算上注意力分数矩阵的规模并不会直接变成原来的 1/H。因为每个 query 头仍然要和共享的 K 做一次点积H 个 query 头就是 H 次点积总的注意力计算量依然是 O(L²D) 级别。真正省下来的是推理时的显存占用和访存带宽。2.2 GQA 的分组平衡术GQAGrouped-Query Attention是 MQA 和 MHA 之间的折中。它的做法是把 query 头分成 G 组每组内部共享一份 K、V。G1 时就是 MQAGH 时就是 MHA所以 GQA 可以看成是一个带旋钮的泛化版本。KV cache 的大小这时候是每 token 2 × G × d。如果 H32、G8那 KV cache 就是 MHA 的 8/32也就是 1/4。上面那个例子里KV cache 大概是 256MB比 MHA 的 1GB 好很多又比 MQA 的 32MB 大一些但换来的是更接近 MHA 的效果。面试时如果被问“G 应该选多大”不要给出固定答案。G 越大越接近 MHA效果越好但 KV cache 也越大G 越小推理越快但效果可能下降。实际选择要看模型规模和部署需求。通常 7B 到 70B 这个区间的开源模型用 8 组或 4 组都比较常见更小的模型可能直接上 MQA 也不会太伤效果。核心原则是在显存和带宽受限时优先保证推理性能在效果敏感场景适当增加组数。2.3 共享 KV 之后的真实效果变化很多人会问一个很自然的问题共享 KV 到底会损失多少效果从实验和实际模型表现来看GQA 在大部分任务上能做到几乎不掉点尤其是模型规模增大之后效果差距会进一步缩小。MQA 在某些任务上会掉得更明显一点但也没有到“不可用”的程度。这里有个很有意思的观察共享 KV 相当于给模型加了一种归纳偏置——不同 query 头需要在同一组信息里寻找不同特征。这反而可能带来一些正则化效果减少过拟合。当然这更多是事后解释不能当作设计 GQA 的原始动机。从部署角度看共享 KV 带来的收益非常实在KV cache 变小batch size 就能开得更大显存带宽压力降低生成吞吐量会明显提升长文本场景下能支持的上下文长度也大大增加。这也是为什么现在主流开源大模型基本都在用 GQA 而不是 MHA大家都很清楚推理成本在实际落地时有多重要。3. MLA 低秩压缩把 KV cache 做成一个瓶颈MQA 和 GQA 的思路是“砍头”MLAMulti-head Latent Attention的思路则更进一步不直接砍头而是把 K、V 先压缩到一个低维空间要用的时候再展开。这就是它的核心创新点。3.1 从降维视角重新理解 KV cacheMHA 里每个 token 的 K、V 是 D 维向量但这么多维度里真的全是有效信息吗不一定。注意力机制真正需要的可能是这 D 维信息里一个低维子空间。既然如此为什么不直接学一个低维表示来存储MLA 的做法就是这样对每个 token先把输入通过一个降维矩阵压缩成一个维度为 c 的向量其中 c 远小于 D。这个 c 维向量就是实际缓存下来的 KV 表示。真正要用 K 和 V 的时候再通过升维矩阵把它展开回 d 维参与注意力计算。这样做的核心价值在于KV cache 不再和 D 成线性关系而是和压缩维度 c 成线性关系。如果把 D4096 压到 c512KV cache 直接变成原来的 1/8。而且这个压缩不是简单地扔掉信息而是通过训练学出来的低秩表示理论上能比 MQA/GQA 那种直接共享保留更多有效信息。3.2 核心公式与矩阵吸收技巧为了面试能讲清楚我按常见论文里的简化表述来拆。设某一层输入是 h_t维度为 D。MLA 先做低秩投影c_t^{KV} h_t W^{DKV}这里 W^{DKV} 的形状是 D × cc 就是压缩维度。得到的 c_t^{KV} 是低维的 KV 隐状态推理时缓存的也是这个 c_t^{KV}。需要计算注意力时再通过上投影展开k_t c_t^{KV} W^{UK} v_t c_t^{KV} W^{UV}同理query 也可以走一个类似的低秩投影得到 q_t。到这里你可能已经发现一个问题如果推理时每个 token 都要先把 c_t^{KV} 展开成完整的 k_t、v_t那计算量不是反而增加了吗这就要提到 MLA 的另一个关键优化——矩阵吸收。因为 q_t 和 k_t 都是通过矩阵乘法得到的所以计算注意力分数 s q_t^T k_j 时可以把 k 的上投影矩阵和 q 的投影矩阵先合并起来这样直接拿低维的 c_j 去算分数而不用先把 c_j 展开成完整 k_j 再做点积。value 那边同理v 的上投影可以往后和输出投影合并避免在注意力内部展开完整的 v。整体效果就是缓存的是低维向量计算时也不需要额外承担高维展开的代价。这个技巧在面试里是很大的加分项。你可以这样解释矩阵乘法满足结合律先合并矩阵再算和先展开再算数学结果是等价的但计算成本完全不同。现实中 MLA 的实现还会配套优化算子减少中间张量的显存占用。3.3 位置编码怎么办RoPE 解耦位置编码是绕不开的问题。很多现代模型都用 RoPE 旋转位置编码它要求把位置信息直接加在 q、k 向量上而且必须是 q、k 同时旋转相应角度才能体现相对位置关系。问题来了MLA 里的 k 是从低维 c 展开出来的如果先展开再施加 RoPE缓存的时候就必须把展开后的完整 k 存下来那压缩的意义就没了如果对低维 c 直接施加 RoPE维度对不上位置信息也会乱。MLA 的做法是把位置相关和语义相关拆开处理。具体来说把 query 和 key 分别拆成两部分一部分走低秩压缩、参与语义注意力计算另一部分单独接 RoPE、负责位置感知。这样位置编码部分只占很小维度不会显著增加 KV cache又能保留相对位置信息。这个设计在专业术语里常叫 decoupled RoPE。面试时能主动提到这一点说明你理解的不只是表面公式而是真正关心位置编码和压缩缓存之间的冲突。3.4 一个可复算的显存对比例子我们沿用前面的参数D4096、L32、T2048、fp16。假设 MLA 的压缩维度 c512那每层每个 token 的 KV cache 就是 2×512 个浮点数总大小大约是2 × 512 × 2048 × 32 × 2 字节 ≈ 128MB作为对比MHA 是 1GBGQAG8是 256MBMQA 是 32MB。MLA 处在 MQA 和 GQA 之间但实验效果通常比同等缓存预算下的 GQA 更接近 MHA甚至更强。这正是 MLA 的价值所在同样的显存预算能换来更好的效果。另外要补充一句MLA 的主要优势集中在推理阶段。训练阶段仍然需要在高维空间里计算完整的注意力升维和降维过程会额外增加一些计算量所以不能简单理解成“MLA 一定比 MHA 训练更快”。面试时把训练和推理分开讲会显得思路更清楚。4. MHA、MQA、GQA、MLA 横向对比与面试追问把四个机制放在一张表里看差异就非常清晰了。前面的公式拆解可能有点长面试时先用这张表抓框架再逐项展开是很高效的表达方式。4.1 一张表说清核心差异机制KV 头配置每 token KV cache相对 MHAK/V 投影参数相对 MHA训练实现复杂度推理访存压力相对效果MHAH 个独立 KV 头11标准高基线MQA1 个共享 KV 头1/H约 1/H简单很低略降GQAG 个共享 KV 头G/H约 G/H简单中低接近 MHAMLA低秩压缩瓶颈 c约 c/D取决于低秩维度较复杂低接近甚至优于 MHA这张表有个地方需要特别解释MQA 和 GQA 的 K/V 投影参数是在“标准实现”前提下才会按比例下降的。如果某个实现只是把原有的大矩阵广播到各个头那参数其实没变只是计算时复用了一份结果。面试时遇到这种实现细节可以直接指出“参数是否真的减少取决于投影矩阵的输出维度有没有跟着缩”这个细节很能体现工程敏感度。MLA 那一行的参数不能简单说变小因为它除了降维矩阵还多了升维矩阵。它的核心收益主要在 KV cache 和访存而不是参数量减少。把这几点说清楚面试官能明显感受到你是真的理解而不是背结论。4.2 面试官最爱追问的 8 个问题我整理了几个高频追问每个都可以用两三句话答到点子上为什么注意力公式要除以 sqrt(d_k) 控制点积数值方差避免 softmax 进入饱和区导致梯度消失。MHA 每个头的 Q、K、V 是怎么得到的 同一个输入 x 分别乘独立的 Q、K、V 投影矩阵每个头用自己那块参数。MQA 的计算量真的比 MHA 低吗 注意力分数计算量基本不变真正降低的是 KV 投影、KV cache 存储和访存开销。GQA 的组数怎么选 G1 就是 MQAGH 就是 MHA实际按效果和推理成本的平衡点选常用 4 或 8。MLA 和直接剪掉一部分维度有什么本质区别 剪枝是去掉已有结构MLA 是让模型自己学一个低秩子空间更灵活信息保留能力更强。RoPE 和低秩 KV 压缩会不会冲突 会所以 MLA 用 decoupled RoPE 把位置维度和语义维度拆开处理。为什么长文本尤其需要 GQA 或 MLA KV cache 随序列长度线性增长上下文越长压缩 KV cache 的收益越大。大模型训练时直接用 GQA 还是先训 MHA 再转换 主流做法是训练时就确定结构并保持一致因为训练和推理结构不一致会导致效果劣化。MLA 更是在训练阶段就要用完整实现确保推理时能正确吸收矩阵。这些问题没有标准答案但回答的思路很重要。面试官想听的不是你背下来的定义而是你能否从计算、显存、带宽、数值稳定性这些角度把因果关系讲清楚。4.3 什么场景选什么机制这不是让你做学术评审而是理解每种机制最适合的位置。如果团队追求极致的推理吞吐且任务对细节要求不算极高MQA 是一个可行的激进选择结构最简单KV cache 最小。如果既要部署成本低又要尽量保住效果GQA 是当前最主流的折中大部分开源模型的生成服务都验证了这个路线。如果要做极长上下文、且显存非常紧张MLA 的低秩压缩在相同缓存预算下通常比 GQA 效果更好代价是实现和训练复杂度偏高。MHA 依然有自己的位置——当模型训练资源充足、推理不是瓶颈、或者任务对注意力多样性要求很高时它仍然是效果上最稳的基线。很多研究实验仍然用 MHA 作为默认配置就是这个原因。5. 备考与实战中容易翻车的三个细节最后这部分是我自己踩过坑、也在模拟面试中看过别人反复摔跟头的地方。提前避开能让你在面试和实际工作中都少走弯路。5.1 “参数共享”不等于“计算量下降”这是我在面试中见到最多的错误说法。很多人会下意识觉得共享了 K、V那注意力计算量当然也变小了。实际上Q K^T 这一步在每个 query 头上都得做MQA 只是让每个头用同一个 K并不是说就少算了多少份点积。我建议这样记忆共享 KV 省的是 KV cache 的存储和读取以及 KV 投影的参数量和计算量省不了的是每个 query 头都必须算一遍的注意力分数。严格来说如果实现时不做优化MQA 的注意力计算量和 MHA 几乎是一样的。推理变快更多是缓存变小之后访存压力降低带来的间接收益。5.2 背公式不如推一遍手推复杂度的正确姿势面试时如果被要求推导复杂度不要直接背结论而是从矩阵形状开始推。把 Q 看成 L×DK^T 是 D×L乘积是 L×L那么 QK^T 这一步就有 L×L×D 次乘加。V 是 L×D加权后得到 L×D再乘输出投影 D×D又是 L×D²。最后把所有步骤加起来就是 O(L²D LD²)。这样推完你还能顺手回答另一个追问序列变长时L²D 会先成为瓶颈模型维度变大时LD² 会更占开销。理解了这个你就知道为什么长文本模型要在注意力机制本身上下功夫为什么 RNN 式的线性注意力会有人研究为什么最近大家都在做长上下文的推理优化。面试中能把复杂度从原理上讲清楚比干巴巴丢出一个大 O 符号有说服力得多。5.3 主动把话题引向长文本是面试加分项我个人的体会是这四个机制的面试本质上考的不只是记忆而是你是否理解大模型推理的结构性瓶颈。如果面试官问你 GQA 和 MLA 的区别最好的回答方式不是只背定义而是从 KV cache 随序列长度线性膨胀这个事实切入然后解释每种机制如何缓解这个问题最后落到“长上下文场景下KV cache 才是主要矛盾”这个结论上。还可以顺手提一下推理系统里除了 KV cache还有激活值显存、权重显存、中间结果、算子调度效率等注意力机制的变体解决的是其中一个环节。能够把视野拉到系统层面面试官会觉得你不只是个调包漏而是真的在思考模型从训练到部署的全链路问题。最后分享一个我常用的准备思路把 MHA 的公式推导一遍把 GQA 的 KV cache 数值手算一遍再用自己的话把 MLA 的矩阵吸收解释给别人听这个过程比单纯刷题库有用得多。真到了面试桌上你能用一张草稿纸把 Q K^T 的形状画清楚把 KV cache 的字节数算出来比任何漂亮话都更有说服力。
阅读完成 · 觉得有帮助?
咨询建站