KVSharer：基于不相似性實現跨層 KV Cache 共享

amei2000go

發布于 2024-10-29 11:47

瀏覽

0收藏

一、背景

本文中我們介紹一種最新的 KV Cache 共享論文 KVSharer，與之前常見的層內共享不同，KVSharer 主要關注跨層共享，并且是整個層的共享。

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區

對應的論文：[2410.18517] KVSharer: Efficient Inference via Layer-Wise Dissimilar KV Cache Sharing

對應的代碼庫：????https://github.com/yangyifei729/KVSharer/tree/main??

二、摘要

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

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

三、方案

3.1 概述

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

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

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區

3.2 搜索策略

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

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

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

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

初始空的共享策略 Z 和空的共享 Layer 個數 P。
依次遍歷排序后的候選共享層 R：

將當前候選 r 添加到共享策略 Z。

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

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

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

如果候選集容量滿足 C，則直接結束。?

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區

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

3.3 推理

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

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區

如下圖所示為代碼中的注釋（KVSharer/llama_real_share/modeling_llama_kvsharer.py#L364-L371）：

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區

四、實驗&結果

4.1 精度

如下圖 Table 1 所示，作者測試幾個常見模型 LLaMA2-7B/13B（PS：如果有 LLaMA3 的結果更有說服力）、InternLM2-7B/20B（中英文能力不錯）和 Mistral-7B 在一些主要任務上的精度。其中 Layer 表示實際計算的層數，Layer 越小，表示共享的越多。Percent 表示共享后的平均精度相比原始平均精度的比例。可以看出，共享越多效果越差，并且各個模型表現各不相同，比如在 LLaMA2-13B（40->30）和 Mistral-7B（32->24）在共享 25% 時，平均精度就不到 90% 了。

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區

其實，從困惑度（PPL）上來看，影響也是挺大的，尤其是與 H2O 和 PyramidInfer 這些層內壓縮結合后：

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區

4.2 推理

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

KVSharer：基于不相似性實現跨層 KV Cache 共享-AI.x社區