1. 面试官到底想从Attention里问出什么面大模型岗位Attention几乎是绕不开的一道坎。尤其是最近两年面试官不再满足于让你背出“Scaled Dot-Product Attention”的公式而是会直接甩一句“手写一下MHA再讲讲MLA和GQA的区别。”很多人准备的时候把这几套机制当成四个孤立的知识点去背结果面试官稍微换个角度追问——比如“为什么MQA能省显存但效果会掉”“MLA的KV压缩到底压在哪一步”——就露馅了。我自己面过也面过别人发现一个规律面试官真正想考察的不是你会不会默写代码而是你脑子里有没有一张“显存-计算-效果”的三角权衡图。MHA、MQA、GQA、MLA这四者本质上是在这张三角图上取不同的点。你只要能说清楚每个点为什么这么取、代价是什么、工程上怎么落地代码反而是最不需要担心的部分。这篇内容我打算按面试实战的顺序来拆先把MHA的完整实现手撕一遍包括那些面试官爱抠的细节再顺着“KV Cache爆炸”这条线引出MQA和GQA然后重点讲MLA这个最近被问得越来越多的结构最后给一套面试现场手写代码的模板和常见追问的应对思路。适合正在准备大模型算法岗面试的同学也适合想把这几个机制真正搞明白的工程师。2. 从MHA手撕开始那些面试官爱抠的实现细节2.1 为什么面试官总让你先写MHAMHA是这一切的基准。面试官让你手写MHA不是想看你背公式而是想看你对张量维度变换的敏感度。我见过太多候选人公式背得滚瓜烂熟一到写代码就在reshape和transpose上卡壳维度对不上就开始瞎试。这种在面试官眼里就是“没真正写过”。先明确MHA的核心逻辑把d_model维的输入通过三组线性变换投影成Q、K、V然后切成h个头每个头独立做attention最后拼回来再过一个输出投影。听起来简单但魔鬼全在维度上。2.2 手撕MHA的完整代码与维度追踪下面这份实现我建议你面试前默写三遍直到不用想就能写出来import torch import torch.nn as nn import math class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout0.1): super().__init__() assert d_model % num_heads 0, d_model必须能被num_heads整除 self.d_model d_model self.num_heads num_heads self.d_k d_model // num_heads # 每个头的维度 self.W_q nn.Linear(d_model, d_model) self.W_k nn.Linear(d_model, d_model) self.W_v nn.Linear(d_model, d_model) self.W_o nn.Linear(d_model, d_model) self.dropout nn.Dropout(dropout) def forward(self, x, maskNone): batch_size, seq_len, _ x.shape # 1. 线性投影 Q self.W_q(x) # (B, L, d_model) K self.W_k(x) V self.W_v(x) # 2. 切多头: (B, L, d_model) - (B, h, L, d_k) Q Q.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) K K.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) V V.view(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2) # 3. 计算attention分数 scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) # scores: (B, h, L, L) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn_weights torch.softmax(scores, dim-1) attn_weights self.dropout(attn_weights) # 4. 加权求和 context torch.matmul(attn_weights, V) # (B, h, L, d_k) # 5. 拼回多头: (B, h, L, d_k) - (B, L, d_model) context context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model) # 6. 输出投影 output self.W_o(context) return output这段代码里有几个面试官必抠的点我一个个说。第一个坑view和transpose的顺序。很多人会写成先transpose再view结果维度全乱。记住view要求内存连续所以必须先view成(B, L, h, d_k)再transpose(1, 2)变成(B, h, L, d_k)。反过来做view会报错或者得到错误的结果。第二个坑contiguous()的位置。在拼回多头的时候transpose之后张量在内存里不连续直接view会出问题所以必须先.contiguous()。这个细节面试官特别爱问因为它是“你到底跑没跑过代码”的试金石。第三个坑缩放因子。为什么除以sqrt(d_k)而不是sqrt(d_model)因为每个头独立做attention点积的方差随d_k增长所以要按d_k缩放。这个在面试里被问到的概率超过50%。2.3 面试现场写MHA的三个实用技巧我自己的经验是面试手写代码时不要一上来就写完整版容易在细节上翻车。可以分三步走先写骨架把__init__里的四个线性层和forward的六个步骤用注释列出来让面试官看到你思路清晰。再填维度变换重点写view和transpose那几行边写边在注释里标维度比如# (B, L, h, d_k)。面试官看到你标维度会觉得你心里有数。最后补细节mask、dropout、缩放因子这些写完主体再补避免一开始就被细节带偏。提示如果面试官让你在白板上写维度注释一定要写清楚。白板代码没有IDE帮你检查维度标注是你唯一的“调试工具”。3. KV Cache这把双刃剑MQA和GQA为什么会出现3.1 推理阶段的显存账本MHA在训练阶段没问题但一到推理就暴露了。自回归生成时每生成一个token都要把之前所有token的K和V缓存下来这就是KV Cache。它的显存占用公式是KV Cache大小 2 × batch_size × num_heads × seq_len × d_k × 精度字节数拿一个具体例子算LLaMA-2 70Bnum_heads64d_k128batch_size1seq_len4096FP16精度。代入公式2 × 1 × 64 × 4096 × 128 × 2 bytes 128 MB单看一个请求还好但线上服务要同时处理几十上百个请求batch_size一上去KV Cache直接吃掉几十GB显存。这就是为什么推理时显存总是不够用的核心原因之一。3.2 MQA把K和V压到极致MQAMulti-Query Attention的思路非常直接所有头共享同一份K和V。Q还是h个头但K和V只有1个头。这样KV Cache直接缩小到原来的1/h。还是上面那个例子MQA下KV Cache变成2 × 1 × 1 × 4096 × 128 × 2 bytes 2 MB从128MB降到2MB这就是MQA的威力。但代价也很明显所有头共享K和V等于强行让不同头关注相同的位置信息表达能力下降训练时也容易不稳定。我实测过MQA在小模型上掉点不明显但模型越大、任务越复杂掉点越明显。3.3 GQAMQA和MHA的折中方案GQAGrouped-Query Attention是Google在LLaMA-2里推广开的方案。它的做法是把h个头分成g组每组共享一份K和V。当g1时退化成MQA当gh时退化成MHA。class GroupedQueryAttention(nn.Module): def __init__(self, d_model, num_heads, num_kv_heads, dropout0.1): super().__init__() self.num_heads num_heads self.num_kv_heads num_kv_heads self.num_groups num_heads // num_kv_heads self.d_k d_model // num_heads self.W_q nn.Linear(d_model, num_heads * self.d_k) self.W_k nn.Linear(d_model, num_kv_heads * self.d_k) self.W_v nn.Linear(d_model, num_kv_heads * self.d_k) self.W_o nn.Linear(num_heads * self.d_k, d_model) def forward(self, x, maskNone): B, L, _ x.shape Q self.W_q(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2) K self.W_k(x).view(B, L, self.num_kv_heads, self.d_k).transpose(1, 2) V self.W_v(x).view(B, L, self.num_kv_heads, self.d_k).transpose(1, 2) # 关键把K和V在头维度上复制num_groups次 K K.repeat_interleave(self.num_groups, dim1) V V.repeat_interleave(self.num_groups, dim1) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) context torch.matmul(attn, V) context context.transpose(1, 2).contiguous().view(B, L, -1) return self.W_o(context)GQA的精髓在repeat_interleave这一步。面试时如果被问到“GQA怎么实现”你一定要点出这个操作K和V在头维度上复制让每个Q头都能找到对应的K、V头。LLaMA-2 70B用的就是num_heads64, num_kv_heads8也就是8个Q头共享一组KVKV Cache降到原来的1/8效果又比MQA好很多。3.4 三者的取舍对照表机制KV头数KV Cache表达能力训练稳定性典型应用MHAh最大最强最稳训练阶段、小模型GQAg (1gh)中等较强较稳LLaMA-2/3、多数开源模型MQA1最小最弱较差推理加速、极端显存场景面试时如果被问“你选哪个”标准答案是训练用MHA或GQA推理部署优先GQA显存极度紧张才考虑MQA。这个回答既体现了你对权衡的理解又贴合工业界的实际选择。4. MLA登场DeepSeek的KV压缩新思路4.1 MLA到底解决了什么问题MLAMulti-head Latent Attention是DeepSeek-V2提出的结构最近面试被问的频率明显上升。它要解决的核心问题和MQA/GQA一样——KV Cache太大但思路完全不同。MQA和GQA是“减少KV的头数”而MLA是“把KV压缩到一个低维潜空间”。具体来说MLA不直接缓存完整的K和V而是缓存一个压缩后的潜向量c_KV用的时候再通过上投影矩阵还原出K和V。这样KV Cache的维度从num_heads × d_k降到了d_c压缩维度而d_c可以远小于num_heads × d_k。4.2 MLA的核心机制拆解MLA的数学表达看起来复杂但拆开就三步第一步压缩。把输入h_t通过下投影矩阵W_DKV压成潜向量c_KV W_DKV h_t第二步还原。用上投影矩阵W_UK和W_UV分别还原出K和Vk W_UK c_KV v W_UV c_KV第三步正常做attention。还原出的K、V和Q做标准的attention计算。关键在于推理时只需要缓存c_KV不需要缓存完整的K和V。c_KV的维度是d_c而完整K、V的维度是num_heads × d_k。DeepSeek-V2里d_c远小于num_heads × d_k所以KV Cache大幅缩小。class MultiHeadLatentAttention(nn.Module): def __init__(self, d_model, num_heads, d_c, rope_dim, dropout0.1): super().__init__() self.num_heads num_heads self.d_k d_model // num_heads self.d_c d_c # 压缩维度 self.rope_dim rope_dim # 解耦的RoPE维度 # Q的投影 self.W_q nn.Linear(d_model, num_heads * self.d_k) # KV压缩投影 self.W_DKV nn.Linear(d_model, d_c) # KV还原投影 self.W_UK nn.Linear(d_c, num_heads * self.d_k) self.W_UV nn.Linear(d_c, num_heads * self.d_k) # 输出投影 self.W_o nn.Linear(num_heads * self.d_k, d_model) def forward(self, x, maskNone): B, L, _ x.shape # Q投影并切头 Q self.W_q(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2) # KV压缩 c_KV self.W_DKV(x) # (B, L, d_c) # 还原K和V K self.W_UK(c_KV).view(B, L, self.num_heads, self.d_k).transpose(1, 2) V self.W_UV(c_KV).view(B, L, self.num_heads, self.d_k).transpose(1, 2) scores torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.d_k) if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn torch.softmax(scores, dim-1) context torch.matmul(attn, V) context context.transpose(1, 2).contiguous().view(B, L, -1) return self.W_o(context)4.3 MLA里最容易被追问的RoPE解耦上面这份代码是简化版实际MLA里有一个非常关键的细节RoPE的解耦。面试官如果懂MLA几乎一定会问这个。问题出在哪RoPE是作用在Q和K上的位置编码它要求对K的特定维度做旋转。但MLA里K是从c_KV还原出来的如果直接对还原后的K做RoPE那c_KV和W_UK的矩阵乘法就没法和RoPE的旋转操作合并推理时的计算图会变得很复杂。DeepSeek的解法是把K拆成两部分一部分走RoPE叫k_rope一部分不走RoPE叫k_nope。k_rope单独从输入投影出来维度很小比如64维k_nope从c_KV还原。最后把两部分拼起来。这样RoPE只作用在小维度上既保留了位置信息又不破坏压缩结构。这个设计在面试里是加分项。如果你能主动提到“MLA的RoPE需要解耦否则压缩和位置编码会打架”面试官基本会认为你是真读过论文的。4.4 MLA、GQA、MQA的KV Cache对比用一个具体配置来对比。假设d_model5120num_heads40d_k128seq_len4096FP16机制缓存内容缓存维度KV Cache大小MHAK, V40×128×280 MBGQA (g8)K, V8×128×216 MBMQAK, V1×128×22 MBMLAc_KV5124 MBMLA用512维的潜向量做到了接近MQA的压缩率但效果比MQA好得多。这就是为什么DeepSeek-V2能在保持效果的同时把推理成本打下来。5. 面试现场手写代码的实战策略5.1 时间分配20分钟怎么用面试手写Attention通常给15到20分钟。我的建议是前3分钟和面试官确认需求。问清楚“要不要写mask”“要不要考虑batch”“缩放因子用不用”。这一步能帮你避免写完才发现方向错了。中间10分钟写主体代码。先写MHA如果面试官要求扩展再改造成GQA或MLA。最后5分钟检查维度、补边界情况、口头解释关键行。注意不要一上来就闷头写。面试官更看重你的思考过程边写边说“这里我切多头维度变成B×h×L×d_k”比默默写完更有说服力。5.2 从MHA改造成GQA的最小改动面试官经常在你写完MHA后说“现在改成GQA。”这时候不要重写只需要改三处W_k和W_v的输出维度从d_model改成num_kv_heads * d_kview的时候用num_kv_heads而不是num_heads在计算scores之前加一行repeat_interleave# 改造点1投影层 self.W_k nn.Linear(d_model, num_kv_heads * self.d_k) self.W_v nn.Linear(d_model, num_kv_heads * self.d_k) # 改造点2view K self.W_k(x).view(B, L, self.num_kv_heads, self.d_k).transpose(1, 2) V self.W_v(x).view(B, L, self.num_kv_heads, self.d_k).transpose(1, 2) # 改造点3复制 K K.repeat_interleave(self.num_groups, dim1) V V.repeat_interleave(self.num_groups, dim1)这样改面试官会觉得你对结构理解得很透而不是死记硬背。5.3 常见追问与应对我把面试里被问最多的几个问题整理成表追问应对要点为什么除以sqrt(d_k)点积方差随d_k增长缩放防止softmax梯度消失mask怎么加在softmax之前用masked_fill填-inf注意因果mask是上三角GQA的repeat_interleave会不会增加显存训练时会推理时K、V本来就要展开不算额外开销MLA为什么能压缩KV共享低维潜空间缓存c_KV而非完整K、VMLA的RoPE为什么要解耦压缩和旋转操作不可交换解耦后RoPE只作用小维度MQA效果为什么掉所有头共享KV位置信息表达受限模型越大越明显5.4 我踩过的几个坑说几个我自己面试和实战里踩过的坑都是血泪教训。第一个坑mask的维度。因果mask的shape要和scores对齐。scores是(B, h, L, L)mask如果是(L, L)广播没问题但如果是(B, L)就会出错。我见过有人把padding mask和causal mask搞混结果生成时能看到未来token训练loss正常但推理全乱。第二个坑repeat_interleave和repeat的区别。GQA里必须用repeat_interleave它是在头维度上逐个复制用repeat会把整个张量复制一遍维度对不上。这个细节我在第一次写GQA时就翻过车。第三个坑MLA的d_c选择。d_c不是越小越好。太小了压缩损失大效果掉得厉害太大了KV Cache省不下来。DeepSeek-V2里d_c大约是d_model的1/10到1/8这个比例可以参考。第四个坑面试时不要过度设计。有人一上来就想写FlashAttention的优化版结果基础版都没写对。先把标准实现写扎实面试官如果问优化再展开。6. 把这四个机制串成一条线面到最后面试官可能会问一个开放问题“你怎么理解这几个Attention机制的关系”这时候你需要一条清晰的线索把它们串起来。我的回答框架是这样的MHA是基准它的问题在推理阶段暴露——KV Cache随序列长度和头数线性增长。MQA和GQA从“减少KV头数”这个方向解决MQA压到极致但掉效果GQA折中所以成了主流。MLA换了个方向从“压缩KV维度”入手用低维潜空间缓存配合RoPE解耦在压缩率和效果之间找到了更好的平衡点。这条线索背后是一个更本质的认知Attention的演进史就是显存、计算、效果三者不断重新平衡的历史。你理解了这一点面试时不管被问到哪个机制都能从这张三角图里找到位置。最后分享一个我自己的习惯每次准备面试我都会拿一张白纸不看任何资料把MHA、GQA、MLA的代码各默写一遍然后自己给自己讲一遍每个设计选择的理由。能讲清楚“为什么”比能写出来更重要。因为代码可以背但面试官追问的那一层背是背不出来的。