一、背景

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

對應(yīng)的論文為：[2409.15012] Inference-Friendly Models With MixAttention

LLM 稀疏化相關(guān)工作可以參考：

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

二、方案

2.1 Character.AI 方案

如下圖所示為 Character.AI 的方案，左側(cè)為標(biāo)準(zhǔn)的 Transformer Layer，全部是 Global Attentio；右側(cè)為 Character.AI 的方案，結(jié)合了跨層 KV Cache 共享和 Sliding Window Attention：

• 藍(lán)色的 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 與標(biāo)準(zhǔn) Transformer Attention 以及 Sliding Window Attention 的區(qū)別。基本與上述的 Character.AI 的方案一致，只不過共享的位置不太一樣。其中紅點(diǎn)表示被共享的 Global Attention，藍(lán)點(diǎn)表示被共享的 Sliding Window Attention。

• MA：與 Character.AI 方案一致。
• MA-Offset：起始的幾個(gè) Layer 先使用 Sliding Window Attention，關(guān)注局部；然后才會有 Global Attention。
• MA-EndSlide：和 MA-Offset 相反，在結(jié)束的 Layer 也采用 Sliding Window Attention。如下圖 Figure 3 所示。這個(gè)主要是為了評估最后一層 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 為前綴，如下圖所示：

三、實(shí)驗(yàn)和結(jié)果

3.1 訓(xùn)練

訓(xùn)練分為 3 個(gè)階段：

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

3.2 評估

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

3.3 推理速度

如下圖 Figure 8 所示，作者在單個(gè) H100 GPU 上使用 SGLang 驗(yàn)證了不同模型的推理速度，使用 300 個(gè) Prompt，輸入長度 31K，輸出長度 1K?？梢钥闯?，MA 相關(guān)的方案在速度上都有比較明顯的提升，大約 2x-3x。此外，支持的最大 Token 數(shù)目也更多，不過其中 Sliding Window Attention 還沒有優(yōu)化，所以支持的最大 Token 數(shù)和標(biāo)準(zhǔn) LLM 相同。

PS：這里的實(shí)驗(yàn)有點(diǎn)單薄，只在一個(gè)單一的數(shù)據(jù)場景，也沒有測試不同壓力下的性能。

3.4 總結(jié)

如下圖所示，從各種評估中可以看出本文的 MA-Offset 和 MA-Pairs 在推理速度，長短文本任務(wù)上都獲得了不錯(cuò)的結(jié)果，而標(biāo)準(zhǔn)的 MA 在長文本任務(wù)上性能較差。

四、參考鏈接

1. ???https://arxiv.org/abs/2409.15012????

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