MixAttention：跨層 KV Cache 共享 + 滑動窗口 Attention
一、背景
我們之前的文章中介紹過 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 種優(yōu)化 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。?
2.2 本文方案
如下圖 Figure 2 所示為本文 MixAttention 與標準 Transformer Attention 以及 Sliding Window Attention 的區(qū)別。基本與上述的 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。?
作者也探索了更多連續(xù)層共享 Global KV Cache 的方案，以 MA-Successive 為前綴，如下圖 Figure 9 所示：
除此之外，作者還探索了沒有共享 Global KV Cache 的方案，以 MA-NoShare 為前綴，如下圖所示：
三、實驗和結果
3.1 訓練
訓練分為 3 個階段：
- Stage 1：101B Token 預訓練，Max Sequence Length 為 4K，RoPE 的 theta 為 0.5M。
- Stage 2：9B Token 自然語言和代碼數(shù)據(jù)，Max Sequence Length 擴展到 32K，RoPE 的 theta 擴展到 8M。
- Stage 3：0.5B 長文本合成數(shù)據(jù)，Max Sequence Length 依然是 32K。
3.2 評估
所有模型在前兩個 Stage 上的 Loss 都非常接近，而在 Stage 3 有較大區(qū)別。如下圖 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（非共享）在深層。
3.3 推理速度
如下圖 Figure 8 所示，作者在單個 H100 GPU 上使用 SGLang 驗證了不同模型的推理速度，使用 300 個 Prompt，輸入長度 31K，輸出長度 1K。可以看出，MA 相關的方案在速度上都有比較明顯的提升，大約 2x-3x。此外，支持的最大 Token 數(shù)目也更多，不過其中 Sliding Window Attention 還沒有優(yōu)化，所以支持的最大 Token 數(shù)和標準 LLM 相同。
PS：這里的實驗有點單薄，只在一個單一的數(shù)據(jù)場景，也沒有測試不同壓力下的性能。
3.4 總結
如下圖所示，從各種評估中可以看出本文的 MA-Offset 和 MA-Pairs 在推理速度，長短文本任務上都獲得了不錯的結果，而標準的 MA 在長文本任務上性能較差。
四、參考鏈接
本文轉載自 ??AI閑談??，作者： AI閑談
