一、背景

本文中我們介紹一種最新的 KV Cache 共享論文 KVSharer，與之前常見(jiàn)的層內(nèi)共享不同，KVSharer 主要關(guān)注跨層共享，并且是整個(gè)層的共享。

對(duì)應(yīng)的論文：[2410.18517] KVSharer: Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing

對(duì)應(yīng)的代碼庫(kù)：????https://github.com/yangyifei729/KVSharer/tree/main??

二、摘要

LLM 推理過(guò)程中對(duì) GPU 內(nèi)存的需求不斷增加，而其中 Attention 的 KV Cache 占據(jù)了超過(guò) 80% 的空間。當(dāng)前，大多數(shù)現(xiàn)有的 KV Cache 壓縮算法主要集中在單個(gè) Transformer 層內(nèi)的壓縮，而較少有研究涉及層間壓縮（PS：其實(shí)我們之前介紹過(guò)的 Character.AI 方案和 MixAttention 都是層間共享）。

本文中，作者提供了一種即插即用的方案，稱(chēng)為 KVSharer，通過(guò)在層間共享 KV Cache 來(lái)實(shí)現(xiàn)層間壓縮。作者發(fā)現(xiàn)一個(gè)反直覺(jué)的現(xiàn)象：共享相似性較低的 KV Cache 更能保持模型性能（PS：這里實(shí)現(xiàn)時(shí)其實(shí)有個(gè)約束，要同時(shí)避免對(duì) LM head 之前的 Hidden State 的影響過(guò)大）。實(shí)驗(yàn)表明，KVSharer 能夠在減少 30% 的 KV Cache 計(jì)算，以及內(nèi)存開(kāi)銷(xiāo)，而對(duì)模型性能影響不大，同時(shí)還能實(shí)現(xiàn)至少 1.3x 的生成加速。此外，作者實(shí)驗(yàn)表明，KVSharer 與現(xiàn)有的層內(nèi) KV Cache 壓縮方案兼容，兩者結(jié)合可進(jìn)一步節(jié)約內(nèi)存。

三、方案

3.1 概述

如下圖 Figure 2 所示為本文方案的示例，其主要包含幾個(gè)步驟：

• 校正集上推理，并記錄 KV Cache。
• 計(jì)算任意兩層的歐式距離。
• 排序，按相似性從低到高排列。
• 從排序的組合中依次篩選可以共享的層。
• 在整個(gè)后續(xù)推理中保持離線(xiàn)搜索的共享策略。?

3.2 搜索策略

如下圖 Algorithm 1 為詳細(xì)的搜索算法，其大體包含兩個(gè)階段：

KV Cache 相似性計(jì)算和初始化（1-4）：

• 首先使用校準(zhǔn)集所有樣本進(jìn)行推理，獲得所有 Layer 的 KV Cache。
• 在每一層，對(duì)所有樣本的 KV Cache 求平均。
• 將每一層的 KV Cache 拉平為一維 Embedding。
• 計(jì)算任意兩層 KV Cache Embedding 的歐式距離（距離越大越不相似），構(gòu)成 S。
• 按照距離降序排列（距離越大，越不相似），得到 R。

共享策略搜索（5-18）：

• 初始空的共享策略 Z 和空的共享 Layer 個(gè)數(shù) P。
• 依次遍歷排序后的候選共享層 R：

將當(dāng)前候選 r 添加到共享策略 Z。

替換候選共享策略 Z 中的所有相關(guān)層（保留靠近輸入的層，靠近輸出的層共享靠近輸入的層），然后使用校準(zhǔn)集驗(yàn)證最后一層 Hidden Stage 和未替換時(shí)的相似度。

如果相似度小于閾值 T，表示替換后影響較大，則從 Z 中刪除當(dāng)前后續(xù) r。

如果相似度大于等于閾值 T，則可以作為候選。

如果候選集容量滿(mǎn)足 C，則直接結(jié)束。?

PS：需要說(shuō)明的是，以上是基于貪心搜索策略，這種逐步判斷并添加的方式往往只能找到局部最優(yōu)解，并不一定是全局最優(yōu)的。要想找到全局最優(yōu)，通常需要進(jìn)行全局搜索，比如窮舉所有可能的共享組合，或者產(chǎn)生更復(fù)雜的搜索算法，但是往往會(huì)大幅增加計(jì)算成本，因此這種貪心搜索是計(jì)算效率和效果的折衷。

3.3 推理

如下圖 Figure 3 所示，在推理階段直接使用離線(xiàn)階段搜索出來(lái)的共享策略，可以看出其實(shí)是會(huì)存在多個(gè) Layer 共享一個(gè) Layer 的。但是圖中似乎有問(wèn)題，論文中并沒(méi)有討論共享的傳遞性，比如Layer 2 和 3 共享 Layer 1，而 Layer 4 共享 Layer 3，Layer 5 共享 Layer 4，那豈不是 Layer 2，3，4，5 都共享 Layer 1？

如下圖所示為代碼中的注釋?zhuān)↘VSharer/llama_real_share/modeling_llama_kvsharer.py#L364-L371）：

四、實(shí)驗(yàn)&結(jié)果

4.1 精度

如下圖 Table 1 所示，作者測(cè)試幾個(gè)常見(jiàn)模型 LLaMA2-7B/13B（PS：如果有 LLaMA3 的結(jié)果更有說(shuō)服力）、InternLM2-7B/20B（中英文能力不錯(cuò)）和 Mistral-7B 在一些主要任務(wù)上的精度。其中 Layer 表示實(shí)際計(jì)算的層數(shù)，Layer 越小，表示共享的越多。Percent 表示共享后的平均精度相比原始平均精度的比例?？梢钥闯觯蚕碓蕉嘈Ч讲?，并且各個(gè)模型表現(xiàn)各不相同，比如在 LLaMA2-13B（40->30） 和 Mistral-7B（32->24） 在共享 25% 時(shí)，平均精度就不到 90% 了。

其實(shí)，從困惑度（PPL）上來(lái)看，影響也是挺大的，尤其是與 H2O 和 PyramidInfer 這些層內(nèi)壓縮結(jié)合后：

4.2 推理

如下圖 Table 2 所示，作者以 LLaMA2-13B-Chat 模型為例，對(duì)比了不同方案對(duì) Memory 的開(kāi)銷(xiāo)以及推理的加速比。其中的 KVSharer（25%）表示 25% 的壓縮比，也就意味著 40 層只計(jì)算 30 層的 KV Cache。（PS：這里的結(jié)論很奇怪，25% 的壓縮比理論上最多節(jié)約 25% 內(nèi)存，最多吞吐提升 25%？更何況 Q 和 Attention 還要正常計(jì)算，而實(shí)際上 1024+4096 時(shí)節(jié)約了 36% 的內(nèi)存，吞吐提升 1.53x。）

4.3 消融實(shí)驗(yàn)

如下圖 Figure 6 所示，作者實(shí)驗(yàn)表明：共享相似性較低的 KV Cache 更能保持模型性能。

如下圖 Table 3 所示，使用 Wikipedia 或 BookCorpus 作為校準(zhǔn)集對(duì)精度的影響不是特別大：

如下圖 Table 4 所示，隨機(jī)共享相比本文的 KVSharer 會(huì)導(dǎo)致精度下降比較多，證明本文方法的有效性：

如下圖 Table 5 所示，KVSharer 在幾個(gè)模型的 Base 模型和 Chat 模型上的精度影響比較類(lèi)似，證明了方法的通用性。（PS：不過(guò) PPL 確實(shí)影響挺大的）

五、參考鏈接

• ??https://arxiv.org/abs/2410.18517??
• ??https://github.com/yangyifei729/KVSharer/tree/main???????
• https://github.com/yangyifei729/KVSharer/blob/main/llama_real_share/modeling_llama_kvsharer.py#L364-L371

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談