DeepSeek核心架構-MLA:剖析低秩聯合壓縮優化KV緩存、提升推理效率的技術細節
?DeepSeek的基本架構仍然在Transformer框架內,每個Transformer模塊由一個注意力模塊和一個前饋網絡組成。為實現更高效的推理和更經濟的訓練,在注意力和前饋網絡部分,設計并使用了創新的MLA(Multi-Head Latent Attention)和DeepSeekMoE 架構。
本文將從MLA的提出背景、技術原理、解耦RoPE策略及MHA與MLA的緩存對比方面進行詳細闡述。
MLA是對多頭自注意力機制(MHA)的改進,其核心是對鍵(Keys)和值(Values)進行低秩聯合壓縮,生成壓縮的潛在向量,以降低推理階段的KV緩存。
1.MHA
(1)MHA的詳細原理。
(2)在自注意力機制中,每個位置的查詢Q需要與序列中所有位置的鍵(K)進行相似度計算得到注意力分數,然后加權值(V)獲得最終的輸出值。Transformer模型生成序列時采用自回歸方式,即在生成當前Token時需依賴所有之前生成的Token。因此,對于每個時間步,都需要與之前所有Token進行注意力計算。為避免重復計算之前所有Token的鍵和值,從而需要緩存所有的鍵和值來加速計算。
(3)MHA在推理過程中需要緩存所有的KV來加速推理,但是在模型部署中,這種量級的KV緩存帶來了高顯存占用問題,限制了最大Batch size的大小和序列長度。
2.MLA原理
(1)MLA的完整計算公式和維度變化示意圖如下,可對照查看。
圖片
3.解耦RoPE策略詳解
備注:矩陣吸收計算是指利用矩陣乘法的結合律或低秩分解等線性代數技巧,重新組合某些矩陣因子,使原本需要獨立計算的矩陣乘積合并在一起,從而降低計算復雜度和內存開銷的過程。
04|MHA與MLA的緩存對比
本文轉載自 ??南夏的算法驛站??,作者: 南夏的算法驛站
