前面文章：
《Transformer深入學習（一）：Sinusoidal位置編碼的精妙》

一、MHA、MQA、GQA

為了降低KV cache，MQA、GQA作為MHA的變體，很容易理解。
多頭注意力（MHA）：
多頭注意力是一種在Transformer架構中廣泛使用的注意力機制，通過將查詢、鍵和值分別投影到多個不同的空間上，然后并行計算這些空間上的注意力得分，從而獲得更加豐富和細致的特征表示。

多查詢注意力（MQA）：
多查詢注意力是MHA的一種變種，它通過共享單個key和value頭來提升性能，但可能會導致質量下降和訓練不穩定。MQA在保持速度的同時提高了模型的推理效率，但在某些情況下可能無法達到與MHA相同的效果。

分組查詢注意力（GQA）：
分組查詢注意力是MQA和MHA之間的過渡方法，旨在同時保持MQA的速度和MHA的質量。GQA通過使用中間數量的鍵值頭（大于一個，小于查詢頭的數量），實現了性能和速度的平衡。具體來說，GQA通過分組的方式減少了需要處理的頭數，從而降低了內存需求和計算復雜度。

分組查詢注意力（Grouped-Query Attention，簡稱GQA）是一種用于提高大模型推理可擴展性的機制。其具體實現機制如下：

1、基本概念：GQA是多頭注意力（Multi-Head Attention，MHA）的變種，通過將查詢頭（query heads）分成多個組來減少內存帶寬的需求。每個組共享一個鍵頭（key head）和一個值頭（value head），從而降低了每個處理步驟中加載解碼器權重和注意力鍵/值的內存消耗。

2、實現方式：在實際應用中，GQA將查詢頭分成G個組，每組共享一個鍵頭和一個值頭。例如，GQA-G表示有G個組，而GQA-1則表示只有一個組，這相當于傳統的MQA（Multi-Group Query Attention）。當GQA的組數等于查詢頭的數量時，它等同于標準的MHA。

3、性能與效率平衡：GQA通過這種方式有效地平衡了性能和內存需求。它允許模型在不顯著降低性能的情況下，處理更多的請求并提高推理效率。此外，使用GQA可以避免由于加載大量解碼器權重和注意力鍵/值而導致的內存瓶頸問題

二、MLA

2.1 基礎原理

這張圖，對從MHA、MQA、GQA到MLA，看的很清楚。
GQA就是用了多組KV Cahe，MQA只用了一組KV Cache。
那MLA呢？MLA看起來和MHA是一樣的，只不過存的壓縮后的隱KV，在計算的時候再通過投影倒多個KV參與注意力計算。
為什么會節省KV Cache？蘇神的文章解釋的很清楚。
看下面的公式，MLA公式如下：

其中的c就是壓縮后的隱KV。
但是這樣好像無法節省KV Cache，因為計算和MHA一樣了，關鍵在于下面的轉換公式：

這個公式把注意力的計算做了轉換，k的投影矩陣這樣就可以合并倒q的投影矩陣中。
另外，因為注意力之后的o還有一個投影矩陣，也可以合并到后面的投影矩陣中。
而c作為壓縮后的隱KV，是所有頭共享的，這樣就實現了內存的節省。

2.2 增加RoPE

但是，如上面，矩陣合并之后，就和RoPE不兼容了，具體看蘇神的分析文章。
MLA采取了一種混合的方法——每個 Attention Head的 Q、K 新增 dr個維度用來添加 RoPE，其中 K 新增的維度每個 Head 共享：

因為dr遠小于dk，所以增加的內存空間不大。

2.3 最后的版本

MLA 的最終版本，還將 Q 的輸入也改為了低秩投影形式，可以減少訓練期間參數量和相應的梯度的顯存。

MLA這種方法，在訓練階段還是照常進行，此時優化空間不大；在推理階段，應該可以大幅減少顯存。
見蘇神的分析：“ MLA 在推理階段做的這個轉換，雖然能有效減少 KV Cache，但其推理的計算量是增加的。
那為什么還能提高推理效率呢？這又回到“瓶頸”一節所討論的問題了，我們可以將 LLM 的推理分兩部分：第一個 Token 的生成（Prefill）和后續每個 Token 的生成（Generation）。
Prefill 階段涉及到對輸入所有 Token 的并行計算，然后把對應的 KV Cache 存下來，這部分對于計算、帶寬和顯存都是瓶頸，MLA 雖然增大了計算量，但 KV Cache 的減少也降低了顯存和帶寬的壓力，大家半斤八兩；但是 Generation 階段由于每步只計算一個 Token，實際上它更多的是帶寬瓶頸和顯存瓶頸，因此 MLA 的引入理論上能明顯提高 Generation 的速度。”

三、參考文章

蘇神：《緩存與效果的極限拉扯：從MHA、MQA、GQA到MLA》
https://mp.weixin.qq.com/s/yCczYU0po0PvPTa-eh2pfg

《大模型KV Cache節省神器MLA學習筆記》
https://mp.weixin.qq.com/s/cBMrRUdM1IM0T1ji_ODxng

《注意力機制的變體之MLA》
https://mp.weixin.qq.com/s/dWZk8TBY89re207ZL3GjfA