MLKV：跨層 KV Cache 共享，降低內(nèi)存占用
一、背景
LLM 中 KV Cache 占據(jù)的顯存越來越大，有很多方案開始嘗試跨層共享 K 和 V，比如我們之前介紹的 YOCO、CLA 以及 Layer-Condensed KV Cache 等，本文介紹的方案也極其類似。
對應(yīng)的論文為：[2406.09297] MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding
對應(yīng)的代碼庫為：??https://github.com/zaydzuhri/pythia-mlkv??
PS：感覺本文創(chuàng)新度明顯不足，相關(guān)實驗也非常少，只在一個 160M 模型測試，甚至沒有測試 7B 模型。
二、摘要
Transformer 模型的自回歸推理因為 KV Cache 的存在可以大幅降低計算量，但隨著模型、Batch Size 以及序列長度的增長，KV Cache 大幅增加，導(dǎo)致可能存在內(nèi)存瓶頸。本文中，作者引入了多層 KV（Multi-Layer Key-Value，MLKV）Cache，可以跨 Transformer Layer 實現(xiàn) KV Cache 共享，以減少內(nèi)存占用，甚至可以比 MQA 和 GQA 節(jié)約更多的內(nèi)存占用。作者使用經(jīng)過訓(xùn)練的 Pythia-160M 變體，針對各種 NLP 基準和推理能力的指標(biāo)進行評估，表明 MLKV 可以以最小的性能損失顯著降低內(nèi)存使用量（???），與 MQA 相比，可以將 KV Cache 大小減少 6 倍。這些結(jié)果凸顯了 MLKV 在部署大規(guī)模 LLM 模型方面的潛力。
三、方法
如下圖 Figure 2 所示，其思路很簡單，也和我們之前介紹過的幾個工作很類似，主要區(qū)別如下：
- 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 的個數(shù)，參數(shù)量，以及 Loss：
四、結(jié)果
如下圖所示為不同配置下在各種評估任務(wù)上的結(jié)果，可以看出在同等配置下是弱于 GQA 的，甚至弱于 MQA：
如下圖是相應(yīng)的顯存占用，同樣 Head 數(shù)的方案內(nèi)存占用相同，Head 越少，內(nèi)存占用越少：
如下圖 Figure 5 所示，同樣 Head 下 MLKV 的速度會更快一些，不過差距都不大：
四、參考鏈接
- [2406.09297] MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer Decoding
- ???https://github.com/zaydzuhri/pythia-mlkv???
