DeepSeek 透過工程手段實現了 GPU 的良好利用，這一做法似乎引起了全球 AI 研究人員的共鳴。
近日，來自韓國科學技術院和 DeepAuto.ai 公司的研究人員，研發出一款名為 InfiniteHiP 的大模型推理框架。
InfiniteHiP 能在推理過程中將鍵值（KV，Key – Value）快取置於主機記憶體，從而能夠顯著降低 GPU 記憶體壓力。這一框架能在單個英偉達 L40s 48GB GPU 上處理多達 300 萬個 tokens，並且不會永久丟失上下文資訊。利用這一框架，研究人員在 100 萬個 tokens 上下文中實現了 18.95 倍的注意力解碼加速，並且完全無需額外訓練。
同時，這一框架還能高效地利用長上下文，透過模組化分層 tokens 修剪演算法，能夠動態地消除不相關的上下文 tokens，從而能夠提高模型效率。
（來源：arXiv）
總的來說，InfiniteHiP 讓預訓練大模型能夠實現超長的上下文推理。在不降低質量的前提下，能讓模型超越了其原始上下文長度限制，同時可以透過高效的鍵值快取解除安裝克服 GPU 的記憶體限制。
作為一種無需訓練的解決方案，對於任何基於 Transformer 的預訓練大模型來說，InfiniteHiP 都能充當一種直接的替代方案，因為它不僅可以在模型層面和硬體層面提供更快的推理速度，還能夠延長上下文的可用長度。
從赫赫有名的注意力機制說起
對於基於 Transformer 的生成式大模型來說，擴充套件上下文長度至關重要。但是，對於大模型來說，處理非常長的上下文視窗卻是一個重大挑戰。一方面這會導致推理速度減慢，另一方面這會導致記憶體成本增加。此外，對於大多數已有的預訓練大模型來說，它們往往無法超越其原始訓練序列長度。
之所以存在這一難題，得從注意力機制說起。注意力機制——是模型的基本組成部分。注意力機制能夠計算每個輸入 token 與所有之前 token 之間的關係，隨著輸入序列長度的增加，這會導致計算成本和記憶體成本呈顯出指數級增長。
鍵值快取，則是注意力機制引起的另一個問題。在模型生成內容的過程中，此前計算的注意力鍵和值，會快取在 GPU 記憶體之中，以供在必要時進行重用。但是，鍵值快取大小會隨著上下文長度出現線性變化，從而會給長上下文推理帶來挑戰。
此前，人們已經提出各種方法來降低注意力機制的高成本。在這些方法之中，有一種名為 FlashAttention（FA2）的方法，它是一種用於加速 Transformer 模型注意力計算的技術，透過避免將整個注意力得分矩陣寫入全域性 GPU 記憶體，以此來降低記憶體消耗和頻寬利用率。然而，FA2 並沒有降低計算成本。
為了減輕鍵值快取的記憶體負擔，業內人士做出了不少努力。比如，一種名為“KV 快取驅逐”的方法，透過選擇性地“遺忘”此前上下文來節省 GPU 記憶體。
然而，這些方法都會永久地刪除此前上下文，而這些上下文很有可能還會被用到。
此前，有研究人員提出了 HiP 注意力機制，它能將不常訪問的“冷 token”解除安裝到更大、更便宜的主機記憶體中，僅在需要使用時將它們提取回 GPU，同時只將那些需要經常訪問的“熱 token”保留在 GPU 上。
儘管上述最佳化起到一定作用，但是上下文擴充套件的所面臨的另一個問題仍然存在：即預訓練大模型無法處理超過其訓練上下文長度的輸入。
由於注意力機制是排列不變的，於是人們開始利用旋轉位置嵌入（RoPE，Rotary Positional Embeddings）等位置嵌入方法來模擬 tokens 的時間順序。
然而，由於大模型通常在固定長度的序列上進行預訓練，因此當遇到更長的上下文提示時，它們無法適應那些沒有見過的位置。
為了解決這一問題，人們採用了長上下文微調的方法，即在一組較長的輸入上針對模型進行微調。然而，進行微調特別是針對長序列進行微調，既需要高昂的訓練成本，也需要高質量的訓練資料。
正因此，超長泛化（OOL，out-of-length）能力變得越來越重要。超長泛化指的是：在沒有訓練的情況下，預訓練模型的能力遠遠超出其預訓練極限。而作為一款長上下文大模型框架，InfiniteHiP 綜合了上述所有方法的優點。
實現 7.24 倍的端到端解碼加速
具體來說，為了減輕注意力的計算負擔，研究人員基於 InfiniteHiP 提出一種模組化稀疏注意力新方案，該方案能夠最大限度地減少不太重要的上下文的計算量。
為了最佳化鍵值快取解除安裝，InfiniteHiP 透過基於 LRU 的快取策略，增強了 HiP 注意力的解除安裝策略。（注：LRU 是 Least Recently Used 的縮寫，中文意思為“最近最少使用”，這是一種常用的置換演算法。）
基於大模型內部的注意力模式，InfiniteHiP 能夠在不同元件之內，透過仔細應用各種 RoPE 調整策略來實現超長泛化。
InfiniteHiP 的另一個獨特之處，在於它對修剪模組的創新使用。對於修剪模組來說，它能夠透過採用一種新穎的模組化分層修剪演算法，以便有選擇性地丟棄不太重要的輸入 token。
每個剪枝模組都能將輸入序列劃分為長度為固定值的塊，同時也能並行識別出每個塊中注意力得分最高的近似 token。而只有最重要的前 K 個塊（其中 K 為常數）會被傳遞到下一個模組，其餘的則會被丟棄。
透過堆疊多個剪枝模組，InfiniteHiP 能夠針對塊稀疏注意力掩碼加以迭代最佳化。（注：塊稀疏注意力掩碼的英文是 Block Sparse Attention Mask，這是一種在處理長序列資料時用於最佳化注意力機制的技術，它透過將輸入序列劃分為較小的塊（blocks），並在這些塊內應用稀疏注意力，從而減少計算複雜度和記憶體消耗。）
（來源：arXiv）
儘管 InfiniteHiP 是基於前人所提出的 HiP 打造而來，但它引入了幾個關鍵性改進：
其一，與 HiP 框架基於啟發式的分層修剪相比，InfiniteHiP 的分層修剪模組實現了更高的準確性。
其二，由於並行性得到增強，每個模組內的修剪演算法速度明顯變得更快。
其三，InfiniteHiP 採用模組化設計的方案，故能針對修剪階段的快取進行細粒度控制，從而能夠實現更快的解碼。
在 SGLang 大模型服務框架上，研究人員證明了 InfiniteHiP 的有效性。實驗結果顯示：在 3M token 上下文上，InfiniteHiP 能夠實現 7.24 倍的端到端解碼加速。
如下表所示，與表現最佳的基線大模型相比，基於 InfiniteHiP 這一框架，讓 Llama 3 的相對得分提高 7.17%，讓 Mistral 0.2 的相對得分提高 3.19%。
（來源：arXiv）
與表現最佳的基線大模型相比，InfiniteHiP 在上述兩個模型中，透過稀疏注意力所處理的金鑰 tokens 減少了四分之三，從而能夠產生更好的解碼延遲。
為了進一步展示 InfiniteHiP 的超長泛化能力，研究人員將 Bench 的英語多項選擇得分，在不同上下文長度下與 Llama 3.1 8B 模型進行對比。
（來源：arXiv）
結果發現：隨著上下文長度的增加，InfiniteHiP 的效能持續提升。而那些不具備超長泛化能力的基線模型，在超出預訓練上下文長度（128K）後，效能會顯著下降。
同時，研究人員還針對其他短上下文大型語言模型進行了實驗。在這些短上下文模型中，研究人員在擴充套件上下文環境下觀察到了十分顯著的效能提升。
例如：與 FA2 方法相比，當在 Gemma2 模型中使用 InfiniteHiP 時，其在英語多項選擇任務中得到 24.45% 以上的效能提升，在英語問答任務中得到 22.03% 以上的效能提升。
（來源：arXiv）
解決長上下文推理問題，能透過鍵值快取解除安裝節約 GPU 記憶體
與此同時，研究人員分析了 InfiniteHiP 中各個元件的延遲情況以及它們所產生的影響。
具體來說，研究人員分析了 InfiniteHiP 在 100 萬 tokens 上下文中的延遲情況，並將其與能夠產生相似基準測試分數的基線模型進行比較。
研究人員先是測量了注意力方法的延遲，在進行 100 萬 tokens 的預填充時，其發現 InfiniteHiP 分別比 FA2 快 20.29 倍，比 InfLLM 出快 6%，並且與 HiP 的延遲相近。
在進行 100 萬 tokens 上下文的解碼時，InfiniteHiP 分別比 FA2 方法快 19.85 倍、比 InfLLM 快 4.98 倍、比 HiP 快 92%。
當啟用上下文擴充套件時，針對預計算的餘弦向量和正弦向量進行額外記憶體讀取會產生開銷，而 InfiniteHiP 在預填充時的速度會減慢約 1.6 倍，在解碼時會減慢 5%。
在啟用上下文擴充套件的預填充過程中，InfiniteHiP 比 InfLLM 慢 50%，但是在解碼過程中卻明顯更快，這是因為解碼過程會受到記憶體的限制。在 3k tokens 的上下文視窗之下，InfiniteHiP 讀取的上下文 tokens 數量，比 InfLLM 在 12k tokens 的上下文視窗下讀取的要少。
與此同時，研究人員在一個密碼金鑰檢索任務樣本上，測量了啟用鍵值快取解除安裝時的解碼延遲。
在 256K 上下文解碼中，InfiniteHiP 比 InfLLM 快 3.64 倍。InfLLM 在執行其注意力核時，採用的是不訪問 CPU 記憶體的方法，因此不得不犧牲前 k 個估計算法的精度，這就需要更大的塊大小和上下文視窗大小，來維持模型在下游任務上的效能。
相比之下，InfiniteHiP 在注意力核執行期間訪問 CPU 記憶體，這為演算法設計提供了更大的靈活性，在下游自然語言理解任務表現更佳。
在下圖中，研究人員展示了使用英偉達 RTX 4090（24GB）和英偉達 L40S（48GB）GPU 的方法的解碼吞吐量。在英偉達 RTX 4090（24GB）上，與 SRT（SGlang Runtime with FlashInfer）的預估解碼吞吐量相比，研究人員的方法在 1M 上下文長度下實現了 3.20 倍的吞吐量。在英偉達 L40S（48GB）GPU 上，研究人員的方法在 3M 上下文長度下比 SRT 高出 7.25 倍。
（來源：arXiv）
在下圖中，研究人員證明 InfiniteHiP 對於重要 tokens 具有更好的覆蓋範圍，這意味著關鍵 tokens 的注意力機率具有更高的召回率（recall）。更好的召回率表明，InfiniteHiP 比基線方法能夠更緊密地遵循預訓練的注意力模式。
（來源：arXiv）
而在下圖中，研究人員透過進行消融研究，找到了每種設定的延遲效能最優修剪模組組合。
（來源：arXiv）
這些實驗均說明：InfiniteHiP 是一個無需訓練的大模型推理框架，它能夠被用於高效的長上下文推理，支援超線性泛化和動態鍵值快取解除安裝。
同時，它也有效解決了長上下文大模型推理中的三個主要挑戰：不僅解決了長上下文的高效推理的問題，也解決了超長泛化的問題，並能在不“遺忘”此前上下文的同時，透過鍵值快取解除安裝來節約 GPU 記憶體。
參考資料：
https://arxiv.org/pdf/2502.08910
運營/排版：何晨龍