资讯动态

MHA到MLA:大模型注意力机制KV Cache显存优化全解析

发布时间:2026/10/20 23:24:09 来源:尧图企业网站定制
面试复盘做了几轮之后我发现一个挺普遍的现象大家准备“大模型面试八股”时都会背“MHA、MQA、GQA、MLA”这一串缩写但真被面试官追着问“为什么会有这些变体”“它们的KV Cache分别是什么量级”“MLA凭什么比MHA省那么多”能答到点子上的不到两成。问题其实不在于背没背过而在于大多数人只记住了结论没跟上这一系列方案背后真正要解决的那道题——推理阶段KV Cache的显存瓶颈。这篇文章就把这条线从头到尾捋一遍从Transformer最基本的自注意力计算开始把MHA的账单算清楚再依次拆解MQA、GQA、MLA的动机、原理、参数与缓存变化最后给出四种机制的横向对比和面试里常见的追问套路。不管你是正在准备面试的算法工程师还是做推理服务时被显存卡到头疼的开发这篇都能给你一套可以“抄作业”的理解框架。1. 从一次面试追问说起注意力机制的进化主线1.1 那个让大多数人卡壳的追问我模拟面试时经常从一个很基础的题目切入“你知道MHA和MQA的区别吗”多数人能答上来“MHA是每个注意力头都有自己的K、V投影MQA是所有Q头共享同一份K、V。”这个回答本身没错但只能算及格线。面试官听到这里通常会继续追一句“那GQA呢它为什么能兼顾质量和速度”再往深处问一句“你知道它们在KV Cache上的差异具体是多少倍吗”这时候能稳住的人就明显变少了。真正值钱的不是那个缩写而是缩写背后那条演化逻辑。MHA是源头MQA是“每个Q头共享一组KV”的激进压缩GQA是“分组共享”的折中MLA干脆换了思路不碰头数而是把KV压进一个低维潜在空间再用上投影恢复出来。把这几个方案串起来你会发现它们都在回答同一个问题怎么在不明显掉点的情况下把推理时KV Cache的显存负担降下来。1.2 变体出现的真正原因KV Cache是推理阶段的显存大头先问个更基础的问题为什么推理阶段会突然冒出来“KV Cache”这种东西Transformer解码是自回归的生成第T个token时需要让第T个token的Query去和前面所有token的Key、Value做注意力。Key和Value是前面已经算好的如果你不缓存每次生成一个新token都要把前面所有历史token的Q、K、V重新算一遍。那复杂度就不是线性增长而是平方增长这也是很多初代Transformer推理速度慢的直接原因。所以工程上一定会选择缓存每算完一个token的K和V就把它存到显存里下次直接读。问题来了这个缓存不但存在而且很大。它的大小是2 × 层数 × batch × 序列长度 × 隐藏维度和模型权重一样也几乎和模型规模成正比。实际推理时当序列上到几千、上下文窗口开到几万KV Cache占的显存很轻松就能超过模型权重本身。这是个硬瓶颈单纯堆显卡不是解决办法。理解了这一点后面所有的方案就都说得通了MQA在K、V头上做文章GQA在分组上做文章MLA在表示维度上做文章——本质全是冲着“一颗token的KV缓存到底要占多少个浮点数”去的。1.3 一条主线从“换头”到“换KV的存法”我把这条线概括成一句话MHA是“每头一套KV”MQA是“所有头一套KV”GQA是“每几个头一套KV”MLA是“用低维向量来存KV”。前三个方案的操作对象都是注意力头只是头之间共享程度不同所以它们的KV Cache和头数是线性挂钩的。MLA则跳出了这个框架它不谈你有多少个KV头而是说——先把K、V的信息一起压缩进一个低维向量等要算注意力的时候再上投影还原。这个概念转变很关键面试时如果能把这一层讲出来基本上就已经超过绝大多数候选人了。下面先把地基打牢不然后面谈对比都是飘的。2. 先把地基夯实自注意力完整计算与参数推导2.1 一次注意力计算的完整流程假设输入是一个长度为S的序列每个token的向量维度是h那么输入可以写成形状[S, h]的矩阵。自注意力计算四步走分别乘上三个权重矩阵得到Query、Key、Value Q XW_Q, K XW_K, V XW_V算注意力分数矩阵QK^T再除以根号下每个头的维度过softmax把分数变成概率分布用这个分布去加权求和Value得到当前token关注完整上下文后的输出最后通常还要过一个输出投影矩阵W_O把注意力头拼接后的结果再映射回h维。拿生活里的场景类比Q像你脑子里“想找什么”K像图书馆每本书的索引标签V像书的正文内容。你拿着Q去和所有K匹配匹配度高的书权重就大最后把多本“正文”按权重混合起来这就是注意力的输出。2.2 为什么除以根号dk这是面试里高频出现的一个点也是很多刚上手的人容易忽略的细节。Q和K的维度是dk假设Q、K向量里的元素都是近似独立、均值为0、方差为1的分布那么它们做点积之后结果的均值还是0但方差会变成dk。也就是说维度越大点积结果波动越大。如果dk很大点积值很容易跑到很大或者很小的绝对值区间softmax在这些区域会进入饱和区梯度会变得非常小训练基本推不动。除以√dk之后点积的方差被拉回1量级softmax能保持在一个对梯度友好的区间。这个操作不需要额外学习参数纯属数值稳定性设计。2.3 Q/K/V投影的参数量到底是多少很多文章直接说“自注意力的参数量约等于4h²”这个结论经常被拿来背但推导过程值得自己走一遍。如果只算单个注意力头它的Q、K、V投影矩阵分别是W_Q、W_K、W_V每个都是h × dk。但实际实现里通常有n_heads个头所有头的Q投影合并成一个大矩阵后形状是h × (n_heads × dk) h × hK和V同理。最后再加一个W_O也是h × h。所以一共四个h × h矩阵参数量就是4h²。总有朋友问“既然每个头的维度是dk为什么不直接写成4h × dk”如果你逐头去写每个头确实是3h × dk h × dkO投影但N个头拼起来就回到了4h × (n_heads × dk) 4h²。这俩是一回事只是写法粒度不同。面试时能讲清楚“合并视角”和“逐头视角”反而显得你理解得扎实。到这里基础计算就齐了下面可以安心聊MHA的账单了。3. MHA标准多头注意力的计算账单与KV Cache瓶颈3.1 多头切分与表示子空间MHAMulti-Head Attention的做法是把h维的Q、K、V分别切成n_heads份每份dk h / n_heads维每份单独计算注意力最后拼接。为什么多头有效一个非常直白的解释是不同的头可以关注不同类型的模式有的头关注词的语法位置有的头关注指代关系有的头关注语义相似度。如果只用单头所有信息都挤在一组Q、K、V里互相干扰多头等于给模型开了多条并行通道每条通道在各自的子空间里做检索最终再融合。在MHA里N个查询头对应N个Key头、N个Value头互不共享。这个结构表达能力最强但代价就是KV Cache也最大。它是我们后面所有对比的基准。3.2 KV Cache的计算公式与爆炸症结推理过程中KV Cache的占用公式可以写成M 2 × L × B × S × h × bL模型层数Bbatch sizeS当前序列长度h隐藏维度b每个参数占用的字节数前面的2是Key和Value各一份这里有个似乎理所当然但值得跟面试官确认的点为什么不缓存Query因为Query是随每个新token实时产生的算完就用掉了下次不会再被拿来检索而Key和Value是“历史信息”每个后续token都要反复查询它们。所以真正吃显存的是K和V。3.3 动手算一笔账1.5GB的缓存是怎么来的套一个典型的中型模型配置h 4096n_heads 32dk 128L 48层S 2048B 1用FP16存储也就是每个参数2字节。代入公式2 × 48 × 1 × 2048 × 4096 × 2算一下2 × 48 9696 × 2048 196608196608 × 4096 805306368再乘2字节约等于1.5GiB也就是约1.6GB。这只是一个请求。如果把序列长度拉到8192直接翻4倍变成约6GB如果开并发batch变成8那就是48GB。对比一下这时候模型权重才几个GBKV Cache反而成了显存里的巨无霸。这句话是面试里很提气的一句结论长上下文和并发场景下显存瓶颈往往不是模型权重而是KV Cache。有了这个认知你自然就能理解为什么后面每一个方案都盯着KV Cache做减法了。4. MQA单一KV头的激进压缩4.1 一个KV头服务所有查询头的设计MQAMulti-Query Attention是2019年一篇关于加速Transformer解码的论文提出的方案。它的核心动作只有一句话所有Query头共享同一个Key头、同一个Value头。也就是说在MHA里你有32个Key头MQA里只剩1个Key头和1个Value头。每个Query头在计算注意力时用的都是这同一份K、V。这是个非常激进的压缩策略因为KV的容量从“每头一套”变成了“全局一套”。4.2 参数与缓存的巨幅收缩从参数上看MHA的K、V投影是两个h × h矩阵MQA的K、V投影是两个h × dk矩阵。还是用上面的配置算h 4096h × h是1677万个参数h × dk只有52万个参数K和V合计一下光投影参数就从3355万降到105万少了97%左右。KV Cache同样夸张从约1.6GB降成2 × 48 × 1 × 2048 × 128 × 2约48MB这也就是MHA的三十二分之一。速度提升非常明显。但注意KV Cache缩小只是省显存实际计算时Key、Value还是会被广播到所有Q头去用所以FLOPs并没有减少那么多MQA的真正卖点是显存占用锐减。4.3 质量下降与训练稳定性问题MQA最大的代价是质量损失。直观理解K、V承载着句子的关键信息如果所有查询头都用同一个Key头、同一个Value头等价于强迫模型把“所有位置的语义信息”都塞进一个头里容量一下就满了。实验和工程实践都表明MQA在部分自然语言理解任务上会比MHA掉点尤其是在训练稳定性上表现不如MHA。后来也确实没有成为绝对主流但它最大的价值是把“共享KV头”这条路给趟开了后续GQA的出现几乎就是对MQA“太激进”的修正。面试里提到MQA时如果能补一句“它证明了压缩KV是可行的但单头容量受限”会显得你不是在死记硬背。5. GQA分组共享的折中方案与工程实践5.1 从“完全共享”到“分组共享”GQAGrouped Query Attention在2023年提出本质上是MHA和MQA之间的插值。它把n_heads个Query头分成g组每组内所有Query头共享一个Key头和一个Value头。当g n_heads时GQA退化成MHA当g 1时GQA退化成MQA。绝大多数情况下g取中间值比如8或16。这个“可调节”的属性很讨巧你想多省显存就把组数调小你想保质量就把组数调大。KV Cache的公式就变成了M 2 × L × B × S × (n_heads / g) × dk × b沿用前面的例子n_heads 32dk 128如果g 8相当于只有32 / 8 4个KV头KV Cache从1.6GB降为约0.4GB省了4倍。如果g 2那就是16个KV头Cache降为0.8GB左右省2倍。5.2 选组数g的学问这里面试官特别喜欢追问“g到底取多少合适”工程经验是g最好能整除n_heads。8组在中等规模模型上很常见质量接近MHA显存又能省不少如果模型很大、显存压力极大可以选4组如果任务对质量极敏感、又不缺显存16组也不是不行。关键是要理解这个trade-off组数越大容量越接近MHA但省显存越少组数越小越接近MQA显存省了但容量不足。5.3 代码级实现细节repeat还是reshape实现GQA时有一点需要留意虽然KV头数变少了但每个Query头还是要能找到自己对应的Key和Value。常见的做法是先把KV输出resize成[batch, kv_heads, seq, dk]然后通过repeat或者expand复制到[batch, n_heads, seq, dk]再走标准多头注意力。这里有一个坑如果直接调用repeat会真的在内存里复制一份完整数据显存反而没省只是省了KV Cache不是省中间计算显存。更推荐用expand或者“索引广播”的方式让多个Q头共享同一份物理内存。很多深度学习库已经封装好了这个逻辑但你自己手写时容易踩雷。面试时能把这一点讲出来说明你真的手写过GQA而不是只在论文里见过。GQA如今已经成为很多大模型的事实标准配置某70B级开源模型和后续的若干开源系列都采用了GQA。它用很小甚至可忽略的质量代价换来了实打实的显存下降属于性价比极高的工程选择。6. MLA把人引向潜在空间的更优解6.1 先看MLA要解决的三个问题MQA省得多但掉点GQA折中但还是在“头数”这个框架里打转。2024年初国内某开源大模型V2的技术报告提出了MLAMulti-head Latent Attention直接改变了游戏规则。MLA要解决的问题可以列三个进一步压缩KV Cache让它不再随“KV头数量”线性增长。避免MQA那样因为单KV头容量不够带来的质量损失。要能兼容旋转位置编码RoPE否则落地不了长文本场景。你可能会想GQA不是已经在工程上够用了吗为什么还要搞MLA因为GQA的KV Cache下界是“至少1个KV头”无论怎么压都跟头的数量挂钩。MLA干脆绕开这个量纲——我不跟你数头了我用低维向量存KV的“压缩态”。6.2 潜在向量压缩与上下投影MLA的核心结构可以分成两步看。第一步对当前token的隐藏状态h_t过一个下投影矩阵W_DKV压成一个低维潜在向量c_t W_DKV · h_t这个c_t的维度dc通常只有几百维比如512但它把K、V需要的信息都打包进去了。第二步计算注意力时再把这个潜在向量分别上投影回K空间和V空间k_t W_UK · c_tv_t W_UV · c_t这里的W_UK和W_UV都是上投影矩阵分别负责从低维潜在空间中恢复出完整的Key、Value。因为恢复过程是全连接投影模型可以学到把冗余信息压缩进潜在向量里再重建。它和MQA/GQA最大的区别在于MQA/GQA是“多个Q头共享有限的KV头”信息在头维度上被硬性切分MLA是在连续的低维空间里做信息压缩信息并没有简单丢弃而是被“编码再解码”。所以理论上MLA可以用更少的缓存容纳更多的信息量。6.3 最难处理的点RoPE与压缩空间的兼容真正让MLA实现变复杂的是旋转位置编码RoPE。RoPE是根据token位置生成旋转角、直接作用在Q和K向量上的。它要求K在计算注意力前就带好位置信息。但MLA的K是从低维潜在向量c_t上投影出来的而c_t本身被设计成所有位置共享一个编码空间。你没办法在c_t这个低维空间里直接做RoPE因为同一份c_t要被不同位置的token复用。MLA的解法是单独加一个带RoPE的Key分支。也就是说除了用c_t重建出主K外模型再另外维护一个k_t^R它是由一个专门的小投影矩阵从c_t或不带压缩的某分量映射出来、然后施加RoPE的。这样实际注意力计算时既要拼接主K和位置相关的那部分Key也要和Q同时带RoPE的部分计算匹配。所以MLA在推理时缓存的不是一个压缩向量而是潜在向量c_t加上带RoPE的k_t^R。多缓存一小段位置编码分支是为了能继续用RoPE这个细节如果面试时能说出来含金量很高。6.4 MLA的KV Cache到底省了多少每token每层的KV CacheMHA是2 × n_heads × dk也就是2h维。而MLA只需要缓存dc dk_R维。继续沿用前面的例子h 4096n_heads 32dk 128再看MLA如果取dc 512、dk_R 64MHA每token每层缓存2 × 32 × 128 8192维MLA每token每层缓存512 64 576维大约缩减到MHA的十四分之一。在某些KV头配置更大的模型里缩减比例能到几十倍。而且最关键的是MLA省缓存靠的是“潜在维度”而不是“砍头数”所以它在质量上的损失明显小于MQA甚至该模型的技术报告里提到在相同计算预算下效果还能小幅优于传统密集模型。代价是什么是解码阶段多了一次上投影计算也就是从c_t恢复K、V时多出的矩阵乘。但在推理场景里显存往往比算力更稀缺多花一点算力去换显存这笔账很划算。这也解释了为什么MLA一出来很多推理框架都抢着做专门的内核优化。7. 四者横向对比与面试官视角的考核要点7.1 一张表看清四种机制的差异我把四种机制放在一起对比用一句话概括各自的精髓机制核心做法每token每层KV Cache代表场景主要特点MHA每个Q头配独立的K、V2×n_heads×dk标准Transformer及训练初期表达力最强Cache最大MQA所有Q头共享一套K、V2×dk早期的解码加速研究最省但容量受限GQA每组Q头共享一套K、V2×(n_heads/g)×dk多种大模型的标配质量和速度平衡好MLAK、V共同压入低维潜在向量dc dk_R面向长文本的极致压缩打破头数限制效果优这里有几个细节值得强调MHA的KV Cache和头数严格成正比n_heads越大表达能力越强但缓存压力越大。MQA把KV头压到1个Cache是MHA的1/n_heads代价是头部容量不足。GQA的KV头数是n_heads/g工程上用4到8个KV头很常见。MLA的缓存不依赖头数而是依赖潜在维度dc和RoPE分支维度dk_R。7.2 面试时的高频追问与高分回答框架面试官通常不会只让你背定义他们更想听推导和取舍。我把高频追问整理成几个方向问为什么KV Cache只缓K和V不缓Q回答框架Q是当前生成token的查询向量用完即弃不会被后续任何一个token再次需要K和V是所有历史token的“记忆”后续每个新token都要拿自己的Q去检索它们所以必须存下来。问GQA的组数g怎么选回答框架g要能整除n_heads实际工程常用8组左右。g太小接近MQA质量损失明显g太大接近MHA显存节省有限。需要根据上下文长度、batch和显存预算倒推。问MQA/GQA看起来只是共享了KV头训练时和MHA有什么不同回答框架训练逻辑基本一致GQA在计算注意力前要把KV广播复制到对应的Q头上训练时主要差异是KV投影参数变少前向计算因为广播略有额外开销但总体影响不大。真正收益在推理阶段因为KV Cache变小了。问MLA的KV Cache为什么是dc加dk_R回答框架dc是被压缩的潜在向量负责保存K、V的主要信息dk_R是为了兼容RoPE额外保留/计算的一段带位置编码的Key分量。V不需要单独缓存因为它可以从潜在向量上投影恢复。问为什么除以根号dk回答框架稳定softmax的数值区间避免点积方差随维度增大导致梯度饱和。如果面试官追问“那不用会怎样”就回答“训练可能不收敛或收敛极慢尤其在长序列上更明显”。7.3 我给准备面试者的三条实操建议第一条不要只背结论把KV Cache公式亲手推一遍。面试官让你算“batch4、seq4096、h5120、32层、fp16时MHA的KV Cache是多少”这类问题如果你能当场写出公式并算出量级通常会留下很好的印象。第二条准备一张“对比图”在脑子里。不用画得多精美但要能边画边讲MHA是N套KVMQA是一套KVGQA是分组共享MLA是低维压缩。讲图比背文字更有说服力。第三条结合一个具体场景谈选型。比如“我要部署一个上下文32K的大模型显存紧张我会选GQA还是MLA”我的建议是框架支持好MLA就用MLA省得最多如果只想最小化改动GQA是稳定可靠的选择。8. 写在最后从理论到实践的一点体感我自己的体会是能把这四种机制讲清楚的人通常不是靠背八股而是真的被显存卡过、真的去看过推理框架的日志、真的改过注意力实现代码。KV Cache这玩意儿纸上算出来再大你也不一定有感觉直到你部署一个长文本模型看着显存占用随序列长度一路飙升你才会对每一个能省几个字节的设计都心存敬意。所以给正在准备的读者一个额外建议看完这篇文章别急着合上去找一个开源代码库里GQA或MLA的实现哪怕只读关键几行看它是怎么管理头的、怎么处理广播的、怎么缓存潜在向量的。面试时一旦聊起来你能从“听说过”变成“我读过”这层差距在面试官看来是非常明显的。理论讲完最后还是落回那个最原始的出发点我们做这么多变体不是为了炫技而是想让模型在更长的上下文里跑得更快、更省。想明白这个初衷你就永远不会被某一串缩写框住。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑