一、背景

LLM 中 KV Cache 占據的顯存越來越大，有很多方案開始嘗試跨層共享 K 和 V，比如我們之前介紹的 YOCO、CLA 以及 Layer-Condensed KV Cache 等，本文介紹的方案也極其類似。

對應的論文為：[2406.09297] MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding

對應的代碼庫為：??https://github.com/zaydzuhri/pythia-mlkv??

PS：感覺本文創新度明顯不足，相關實驗也非常少，只在一個 160M 模型測試，甚至沒有測試 7B 模型。

二、摘要

Transformer 模型的自回歸推理因為 KV Cache 的存在可以大幅降低計算量，但隨著模型、Batch Size 以及序列長度的增長，KV Cache 大幅增加，導致可能存在內存瓶頸。本文中，作者引入了多層 KV（Multi-Layer Key-Value，MLKV）Cache，可以跨 Transformer Layer 實現 KV Cache 共享，以減少內存占用，甚至可以比 MQA 和 GQA 節約更多的內存占用。作者使用經過訓練的 Pythia-160M 變體，針對各種 NLP 基準和推理能力的指標進行評估，表明 MLKV 可以以最小的性能損失顯著降低內存使用量（???），與 MQA 相比，可以將 KV Cache 大小減少 6 倍。這些結果凸顯了 MLKV 在部署大規模 LLM 模型方面的潛力。 

三、方法

如下圖 Figure 2 所示，其思路很簡單，也和我們之前介紹過的幾個工作很類似，主要區別如下：

• MHA：原始的 Multi Head Attention，每一層的每一個 Head 都有獨立的 K 和 V。
• MQA：Multi Query Attention，每一層的所有 Head 共享 K 和 V.
• GQA：Grouped Query Attention，MHA 和 MQA 的折衷，每一層的 Head 分為多組，每一組共享 K 和 V.
• MLKV：多個層共享 K 和 V，并且可以與上述 MQA 和 GQA 兼容。?

如下圖 Table 2 所示為不同配置下總共 KV Head 的個數，參數量，以及 Loss：

四、結果

如下圖所示為不同配置下在各種評估任務上的結果，可以看出在同等配置下是弱于 GQA 的，甚至弱于 MQA：

如下圖是相應的顯存占用，同樣 Head 數的方案內存占用相同，Head 越少，內存占用越少：

如下圖 Figure 5 所示，同樣 Head 下 MLKV 的速度會更快一些，不過差距都不大：

四、參考鏈接

1. [2406.09297] MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding
2. ???https://github.com/zaydzuhri/pythia-mlkv???