MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention

amei2000go

發布于 2024-10-8 16:03

瀏覽

0收藏

一、背景

我們之前的文章中介紹過 Character.AI 的 LLM 推理最佳實踐，其在 1 年多的時間里將推理成本降低了 33 倍。其中一個關鍵技術是對 KV Cache 的跨層共享以及與 Local Attention 的結合。本文我們介紹 MixAttention，其思路和上述方案完全一致，不過針對長文本場景做了更多實驗和調整。

對應的論文為：[2409.15012] Inference-Friendly Models With MixAttention

LLM 稀疏化相關工作可以參考：

??SnapKV: KV Cache 稀疏化，零微調加速長序列 LLM 推理??
??TriForce：KV Cache 稀疏化+投機采樣，2.3x LLM 無損加速??
??33 倍 LLM 推理性能提升：Character.AI 的最佳實踐??
??微軟 MInference：百萬 Token 序列，10x 加速??
??MLKV：跨層 KV Cache 共享，降低內存占用??
???MiniCache 和 PyramidInfer 等 6 種優化 LLM KV Cache 的最新工作???

二、方案

2.1 Character.AI 方案

如下圖所示為 Character.AI 的方案，左側為標準的 Transformer Layer，全部是 Global Attentio；右側為 Character.AI 的方案，結合了跨層 KV Cache 共享和 Sliding Window Attention：

藍色的 1,7,13 使用 Global Attention，并且 7 和 13 共享 1 的 KV Cache。
綠色的 2,4,8,10 和紅色的 3,5,6,9,11,12 使用 Local Attention，并且紅色的 3 會共享綠色2 的 KV Cache，紅色的 5 和 6 會共享綠色4 的 KV Cache。?

MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention-AI.x社區

2.2 本文方案

如下圖 Figure 2 所示為本文 MixAttention 與標準 Transformer Attention 以及 Sliding Window Attention 的區別。基本與上述的 Character.AI 的方案一致，只不過共享的位置不太一樣。其中紅點表示被共享的 Global Attention，藍點表示被共享的 Sliding Window Attention。

MA：與 Character.AI 方案一致。
MA-Offset：起始的幾個 Layer 先使用 Sliding Window Attention，關注局部；然后才會有 Global Attention。
MA-EndSlide：和 MA-Offset 相反，在結束的 Layer 也采用 Sliding Window Attention。如下圖 Figure 3 所示。這個主要是為了評估最后一層 Global Attention 對長序列的影響有多大。
MA-Pairs：Global Attention 也采用 Pair 的方式。在 MA 和 MA-Offset 只會有一層的全局 KV Cache，在 MA-Pairs 中會有多層的全局 KV Cache。?

MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention-AI.x社區

作者也探索了更多連續層共享 Global KV Cache 的方案，以 MA-Successive 為前綴，如下圖 Figure 9 所示：

MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention-AI.x社區

除此之外，作者還探索了沒有共享 Global KV Cache 的方案，以 MA-NoShare 為前綴，如下圖所示：

MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention-AI.x社區

三、實驗和結果

3.1 訓練

訓練分為 3 個階段：

Stage 1：101B Token 預訓練，Max Sequence Length 為 4K，RoPE 的 theta 為 0.5M。
Stage 2：9B Token 自然語言和代碼數據，Max Sequence Length 擴展到 32K，RoPE 的 theta 擴展到 8M。
Stage 3：0.5B 長文本合成數據，Max Sequence Length 依然是 32K。

3.2 評估

所有模型在前兩個 Stage 上的 Loss 都非常接近，而在 Stage 3 有較大區別。如下圖 Figure 4 所示，MA、Sliding Window Attention 和 MA-EndSlide 的效果明顯差于其他模型，在長文本 RULER 評估上也有類似的結論。作者也分析了相關原因，MA 和 MA-EndSlide 的 Global Attention KV Cache（非共享）都是在第 1 層，而 MA-Offset 和 MA-Pairs 至少有一個 Global Attention KV Cache（非共享）在深層。

MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention-AI.x社區

3.3 推理速度

如下圖 Figure 8 所示，作者在單個 H100 GPU 上使用 SGLang 驗證了不同模型的推理速度，使用 300 個 Prompt，輸入長度 31K，輸出長度 1K。可以看出，MA 相關的方案在速度上都有比較明顯的提升，大約 2x-3x。此外，支持的最大 Token 數目也更多，不過其中 Sliding Window Attention 還沒有優化，所以支持的最大 Token 數和標準 LLM 相同。

MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention-AI.x社區