KV Cache

KV Cache是大模型標配的推理加速功能，也是推理過程中，顯存資源巨大開銷的元兇之一。在模型推理時，KV Cache在顯存占用量可達30%以上。

目前大部分針對KV Cache的優化工作，主要集中在工程上。比如著名的VLLM，基于paged Attention，最大限度地利用碎片化顯存空間，從而提升了空間利用率。

但是這些方案并沒有從根本上改變KV Cache占用空間巨大的問題。

我們先來看看KV Cache的基本原理，然后在文章后面詳細介紹DeepSeek MLA機制時，再來看DeepSeek-V2是怎么解決這個問題的。
KV Cache基本原理

如果不熟悉Transformer的生成過程，這部分推薦看看國外博主Jay Alammar的這篇文章：https://jalammar.github.io/illustrated-gpt2/

我們先回到Transformer計算Attention的公式，

大模型的推理是一個自回歸的過程，即一個從頭到尾逐步生成的過程。下一步的輸出，取決于上一步。

假如我們需要輸出Robot must obey orders這四個字。

模型生成第一步Robot時，會接收一個特殊字符，作為第一步的輸入，然后輸出Robot。接著將“ Robot”作為第二步的輸入，生成must，以此類推，直到模型輸出遇到最大長度限制，或者輸出了停止字符，則停止輸出過程。

我們來模擬一個輸出過程中每一步全局的Masked Attention的計算過程。這里公式中忽略了

以便展示。

每一行代表一個查詢向量與所有鍵向量的點積，但由于掩碼的存在，每個查詢只能訪問它自己和之前的鍵（因果關系）。畢竟在生成的過程中，是不能看到后續的值的。

最后輸出的softmax結果大概如下：

然后再將這個softmax結果，和對應的V值進行計算。在這個過程中，上面的矩陣計算，四行的Attention拆解下來如下,

我們可以發現這么一些規律：

每一個

的計算，只取決于當前步的，不需要以前的

。
之前的K和V，在后面會被重復利用。

那么這里就很清楚了，隨著生成序列的增加，計算Attention的增多，前面步驟的K和V都需要參與后續的計算。

所以我們如果將之前的K和V都存儲下來，就不用在當前這一步，再重新生成計算一次以前生成過的K和V。這就是KV Cache。一種空間換時間的做法。

所以可以看出，隨著序列增長，需要存儲的K和V逐漸增多，因此推理中產生的開銷也會越大。
MOE架構

MOE層替代的是傳統Transformer架構中FFN這一層。就是圖中框紅圈的Feed Forward層。

之前有寫過一篇介紹大模型推理加速的文章，
隊長：大模型推理加速，從逮蝦戶說起34 贊同 · 0 評論文章

闡述了FFN層

的幾個問題，

第一，是FFN層參數量巨大，大概占了整個模型參數量的三分之二左右。

第二，FFN層的激活存在較大的稀疏性，也就是針對某些問題的輸入，FFN層只有部分的參數是有用的。

MOE架構主要針對第二個問題進行了較大的改進。簡而言之就是將一個完整的FFN，替換成由多個“專家網絡

”組成的神經網絡，這個神經網絡可以是簡單的FFN，或者其它結構。從而在推理或者訓練時，能夠針對不同的數據進行解耦，增加效率。

從上圖的結構可以看出，MOE架構包含兩個部分：

一個是之前說的多個專家網絡組成的稀疏MOE層。
另一個是門控網絡，決定了模型的輸入，需要由哪個專家網絡進行處理。

這么做的好處在推理中體現如下，

首先，提升了模型的拓展性，MoE的每個專家都負責處理一部分數據。這意味著整個模型可以包含大量的專家（和相應的參數），而每次只有少部分專家被激活處理特定的數據。這樣，我們可以擴展模型的能力而不必擔心處理速度變慢。

只有被門控機制選中的專家才會參與到當前輸入的處理中。這不僅使模型在處理單個任務時更為高效，而且大大減少了不必要的計算，每個專家只計算與其專長相關的數據。

不過，也有隨之而來的問題，最主要的問題是，需要設計一個高效和公平的門控機制，即負載均衡問題。它需要確保正確的專家被激活，同時避免某些專家被過度使用而其他專家則幾乎閑置。

此外，還得確保每個專家都能得到足夠的訓練，以避免過擬合或欠擬合的問題，否則輸出的穩定性會受到極大影響。
DeepSeek-V2詳解

說完了背景知識，KV Cache和MOE，以及它們存在的一些問題。

我們可以來看看DeepSeek-V2對上述兩個核心部分到底做了哪些改進。
Multiple Latent Attention

MLA是對傳統多頭注意力做的改進，其目的有兩個：

降低推理過程中的KV Cache資源開銷。
緩解MQA、MGA對性能的損耗。

首先來說第一點，之前介紹KV Cache中，提到每一步都需要將K和V緩存下來。假設單個Attention Block塊中的多頭注意力，有n個頭，每個k和v的維度為d，則每一步需要緩存的參數量為

，l為block的塊數。

因此，MLA立足于在推理中，降低

。對Key和Value進行了一個低秩聯合壓縮。

簡單理解就是，假設有個矩陣的維度是
，那么可以將其分解為兩個的矩陣相乘，而

。這樣就降低了存儲量。

具體來看看DeepSeek中的具體實現公式：

是對Key和Value壓縮后的隱向量，通過一個降維映射矩陣和模型輸入得到。的維度，遠小于多頭key和value的原始維度
。
得到這個
后，具體的key和value，由兩個對應的升維矩陣和

還原。

在推理的過程中，只需要緩存每一步的
，然后再計算還原回原始的K和V即可。由于的維度遠小于K、V。因此每一步token的推理產生的緩存由之前的，變成

。

另外，之前提到KV Cache中，Q的作用只發生在當下，但是在模型訓練的過程中，每個輸入的token會通過多頭注意力機制生成對應的query、key和value。這些中間數據的維度往往非常高，因此占用的內存量也相應很大。所以論文中也提到，為了降低訓練過程中的激活內存activation memory，DeepSeek-V2還對queries進行低秩壓縮，即便這并不能降低KV Cache。

對Q的壓縮方式和K、V一致。

至此，架構圖中紅圈部分已解釋完畢。
位置編碼解耦

從架構圖中發現，DeepSeek-V2的q和k各自都有2個部分。1個部分是剛剛解釋過的壓縮部分，而另外的1個部分，加上了RoPE位置編碼

。做了一個位置編碼的解耦。

在RoPE的實現中，如果我們要讓Q、K帶上位置信息，會分別乘以相應的位置編碼矩陣。

如果計算

時，就變成了

DeepSeek-V2對Q和K都進行了壓縮，則整個過程變成：

這里，
和 分別是用于從低秩表示恢復到原始維度的解壓縮矩陣。問題在于，由于低秩表示已經是壓縮了的狀態，直接在 和 上應用 和

不再等價于在完整的Q和K上應用位置編碼。這是因為壓縮操作可能已經丟失了某些信息，使得位置編碼不能直接和有效地反映原始Q和K的位置關系。

為了解決這個問題，Deepseek-V2設計了兩個pe結尾的變量用于儲存旋轉位置編碼的信息，將信息存儲和旋轉編碼解耦開。

最后將這四個變量分別拼接起來，形成帶信息壓縮的Q、K，以及帶位置信息的Q、K，進行最后的計算。

最終，單個Token產生的緩存包含了兩個部分，即

。

同時，與GQA和MQA相比，不同于MQA和GQA可能因合并或分組而丟失細節信息，MLA的壓縮是基于保持盡可能多的原始信息的前提下進行的。這使得模型在執行注意力操作時，能夠利用到更精確的信息，從而提高整體的性能。
DeepSeek MOE

DeepSeekMoE引入了兩個主要策略：

細粒度專家分割（Fine-Grained Expert Segmentation）：通過將每個FFN專家進一步細分，這允許模型在保持參數總數不變的情況下，激活更多的、更細粒度的專家。這種策略使得各個專家能夠專注于更細致的知識領域，提高了專家的專業化程度。
共享專家隔離（Shared Expert Isolation）：設置一部分專家作為“共享專家”，這些專家總是被激活，用于捕捉和整合常見的跨上下文知識。這樣可以減少路由專家之間的知識冗余，每個路由專家可以更專注于獨特的知識領域。

上圖摘取自DeepSeek MOE的技術報告，從左到右分別是傳統的MOE，細粒度專家分割，細粒度專家分割+共享專家隔離。

這兩個策略是對緩解傳統MoE的缺陷進行的改進，

知識雜糅（Knowledge Hybridity）：傳統的MoE模型中，每個專家往往需要處理多種類型的知識，這使得專家難以形成專門化的知識結構。
知識冗余（Knowledge Redundancy）：不同的專家在處理不同的輸入時可能需要相同的知識，導致多個專家中存在重復的知識，浪費了模型參數。

針對上面的缺陷，業界其實有挺多人做過實驗，例如這篇文章，

作者對Mixtral-8x7B進行了實驗，想知道模型中不同的專家是否貢獻相同。于是作者逐一去掉了網絡中的專家，再對模型進行測試。最后發現expert 3的貢獻遠大于其它專家。

呈現出了一人干活，旁人看戲的特點。

顯然，Mixtral-8x7B可能負載均衡沒做好，導致了這個問題。

那么Deepseek MoE是否也有類似的問題呢？

從技術報告的說明來看，似乎是沒有，因為報告里專門寫了一下，當只有4個路由專家被激活時，DeepSeekMoE 的表現（Pile 損失）也與基于GShard架構的同等參數量模型相當。