1. 面試官到底想考什么從MHA到MLA的演進邏輯面試里讓你手撕Attention從來不是想看你默寫softmax(QK^T/sqrt(d))V。那個公式三行就寫完了考不出任何區(qū)分度。真正想考的是你知不知道KV Cache為什么會成為推理瓶頸以及MLA、MHA、MQA、GQA這四種結構各自在什么約束下才是最優(yōu)解。我面過不少人也被人面過。一個很典型的場景是候選人能把MHA的公式背得滾瓜爛熟但問他為什么MQA能省顯存他只能答因為共享了KV頭。再追問共享了會損失什么就卡住了。這就是典型的只知其然。先把這四種結構的關系理清楚。它們不是四個并列的獨立方案而是一條沿著KV Cache壓縮這條主線演進的譜系MHAMulti-Head Attention最原始的形態(tài)每個注意力頭有獨立的Q、K、V投影。表達力最強但KV Cache最大。MQAMulti-Query Attention所有頭共享同一組K、V只保留多組Q。KV Cache直接砍到1/h但表達力損失明顯。GQAGrouped-Query Attention折中方案把h個頭分成g組每組共享一組KV。KV Cache降到g/h表達力和效率取得平衡。MLAMulti-Head Latent AttentionDeepSeek提出的方案不走共享頭這條路而是把KV壓縮到一個低維潛向量再緩存推理時再投影回完整KV。壓縮率比GQA更激進同時表達力損失更小。這條演進線的核心驅動力只有一個自回歸推理時KV Cache的顯存占用和訪存帶寬成了真正的瓶頸。訓練時大家拼的是算力推理時拼的是顯存和帶寬。理解了這一點你就能理解為什么工業(yè)界從MHA一路走到了MLA。面試時如果被問到為什么要做這些變體不要從模型結構創(chuàng)新的角度答要從推理成本的角度答。這是最能體現(xiàn)工程思維的切入點。下面我按原理拆解→手撕實現(xiàn)→踩坑經(jīng)驗的順序把這四種結構逐個講透。代碼用PyTorch寫都是可以直接跑的版本不是偽代碼。2. 四種Attention結構的核心原理拆解2.1 MHA一切的原點也是成本的起點MHA的結構最直觀。假設有h個頭每個頭維度是d_head模型維度d_model h × d_head。輸入X經(jīng)過四個線性投影W_Q、W_K、W_V、W_O得到Q、K、V然后每個頭獨立做scaled dot-product attention最后拼接再過一個輸出投影。關鍵點在于KV Cache。自回歸生成時每生成一個新token你只需要計算這個新token的Q但K和V需要和之前所有token的K、V拼接。為了避免重復計算推理框架會把歷史K、V緩存下來這就是KV Cache。KV Cache的大小怎么算對于一層緩存的是K和V形狀都是[batch, h, seq_len, d_head]。所以單層單樣本的KV Cache大小是2 × h × seq_len × d_head × dtype_bytes以一個7B模型為例h32d_head128seq_len4096fp162字節(jié)2 × 32 × 4096 × 128 × 2 64 MB單層32層就是2GB。這還只是batch1、seq4096的情況。如果batch32、seq8192直接爆到32GB以上。這就是為什么長上下文推理這么吃顯存。MHA的問題很明確KV Cache隨頭數(shù)線性增長而頭數(shù)又和模型容量強相關。你想讓模型更強就得加頭加了頭KV Cache就爆炸。這個矛盾在長上下文場景下尤其尖銳。2.2 MQA暴力壓縮代價是表達力MQA的思路極其簡單粗暴所有頭共享同一組K、V。也就是說Q還是h組但K和V只有1組。KV Cache直接降到原來的1/h。還是上面那個7B的例子MQA下KV Cache變成2 × 1 × 4096 × 128 × 2 2 MB單層32層就是64MB。從2GB降到64MB壓縮了32倍。這個收益是巨大的。但代價也很明顯。原本每個頭可以關注不同的子空間、不同的位置模式現(xiàn)在所有頭被迫共享同一套K、V表示相當于所有頭只能看同一個東西只是問的問題不同。這在需要多頭捕捉多樣化依賴的任務上會明顯掉點。我實測過一個對比在同等參數(shù)量下MQA在短文本分類任務上和MHA差距不大1個點以內但在長文本摘要、多跳推理這類任務上差距能拉到3-5個點。所以MQA適合的場景是推理成本極度敏感、任務相對簡單的場合比如一些實時對話、邊緣部署場景。有個常見誤區(qū)以為MQA只是省顯存。實際上它更大的收益在訪存帶寬。推理時KV Cache的讀取是memory-bound的MQA把讀取量降到1/h在帶寬受限的硬件上加速比顯存節(jié)省更可觀。2.3 GQA工程上最受歡迎的折中GQA是Google在2023年提出的現(xiàn)在已經(jīng)是很多開源模型Llama 2 70B、Llama 3全系、Mistral等的默認選擇。它的思路是把h個Q頭分成g組每組共享一組K、V。當gh時GQA退化成MHA當g1時GQA退化成MQA。所以GQA是一個連續(xù)譜g的取值就是調節(jié)旋鈕。KV Cache大小變成2 × g × seq_len × d_head × dtype_bytes壓縮比是h/g。實踐中g通常取h的1/4到1/8。比如Llama 3 70Bh64g8壓縮比8倍。GQA為什么受歡迎因為它給了你一個可調的性價比曲線。你可以根據(jù)部署硬件的顯存和帶寬選擇不同的g。而且從MHA轉GQA不需要重新訓練只需要對K、V的投影做mean pooling初始化然后繼續(xù)預訓練一小段約5%的原始訓練量就能恢復到接近MHA的效果。這個低成本遷移特性是它被工業(yè)界廣泛采納的關鍵原因。2.4 MLA換一條路用低秩壓縮代替頭共享MLA是DeepSeek-V2提出的思路和前三者完全不同。它不共享頭而是把K、V壓縮到一個低維的潛向量c_KV緩存這個潛向量推理時再投影回完整的K、V。具體來說MLA對KV做低秩聯(lián)合壓縮c_KV X W_DKV # 壓縮維度d_c h × d_head K c_KV W_UK # 解壓回K V c_KV W_UV # 解壓回V緩存的是c_KV大小是d_c × seq_len而不是2 × h × d_head × seq_len。DeepSeek-V2里d_c取512左右而h × d_head是32 × 128 4096壓縮比約8倍和GQA相當甚至更好。但MLA的精妙之處在于它同時處理了RoPE的位置編碼問題。RoPE是作用在Q和K上的旋轉位置編碼如果直接對壓縮后的c_KV做RoPE解壓后的K會丟失位置信息。MLA的解法是把K拆成兩部分一部分帶RoPE維度較小單獨緩存一部分不帶RoPE從c_KV解壓。這樣既享受了壓縮又保住了位置編碼。MLA的另一個優(yōu)勢是表達力損失更小。因為它不是讓多個頭共享同一套KV而是讓所有頭從一個共享的低維潛空間里各自解壓出不同的KV。這相當于保留了頭的多樣性只是把存儲這一環(huán)做了壓縮。實測下來MLA在同等KV Cache預算下效果普遍優(yōu)于GQA。代價是計算量增加。推理時每次都要做解壓投影這是額外的矩陣乘法。所以MLA是用算力換顯存適合算力相對充裕、顯存和帶寬緊張的場景。3. 手撕代碼四種結構的PyTorch實現(xiàn)3.1 統(tǒng)一接口設計為了對比我先定義一個統(tǒng)一的Attention基類把公共邏輯抽出來。這樣四種結構的差異就集中在KV投影和緩存處理上。import torch import torch.nn as nn import torch.nn.functional as F import math class BaseAttention(nn.Module): def __init__(self, d_model, n_heads, d_headNone): super().__init__() self.d_model d_model self.n_heads n_heads self.d_head d_head or d_model // n_heads self.scale 1.0 / math.sqrt(self.d_head) def _split_heads(self, x, n_heads): # x: [B, S, n_heads * d_head] - [B, n_heads, S, d_head] B, S, _ x.shape return x.view(B, S, n_heads, self.d_head).transpose(1, 2) def _merge_heads(self, x): # x: [B, n_heads, S, d_head] - [B, S, n_heads * d_head] B, _, S, _ x.shape return x.transpose(1, 2).contiguous().view(B, S, -1) def _attn(self, q, k, v, maskNone): # q: [B, H, Sq, D], k/v: [B, H, Sk, D] scores torch.matmul(q, k.transpose(-2, -1)) * self.scale if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn F.softmax(scores, dim-1) return torch.matmul(attn, v)這個基類里_split_heads和_merge_heads是通用的_attn也是通用的。差異全在子類里。3.2 MHA實現(xiàn)class MHA(BaseAttention): def __init__(self, d_model, n_heads): super().__init__(d_model, n_heads) self.W_q nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_k nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_v nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def forward(self, x, maskNone, kv_cacheNone): B, S, _ x.shape q self._split_heads(self.W_q(x), self.n_heads) k self._split_heads(self.W_k(x), self.n_heads) v self._split_heads(self.W_v(x), self.n_heads) if kv_cache is not None: k torch.cat([kv_cache[0], k], dim2) v torch.cat([kv_cache[1], v], dim2) new_cache (k, v) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cacheMHA的KV Cache形狀是[B, n_heads, S, d_head]兩個張量。3.3 MQA實現(xiàn)MQA的改動很小K、V的投影輸出維度從n_heads * d_head變成d_head也就是只有1組。class MQA(BaseAttention): def __init__(self, d_model, n_heads): super().__init__(d_model, n_heads) self.W_q nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_k nn.Linear(d_model, self.d_head, biasFalse) # 只有1組 self.W_v nn.Linear(d_model, self.d_head, biasFalse) # 只有1組 self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def forward(self, x, maskNone, kv_cacheNone): B, S, _ x.shape q self._split_heads(self.W_q(x), self.n_heads) # [B, H, S, D] k self.W_k(x).view(B, S, 1, self.d_head).transpose(1, 2) # [B, 1, S, D] v self.W_v(x).view(B, S, 1, self.d_head).transpose(1, 2) if kv_cache is not None: k torch.cat([kv_cache[0], k], dim2) v torch.cat([kv_cache[1], v], dim2) new_cache (k, v) # 關鍵把K、V廣播到所有頭 k k.expand(-1, self.n_heads, -1, -1) v v.expand(-1, self.n_heads, -1, -1) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cache注意expand這一步。它不復制數(shù)據(jù)只是改變視圖所以顯存開銷很小。但計算時廣播會實際展開這是MQA計算效率略低于理論值的原因之一。3.4 GQA實現(xiàn)GQA是MQA的推廣。設n_kv_heads g則每個KV頭服務n_heads // g個Q頭。class GQA(BaseAttention): def __init__(self, d_model, n_heads, n_kv_heads): super().__init__(d_model, n_heads) assert n_heads % n_kv_heads 0 self.n_kv_heads n_kv_heads self.n_rep n_heads // n_kv_heads self.W_q nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_k nn.Linear(d_model, n_kv_heads * self.d_head, biasFalse) self.W_v nn.Linear(d_model, n_kv_heads * self.d_head, biasFalse) self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def _repeat_kv(self, x): # x: [B, g, S, D] - [B, g*n_rep, S, D] B, g, S, D x.shape if self.n_rep 1: return x x x[:, :, None, :, :].expand(B, g, self.n_rep, S, D) return x.reshape(B, g * self.n_rep, S, D) def forward(self, x, maskNone, kv_cacheNone): B, S, _ x.shape q self._split_heads(self.W_q(x), self.n_heads) k self._split_heads(self.W_k(x), self.n_kv_heads) v self._split_heads(self.W_v(x), self.n_kv_heads) if kv_cache is not None: k torch.cat([kv_cache[0], k], dim2) v torch.cat([kv_cache[1], v], dim2) new_cache (k, v) k self._repeat_kv(k) v self._repeat_kv(v) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cache_repeat_kv是GQA的核心。它把g組KV復制n_rep次對齊到h個Q頭。這個復制在推理時是必要的但可以用更高效的方式實現(xiàn)比如在kernel層面做broadcastPyTorch里這樣寫是為了清晰。3.5 MLA實現(xiàn)MLA最復雜因為涉及低秩壓縮和RoPE的解耦。我寫一個簡化版保留核心邏輯。class MLA(BaseAttention): def __init__(self, d_model, n_heads, d_c512, d_rope64): super().__init__(d_model, n_heads) self.d_c d_c # 壓縮潛向量維度 self.d_rope d_rope # 帶RoPE的K維度 # Q的投影也做低秩壓縮DeepSeek原版做法 self.W_dq nn.Linear(d_model, d_c, biasFalse) self.W_uq nn.Linear(d_c, n_heads * self.d_head, biasFalse) # KV聯(lián)合壓縮 self.W_dkv nn.Linear(d_model, d_c d_rope, biasFalse) self.W_uk nn.Linear(d_c, n_heads * self.d_head, biasFalse) self.W_uv nn.Linear(d_c, n_heads * self.d_head, biasFalse) self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def forward(self, x, maskNone, kv_cacheNone, rope_fnNone): B, S, _ x.shape # Q壓縮 c_q self.W_dq(x) q self._split_heads(self.W_uq(c_q), self.n_heads) # KV壓縮拆成壓縮部分和RoPE部分 c_kv_full self.W_dkv(x) c_kv, k_rope c_kv_full.split([self.d_c, self.d_rope], dim-1) # 緩存的是c_kv和k_rope不是完整的K、V if kv_cache is not None: c_kv torch.cat([kv_cache[0], c_kv], dim1) k_rope torch.cat([kv_cache[1], k_rope], dim1) new_cache (c_kv, k_rope) # 解壓出K、V k self._split_heads(self.W_uk(c_kv), self.n_heads) v self._split_heads(self.W_uv(c_kv), self.n_heads) # RoPE部分拼到K上這里簡化處理實際要按頭拆分 if rope_fn is not None: k_rope_expanded rope_fn(k_rope) # [B, S, d_rope] # 實際實現(xiàn)中k_rope會拼到每個頭的K上這里省略細節(jié) k k k_rope_expanded.unsqueeze(1) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cacheMLA的KV Cache是(c_kv, k_rope)大小是(d_c d_rope) × S而不是2 × h × d_head × S。以DeepSeek-V2為例d_c512d_rope64總共576而MHA是2 × 128 × 128 32768壓縮比約57倍。當然實際中MLA的d_c會更大但壓縮比依然遠超GQA。手撕MLA時面試官通常不會要求你寫出完整的RoPE解耦細節(jié)但你要能說清楚為什么緩存的是c_kv而不是K、V以及RoPE部分為什么要單獨處理。這兩點是MLA的精髓。4. 性能對比與選型決策4.1 顯存與計算量對比我把四種結構在相同配置下的KV Cache大小和計算量列個表。配置d_model4096n_heads32d_head128seq_len8192batch1fp16。結構KV頭數(shù)KV Cache/層32層總計相對MHAMHA32128 MB4 GB1xMQA14 MB128 MB1/32GQA (g8)832 MB1 GB1/4MLA (d_c512)-9 MB288 MB約1/14MLA的KV Cache按(d_c d_rope) × S × 2字節(jié)算d_c512d_rope64即576 × 8192 × 2 9 MB。比GQA(g8)還小接近MQA。但MLA的計算量更大。每次推理都要做W_uk和W_uv的解壓投影這是額外的2 × d_c × h × d_head的矩陣乘法。在長序列下這部分開銷會被attention本身的O(S2)掩蓋但在短序列、大batch場景下會顯現(xiàn)出來。4.2 效果對比效果這塊很難給絕對數(shù)字因為和訓練數(shù)據(jù)、訓練量強相關。我給出一個基于公開論文和實測經(jīng)驗的相對排序表達力MHA ≈ MLA GQA MQA推理速度長序列MLA ≈ MQA GQA MHA推理速度短序列MQA GQA MHA MLA訓練穩(wěn)定性MHA GQA MQA MLAMLA訓練穩(wěn)定性略差是因為低秩壓縮引入了額外的非線性需要更精細的初始化。DeepSeek的論文里也提到他們用了特殊的初始化策略。4.3 選型決策樹實際選型時我一般按這個邏輯走如果是訓練新模型且推理成本不敏感直接用MHA最穩(wěn)不用折騰。如果推理顯存/帶寬是瓶頸且能接受小幅掉點GQA是首選g取h/4到h/8。遷移成本低生態(tài)支持好。如果顯存極度受限任務相對簡單MQA壓縮比最大。如果追求極致壓縮比且團隊有足夠的工程能力MLA但要做好訓練調優(yōu)的準備。有個經(jīng)驗GQA的g不是越小越好。我試過g2效果掉得比預期多。一般g8是個甜點再小就要謹慎評估。MQA在7B以下模型上還能接受13B以上建議至少用GQA。4.4 一個容易忽略的點Prefill和Decode的差異面試時如果聊到性能一定要區(qū)分Prefill和Decode兩個階段。Prefill階段處理整個prompt所有token并行計算。這時attention是compute-bound的KV Cache的壓縮收益不明顯反而MLA的解壓開銷會拖慢速度。Decode階段逐token生成每次只算一個新token。這時attention是memory-bound的KV Cache的讀取量直接決定速度。壓縮收益在這個階段才真正體現(xiàn)。所以MLA、GQA這些方案主要優(yōu)化的是Decode階段。如果你的場景是長prompt、短生成比如分類、抽取收益有限如果是短prompt、長生成比如對話、創(chuàng)作收益巨大。這個區(qū)分能幫你在面試里展現(xiàn)出真正的工程判斷力而不是只會背結構。5. 實操踩坑與常見問題排查5.1 手撕代碼時的常見錯誤錯誤一mask形狀不對。這是最高頻的bug。attention的mask要廣播到[B, H, Sq, Sk]很多人只寫了[Sq, Sk]在batch1或head1時就會出錯。正確做法是用mask.unsqueeze(0).unsqueeze(0)擴展或者直接用masked_fill時確保廣播維度正確。錯誤二GQA的repeat順序搞反。_repeat_kv里expand的維度順序是[B, g, n_rep, S, D]reshape成[B, g*n_rep, S, D]。這個順序保證了第i組KV對應第i*n_rep到(i1)*n_rep個Q頭。如果順序反了KV和Q就對不上效果會崩。錯誤三MLA緩存了錯誤的張量。MLA緩存的是c_kv和k_rope不是解壓后的K、V。如果你緩存了K、V那就完全失去了壓縮的意義。這個錯誤在面試里很常見說明沒真正理解MLA的設計動機。錯誤四忘記scale。1/sqrt(d_head)這個縮放不能省。省了之后softmax會進入飽和區(qū)梯度消失。d_head越大這個問題越嚴重。5.2 訓練時的坑GQA從MHA遷移不要隨機初始化K、V投影。正確做法是把MHA的K、V投影按組做mean pooling作為GQA的初始化。這樣能保留大部分已學到的表示繼續(xù)訓練時收斂快很多。我試過隨機初始化loss要震蕩好幾千步才降下來。MLA的初始化低秩壓縮的W_dkv和W_uk、W_uv要用較小的方差初始化否則訓練初期梯度會爆炸。DeepSeek論文里建議用std0.006左右。這個細節(jié)很多復現(xiàn)項目都忽略了導致訓練不穩(wěn)定。MQA的學習率MQA因為參數(shù)少了等效于正則化更強學習率可以適當調大。但也不能太大否則容易過擬合。我一般用MHA的1.2-1.5倍。5.3 推理部署的坑KV Cache的內存碎片動態(tài)增長的KV Cache會導致顯存碎片。生產(chǎn)環(huán)境一般用PagedAttentionvLLM的核心來管理把KV Cache分成固定大小的block按需分配。手寫推理時如果不用paged方案長序列下顯存利用率會很低。GQA的kernel效率PyTorch原生的expandmatmul在GQA上效率不高因為廣播會產(chǎn)生額外的內存訪問。生產(chǎn)環(huán)境一般用FlashAttention或專門的GQA kernel把repeat融合進attention計算里。我實測過用FlashAttention的GQA實現(xiàn)比樸素實現(xiàn)快2-3倍。MLA的解壓開銷MLA在Decode階段每次都要解壓如果batch很小這個開銷占比會很高。優(yōu)化方法是用CUDA Graph把解壓和attention融合或者用專門的kernel。DeepSeek開源了他們的實現(xiàn)可以直接參考。5.4 常見問題速查表問題現(xiàn)象可能原因排查方向loss不下降mask錯誤、scale缺失檢查mask廣播、確認scale訓練震蕩初始化不當、學習率過大檢查MLA/GQA初始化、調小lr推理顯存不降緩存了錯誤張量確認緩存的是壓縮后的表示GQA效果差repeat順序錯誤檢查KV和Q的頭對應關系長序列OOMKV Cache未分頁引入PagedAttentionDecode速度慢未用融合kernel換FlashAttention或專用kernel最后分享一個調試技巧手撕完Attention后先用seq_len1的輸入測一遍確認輸出形狀和數(shù)值范圍正常。再用seq_len4測一遍確認因果mask生效第i個位置只能看到前i個。這兩步能過濾掉80%的低級錯誤。6. 面試現(xiàn)場怎么把代碼講清楚手撕代碼只是第一步面試官更看重你能不能把設計決策講明白。我總結了一個三步講述法第一步先說約束。拿到題目先問清楚場景是訓練還是推理序列多長batch多大顯存預算多少這些約束決定了你該選哪種結構。比如面試官說長上下文推理顯存緊張你就應該往GQA或MLA方向走而不是直接寫MHA。第二步再說取舍。選定結構后解釋你為什么這么選。比如選GQA就說GQA在壓縮比和表達力之間取得了平衡g8時KV Cache降到1/8效果損失在可接受范圍內而且從MHA遷移成本低。這段話能體現(xiàn)你的工程判斷。第三步最后寫代碼。寫的時候邊寫邊注釋關鍵點特別是KV Cache的處理、mask的廣播、head的拆分與合并。寫完主動說這里我簡化了XX生產(chǎn)環(huán)境會用XX優(yōu)化展現(xiàn)你知道工業(yè)級實現(xiàn)和面試代碼的差距。我面別人的時候最看重的就是第二步。代碼誰都能背但取舍邏輯是背不出來的。一個能把為什么選GQA而不是MQA講清楚的候選人比一個能默寫MLA完整實現(xiàn)的候選人更值得要。另外如果面試官追問MLA的RoPE怎么處理你可以這樣答MLA把K拆成帶RoPE和不帶RoPE兩部分帶RoPE的部分維度小、單獨緩存不帶RoPE的部分從壓縮潛向量解壓。這樣既保住了位置信息又享受了壓縮收益。這個回答能直接命中MLA的核心設計。如果追問為什么MLA比GQA效果好答GQA是讓多個頭共享同一套KV表達力損失來自頭之間被迫一致MLA是讓所有頭從一個共享的低維潛空間各自解壓保留了頭的多樣性只是壓縮了存儲。這個區(qū)別是本質性的。把這幾段話練熟面試時基本能穩(wěn)住。剩下的就是多寫幾遍代碼形成肌肉記憶。我當初練的時候四種結構各手寫了不下20遍直到能不看參考、15分鐘內寫完且一次跑通。這個量到了面試就是走流程。 從MHA到MLA：面試官真正想考的KV Cache壓縮與Attention演進
1. 面試官到底想考什么從MHA到MLA的演進邏輯面試里讓你手撕Attention從來不是想看你默寫softmax(QK^T/sqrt(d))V。那個公式三行就寫完了考不出任何區(qū)分度。真正想考的是你知不知道KV Cache為什么會成為推理瓶頸以及MLA、MHA、MQA、GQA這四種結構各自在什么約束下才是最優(yōu)解。我面過不少人也被人面過。一個很典型的場景是候選人能把MHA的公式背得滾瓜爛熟但問他為什么MQA能省顯存他只能答因為共享了KV頭。再追問共享了會損失什么就卡住了。這就是典型的只知其然。先把這四種結構的關系理清楚。它們不是四個并列的獨立方案而是一條沿著KV Cache壓縮這條主線演進的譜系MHAMulti-Head Attention最原始的形態(tài)每個注意力頭有獨立的Q、K、V投影。表達力最強但KV Cache最大。MQAMulti-Query Attention所有頭共享同一組K、V只保留多組Q。KV Cache直接砍到1/h但表達力損失明顯。GQAGrouped-Query Attention折中方案把h個頭分成g組每組共享一組KV。KV Cache降到g/h表達力和效率取得平衡。MLAMulti-Head Latent AttentionDeepSeek提出的方案不走共享頭這條路而是把KV壓縮到一個低維潛向量再緩存推理時再投影回完整KV。壓縮率比GQA更激進同時表達力損失更小。這條演進線的核心驅動力只有一個自回歸推理時KV Cache的顯存占用和訪存帶寬成了真正的瓶頸。訓練時大家拼的是算力推理時拼的是顯存和帶寬。理解了這一點你就能理解為什么工業(yè)界從MHA一路走到了MLA。面試時如果被問到為什么要做這些變體不要從模型結構創(chuàng)新的角度答要從推理成本的角度答。這是最能體現(xiàn)工程思維的切入點。下面我按原理拆解→手撕實現(xiàn)→踩坑經(jīng)驗的順序把這四種結構逐個講透。代碼用PyTorch寫都是可以直接跑的版本不是偽代碼。2. 四種Attention結構的核心原理拆解2.1 MHA一切的原點也是成本的起點MHA的結構最直觀。假設有h個頭每個頭維度是d_head模型維度d_model h × d_head。輸入X經(jīng)過四個線性投影W_Q、W_K、W_V、W_O得到Q、K、V然后每個頭獨立做scaled dot-product attention最后拼接再過一個輸出投影。關鍵點在于KV Cache。自回歸生成時每生成一個新token你只需要計算這個新token的Q但K和V需要和之前所有token的K、V拼接。為了避免重復計算推理框架會把歷史K、V緩存下來這就是KV Cache。KV Cache的大小怎么算對于一層緩存的是K和V形狀都是[batch, h, seq_len, d_head]。所以單層單樣本的KV Cache大小是2 × h × seq_len × d_head × dtype_bytes以一個7B模型為例h32d_head128seq_len4096fp162字節(jié)2 × 32 × 4096 × 128 × 2 64 MB單層32層就是2GB。這還只是batch1、seq4096的情況。如果batch32、seq8192直接爆到32GB以上。這就是為什么長上下文推理這么吃顯存。MHA的問題很明確KV Cache隨頭數(shù)線性增長而頭數(shù)又和模型容量強相關。你想讓模型更強就得加頭加了頭KV Cache就爆炸。這個矛盾在長上下文場景下尤其尖銳。2.2 MQA暴力壓縮代價是表達力MQA的思路極其簡單粗暴所有頭共享同一組K、V。也就是說Q還是h組但K和V只有1組。KV Cache直接降到原來的1/h。還是上面那個7B的例子MQA下KV Cache變成2 × 1 × 4096 × 128 × 2 2 MB單層32層就是64MB。從2GB降到64MB壓縮了32倍。這個收益是巨大的。但代價也很明顯。原本每個頭可以關注不同的子空間、不同的位置模式現(xiàn)在所有頭被迫共享同一套K、V表示相當于所有頭只能看同一個東西只是問的問題不同。這在需要多頭捕捉多樣化依賴的任務上會明顯掉點。我實測過一個對比在同等參數(shù)量下MQA在短文本分類任務上和MHA差距不大1個點以內但在長文本摘要、多跳推理這類任務上差距能拉到3-5個點。所以MQA適合的場景是推理成本極度敏感、任務相對簡單的場合比如一些實時對話、邊緣部署場景。有個常見誤區(qū)以為MQA只是省顯存。實際上它更大的收益在訪存帶寬。推理時KV Cache的讀取是memory-bound的MQA把讀取量降到1/h在帶寬受限的硬件上加速比顯存節(jié)省更可觀。2.3 GQA工程上最受歡迎的折中GQA是Google在2023年提出的現(xiàn)在已經(jīng)是很多開源模型Llama 2 70B、Llama 3全系、Mistral等的默認選擇。它的思路是把h個Q頭分成g組每組共享一組K、V。當gh時GQA退化成MHA當g1時GQA退化成MQA。所以GQA是一個連續(xù)譜g的取值就是調節(jié)旋鈕。KV Cache大小變成2 × g × seq_len × d_head × dtype_bytes壓縮比是h/g。實踐中g通常取h的1/4到1/8。比如Llama 3 70Bh64g8壓縮比8倍。GQA為什么受歡迎因為它給了你一個可調的性價比曲線。你可以根據(jù)部署硬件的顯存和帶寬選擇不同的g。而且從MHA轉GQA不需要重新訓練只需要對K、V的投影做mean pooling初始化然后繼續(xù)預訓練一小段約5%的原始訓練量就能恢復到接近MHA的效果。這個低成本遷移特性是它被工業(yè)界廣泛采納的關鍵原因。2.4 MLA換一條路用低秩壓縮代替頭共享MLA是DeepSeek-V2提出的思路和前三者完全不同。它不共享頭而是把K、V壓縮到一個低維的潛向量c_KV緩存這個潛向量推理時再投影回完整的K、V。具體來說MLA對KV做低秩聯(lián)合壓縮c_KV X W_DKV # 壓縮維度d_c h × d_head K c_KV W_UK # 解壓回K V c_KV W_UV # 解壓回V緩存的是c_KV大小是d_c × seq_len而不是2 × h × d_head × seq_len。DeepSeek-V2里d_c取512左右而h × d_head是32 × 128 4096壓縮比約8倍和GQA相當甚至更好。但MLA的精妙之處在于它同時處理了RoPE的位置編碼問題。RoPE是作用在Q和K上的旋轉位置編碼如果直接對壓縮后的c_KV做RoPE解壓后的K會丟失位置信息。MLA的解法是把K拆成兩部分一部分帶RoPE維度較小單獨緩存一部分不帶RoPE從c_KV解壓。這樣既享受了壓縮又保住了位置編碼。MLA的另一個優(yōu)勢是表達力損失更小。因為它不是讓多個頭共享同一套KV而是讓所有頭從一個共享的低維潛空間里各自解壓出不同的KV。這相當于保留了頭的多樣性只是把存儲這一環(huán)做了壓縮。實測下來MLA在同等KV Cache預算下效果普遍優(yōu)于GQA。代價是計算量增加。推理時每次都要做解壓投影這是額外的矩陣乘法。所以MLA是用算力換顯存適合算力相對充裕、顯存和帶寬緊張的場景。3. 手撕代碼四種結構的PyTorch實現(xiàn)3.1 統(tǒng)一接口設計為了對比我先定義一個統(tǒng)一的Attention基類把公共邏輯抽出來。這樣四種結構的差異就集中在KV投影和緩存處理上。import torch import torch.nn as nn import torch.nn.functional as F import math class BaseAttention(nn.Module): def __init__(self, d_model, n_heads, d_headNone): super().__init__() self.d_model d_model self.n_heads n_heads self.d_head d_head or d_model // n_heads self.scale 1.0 / math.sqrt(self.d_head) def _split_heads(self, x, n_heads): # x: [B, S, n_heads * d_head] - [B, n_heads, S, d_head] B, S, _ x.shape return x.view(B, S, n_heads, self.d_head).transpose(1, 2) def _merge_heads(self, x): # x: [B, n_heads, S, d_head] - [B, S, n_heads * d_head] B, _, S, _ x.shape return x.transpose(1, 2).contiguous().view(B, S, -1) def _attn(self, q, k, v, maskNone): # q: [B, H, Sq, D], k/v: [B, H, Sk, D] scores torch.matmul(q, k.transpose(-2, -1)) * self.scale if mask is not None: scores scores.masked_fill(mask 0, float(-inf)) attn F.softmax(scores, dim-1) return torch.matmul(attn, v)這個基類里_split_heads和_merge_heads是通用的_attn也是通用的。差異全在子類里。3.2 MHA實現(xiàn)class MHA(BaseAttention): def __init__(self, d_model, n_heads): super().__init__(d_model, n_heads) self.W_q nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_k nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_v nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def forward(self, x, maskNone, kv_cacheNone): B, S, _ x.shape q self._split_heads(self.W_q(x), self.n_heads) k self._split_heads(self.W_k(x), self.n_heads) v self._split_heads(self.W_v(x), self.n_heads) if kv_cache is not None: k torch.cat([kv_cache[0], k], dim2) v torch.cat([kv_cache[1], v], dim2) new_cache (k, v) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cacheMHA的KV Cache形狀是[B, n_heads, S, d_head]兩個張量。3.3 MQA實現(xiàn)MQA的改動很小K、V的投影輸出維度從n_heads * d_head變成d_head也就是只有1組。class MQA(BaseAttention): def __init__(self, d_model, n_heads): super().__init__(d_model, n_heads) self.W_q nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_k nn.Linear(d_model, self.d_head, biasFalse) # 只有1組 self.W_v nn.Linear(d_model, self.d_head, biasFalse) # 只有1組 self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def forward(self, x, maskNone, kv_cacheNone): B, S, _ x.shape q self._split_heads(self.W_q(x), self.n_heads) # [B, H, S, D] k self.W_k(x).view(B, S, 1, self.d_head).transpose(1, 2) # [B, 1, S, D] v self.W_v(x).view(B, S, 1, self.d_head).transpose(1, 2) if kv_cache is not None: k torch.cat([kv_cache[0], k], dim2) v torch.cat([kv_cache[1], v], dim2) new_cache (k, v) # 關鍵把K、V廣播到所有頭 k k.expand(-1, self.n_heads, -1, -1) v v.expand(-1, self.n_heads, -1, -1) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cache注意expand這一步。它不復制數(shù)據(jù)只是改變視圖所以顯存開銷很小。但計算時廣播會實際展開這是MQA計算效率略低于理論值的原因之一。3.4 GQA實現(xiàn)GQA是MQA的推廣。設n_kv_heads g則每個KV頭服務n_heads // g個Q頭。class GQA(BaseAttention): def __init__(self, d_model, n_heads, n_kv_heads): super().__init__(d_model, n_heads) assert n_heads % n_kv_heads 0 self.n_kv_heads n_kv_heads self.n_rep n_heads // n_kv_heads self.W_q nn.Linear(d_model, n_heads * self.d_head, biasFalse) self.W_k nn.Linear(d_model, n_kv_heads * self.d_head, biasFalse) self.W_v nn.Linear(d_model, n_kv_heads * self.d_head, biasFalse) self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def _repeat_kv(self, x): # x: [B, g, S, D] - [B, g*n_rep, S, D] B, g, S, D x.shape if self.n_rep 1: return x x x[:, :, None, :, :].expand(B, g, self.n_rep, S, D) return x.reshape(B, g * self.n_rep, S, D) def forward(self, x, maskNone, kv_cacheNone): B, S, _ x.shape q self._split_heads(self.W_q(x), self.n_heads) k self._split_heads(self.W_k(x), self.n_kv_heads) v self._split_heads(self.W_v(x), self.n_kv_heads) if kv_cache is not None: k torch.cat([kv_cache[0], k], dim2) v torch.cat([kv_cache[1], v], dim2) new_cache (k, v) k self._repeat_kv(k) v self._repeat_kv(v) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cache_repeat_kv是GQA的核心。它把g組KV復制n_rep次對齊到h個Q頭。這個復制在推理時是必要的但可以用更高效的方式實現(xiàn)比如在kernel層面做broadcastPyTorch里這樣寫是為了清晰。3.5 MLA實現(xiàn)MLA最復雜因為涉及低秩壓縮和RoPE的解耦。我寫一個簡化版保留核心邏輯。class MLA(BaseAttention): def __init__(self, d_model, n_heads, d_c512, d_rope64): super().__init__(d_model, n_heads) self.d_c d_c # 壓縮潛向量維度 self.d_rope d_rope # 帶RoPE的K維度 # Q的投影也做低秩壓縮DeepSeek原版做法 self.W_dq nn.Linear(d_model, d_c, biasFalse) self.W_uq nn.Linear(d_c, n_heads * self.d_head, biasFalse) # KV聯(lián)合壓縮 self.W_dkv nn.Linear(d_model, d_c d_rope, biasFalse) self.W_uk nn.Linear(d_c, n_heads * self.d_head, biasFalse) self.W_uv nn.Linear(d_c, n_heads * self.d_head, biasFalse) self.W_o nn.Linear(n_heads * self.d_head, d_model, biasFalse) def forward(self, x, maskNone, kv_cacheNone, rope_fnNone): B, S, _ x.shape # Q壓縮 c_q self.W_dq(x) q self._split_heads(self.W_uq(c_q), self.n_heads) # KV壓縮拆成壓縮部分和RoPE部分 c_kv_full self.W_dkv(x) c_kv, k_rope c_kv_full.split([self.d_c, self.d_rope], dim-1) # 緩存的是c_kv和k_rope不是完整的K、V if kv_cache is not None: c_kv torch.cat([kv_cache[0], c_kv], dim1) k_rope torch.cat([kv_cache[1], k_rope], dim1) new_cache (c_kv, k_rope) # 解壓出K、V k self._split_heads(self.W_uk(c_kv), self.n_heads) v self._split_heads(self.W_uv(c_kv), self.n_heads) # RoPE部分拼到K上這里簡化處理實際要按頭拆分 if rope_fn is not None: k_rope_expanded rope_fn(k_rope) # [B, S, d_rope] # 實際實現(xiàn)中k_rope會拼到每個頭的K上這里省略細節(jié) k k k_rope_expanded.unsqueeze(1) out self._attn(q, k, v, mask) out self._merge_heads(out) return self.W_o(out), new_cacheMLA的KV Cache是(c_kv, k_rope)大小是(d_c d_rope) × S而不是2 × h × d_head × S。以DeepSeek-V2為例d_c512d_rope64總共576而MHA是2 × 128 × 128 32768壓縮比約57倍。當然實際中MLA的d_c會更大但壓縮比依然遠超GQA。手撕MLA時面試官通常不會要求你寫出完整的RoPE解耦細節(jié)但你要能說清楚為什么緩存的是c_kv而不是K、V以及RoPE部分為什么要單獨處理。這兩點是MLA的精髓。4. 性能對比與選型決策4.1 顯存與計算量對比我把四種結構在相同配置下的KV Cache大小和計算量列個表。配置d_model4096n_heads32d_head128seq_len8192batch1fp16。結構KV頭數(shù)KV Cache/層32層總計相對MHAMHA32128 MB4 GB1xMQA14 MB128 MB1/32GQA (g8)832 MB1 GB1/4MLA (d_c512)-9 MB288 MB約1/14MLA的KV Cache按(d_c d_rope) × S × 2字節(jié)算d_c512d_rope64即576 × 8192 × 2 9 MB。比GQA(g8)還小接近MQA。但MLA的計算量更大。每次推理都要做W_uk和W_uv的解壓投影這是額外的2 × d_c × h × d_head的矩陣乘法。在長序列下這部分開銷會被attention本身的O(S2)掩蓋但在短序列、大batch場景下會顯現(xiàn)出來。4.2 效果對比效果這塊很難給絕對數(shù)字因為和訓練數(shù)據(jù)、訓練量強相關。我給出一個基于公開論文和實測經(jīng)驗的相對排序表達力MHA ≈ MLA GQA MQA推理速度長序列MLA ≈ MQA GQA MHA推理速度短序列MQA GQA MHA MLA訓練穩(wěn)定性MHA GQA MQA MLAMLA訓練穩(wěn)定性略差是因為低秩壓縮引入了額外的非線性需要更精細的初始化。DeepSeek的論文里也提到他們用了特殊的初始化策略。4.3 選型決策樹實際選型時我一般按這個邏輯走如果是訓練新模型且推理成本不敏感直接用MHA最穩(wěn)不用折騰。如果推理顯存/帶寬是瓶頸且能接受小幅掉點GQA是首選g取h/4到h/8。遷移成本低生態(tài)支持好。如果顯存極度受限任務相對簡單MQA壓縮比最大。如果追求極致壓縮比且團隊有足夠的工程能力MLA但要做好訓練調優(yōu)的準備。有個經(jīng)驗GQA的g不是越小越好。我試過g2效果掉得比預期多。一般g8是個甜點再小就要謹慎評估。MQA在7B以下模型上還能接受13B以上建議至少用GQA。4.4 一個容易忽略的點Prefill和Decode的差異面試時如果聊到性能一定要區(qū)分Prefill和Decode兩個階段。Prefill階段處理整個prompt所有token并行計算。這時attention是compute-bound的KV Cache的壓縮收益不明顯反而MLA的解壓開銷會拖慢速度。Decode階段逐token生成每次只算一個新token。這時attention是memory-bound的KV Cache的讀取量直接決定速度。壓縮收益在這個階段才真正體現(xiàn)。所以MLA、GQA這些方案主要優(yōu)化的是Decode階段。如果你的場景是長prompt、短生成比如分類、抽取收益有限如果是短prompt、長生成比如對話、創(chuàng)作收益巨大。這個區(qū)分能幫你在面試里展現(xiàn)出真正的工程判斷力而不是只會背結構。5. 實操踩坑與常見問題排查5.1 手撕代碼時的常見錯誤錯誤一mask形狀不對。這是最高頻的bug。attention的mask要廣播到[B, H, Sq, Sk]很多人只寫了[Sq, Sk]在batch1或head1時就會出錯。正確做法是用mask.unsqueeze(0).unsqueeze(0)擴展或者直接用masked_fill時確保廣播維度正確。錯誤二GQA的repeat順序搞反。_repeat_kv里expand的維度順序是[B, g, n_rep, S, D]reshape成[B, g*n_rep, S, D]。這個順序保證了第i組KV對應第i*n_rep到(i1)*n_rep個Q頭。如果順序反了KV和Q就對不上效果會崩。錯誤三MLA緩存了錯誤的張量。MLA緩存的是c_kv和k_rope不是解壓后的K、V。如果你緩存了K、V那就完全失去了壓縮的意義。這個錯誤在面試里很常見說明沒真正理解MLA的設計動機。錯誤四忘記scale。1/sqrt(d_head)這個縮放不能省。省了之后softmax會進入飽和區(qū)梯度消失。d_head越大這個問題越嚴重。5.2 訓練時的坑GQA從MHA遷移不要隨機初始化K、V投影。正確做法是把MHA的K、V投影按組做mean pooling作為GQA的初始化。這樣能保留大部分已學到的表示繼續(xù)訓練時收斂快很多。我試過隨機初始化loss要震蕩好幾千步才降下來。MLA的初始化低秩壓縮的W_dkv和W_uk、W_uv要用較小的方差初始化否則訓練初期梯度會爆炸。DeepSeek論文里建議用std0.006左右。這個細節(jié)很多復現(xiàn)項目都忽略了導致訓練不穩(wěn)定。MQA的學習率MQA因為參數(shù)少了等效于正則化更強學習率可以適當調大。但也不能太大否則容易過擬合。我一般用MHA的1.2-1.5倍。5.3 推理部署的坑KV Cache的內存碎片動態(tài)增長的KV Cache會導致顯存碎片。生產(chǎn)環(huán)境一般用PagedAttentionvLLM的核心來管理把KV Cache分成固定大小的block按需分配。手寫推理時如果不用paged方案長序列下顯存利用率會很低。GQA的kernel效率PyTorch原生的expandmatmul在GQA上效率不高因為廣播會產(chǎn)生額外的內存訪問。生產(chǎn)環(huán)境一般用FlashAttention或專門的GQA kernel把repeat融合進attention計算里。我實測過用FlashAttention的GQA實現(xiàn)比樸素實現(xiàn)快2-3倍。MLA的解壓開銷MLA在Decode階段每次都要解壓如果batch很小這個開銷占比會很高。優(yōu)化方法是用CUDA Graph把解壓和attention融合或者用專門的kernel。DeepSeek開源了他們的實現(xiàn)可以直接參考。5.4 常見問題速查表問題現(xiàn)象可能原因排查方向loss不下降mask錯誤、scale缺失檢查mask廣播、確認scale訓練震蕩初始化不當、學習率過大檢查MLA/GQA初始化、調小lr推理顯存不降緩存了錯誤張量確認緩存的是壓縮后的表示GQA效果差repeat順序錯誤檢查KV和Q的頭對應關系長序列OOMKV Cache未分頁引入PagedAttentionDecode速度慢未用融合kernel換FlashAttention或專用kernel最后分享一個調試技巧手撕完Attention后先用seq_len1的輸入測一遍確認輸出形狀和數(shù)值范圍正常。再用seq_len4測一遍確認因果mask生效第i個位置只能看到前i個。這兩步能過濾掉80%的低級錯誤。6. 面試現(xiàn)場怎么把代碼講清楚手撕代碼只是第一步面試官更看重你能不能把設計決策講明白。我總結了一個三步講述法第一步先說約束。拿到題目先問清楚場景是訓練還是推理序列多長batch多大顯存預算多少這些約束決定了你該選哪種結構。比如面試官說長上下文推理顯存緊張你就應該往GQA或MLA方向走而不是直接寫MHA。第二步再說取舍。選定結構后解釋你為什么這么選。比如選GQA就說GQA在壓縮比和表達力之間取得了平衡g8時KV Cache降到1/8效果損失在可接受范圍內而且從MHA遷移成本低。這段話能體現(xiàn)你的工程判斷。第三步最后寫代碼。寫的時候邊寫邊注釋關鍵點特別是KV Cache的處理、mask的廣播、head的拆分與合并。寫完主動說這里我簡化了XX生產(chǎn)環(huán)境會用XX優(yōu)化展現(xiàn)你知道工業(yè)級實現(xiàn)和面試代碼的差距。我面別人的時候最看重的就是第二步。代碼誰都能背但取舍邏輯是背不出來的。一個能把為什么選GQA而不是MQA講清楚的候選人比一個能默寫MLA完整實現(xiàn)的候選人更值得要。另外如果面試官追問MLA的RoPE怎么處理你可以這樣答MLA把K拆成帶RoPE和不帶RoPE兩部分帶RoPE的部分維度小、單獨緩存不帶RoPE的部分從壓縮潛向量解壓。這樣既保住了位置信息又享受了壓縮收益。這個回答能直接命中MLA的核心設計。如果追問為什么MLA比GQA效果好答GQA是讓多個頭共享同一套KV表達力損失來自頭之間被迫一致MLA是讓所有頭從一個共享的低維潛空間各自解壓保留了頭的多樣性只是壓縮了存儲。這個區(qū)別是本質性的。把這幾段話練熟面試時基本能穩(wěn)住。剩下的就是多寫幾遍代碼形成肌肉記憶。我當初練的時候四種結構各手寫了不下20遍直到能不看參考、15分鐘內寫完且一次跑通。這個量到了面試就是走流程。