ICLR2025|告別Token丟棄：更適合CoT和多輪對話的長文字推理加速方法

近年來，大語言模型（LLMs）展現了在文件問答、長對話、複雜指令遵循等場景下的強大能力。然而，隨著上下文長度的增長，一個關鍵的瓶頸日益凸顯——KV Cache（鍵值快取）帶來的巨大 GPU 視訊記憶體開銷。

為了緩解這一問題，現有方法通常基於注意力稀疏性假設，在推理過程中丟棄（discard）它們認為不重要的 KV Cache。但這帶來了一個新的困境：注意力分數是基當前隱藏狀態計算的，無法完全預示 Token 在未來推理步驟中的重要性。

尤其在多步推理（如 CoT）或多輪對話場景下，早期被判定為不重要的 Token 可能在後續步驟中變得至關重要。丟棄 Token 可能導致關鍵資訊的永久丟失，影響模型效能。

為此，本文提出了 OmniKV，一種無需丟棄 Token、無需訓練的 LLM 高效推理方法。OmniKV 巧妙地利用了層間注意力的相似性，實現了動態的上下文選擇，在不損失效能的前提下，在 Lightllm（https://github.com/ModelTC/lightllm）上實現了 1.7x 相比於 vLLM 的吞吐量提升。

論文標題：

OmniKV: Dynamic Context Selection for Efficient Long-Context LLMs

論文連結：

https://openreview.net/forum?id=ulCAPXYXfa

程式碼連結：

https://github.com/antgroup/OmniKV.git

核心洞察：層間注意力相似性（Inter-Layer Attention Similarity）

OmniKV 的核心創新基於一個關鍵洞察：在單個生成步驟內，被模型高度關注的（即注意力得分高的）Token 集合，在不同的 Transformer 層之間表現出高度的相似性。換句話說，重要的上下文資訊在多個 transformer 層之間存在共識。

▲ 圖1a：層間注意力相似性。即使相隔多個層，重要 Token 的分佈也高度相似。

OmniKV 正是利用了這一特性。它僅選擇少數幾個層（稱為 “Filter 層”）來計算完整的注意力分數並識別重要的 Token 子集，而其他大多數層則直接複用（共享）來自最近 Filter 層識別出的 Token 索引。在計算注意力時，僅載入並計算這個稀疏子集 KV Cache，從而大幅減少計算量和資料傳輸量。

OmniKV 方法

OmniKV 的推理過程分為 Prefill 和 Decode 兩個階段：

Prefill 階段：對輸入的 Prompt 進行編碼，生成完整的 KV Cache。此時，OmniKV 將大部分非 Filter 層的 KV Cache 解除安裝（offload）到 CPU 記憶體，僅保留少量 Filter 層的 KV Cache 在 GPU 上。

Decode 階段（生成階段）：

Filter 層：計算完整注意力，並使用 Context Selector 動態選擇當前步驟最重要的 Top-K 個 Token 索引。
非 Filter 層：直接從 CPU 載入（load）由前一個 Filter 層選擇出的 Token 索引對應的 KV Cache 子集，並在該子集上執行稀疏注意力計算。
Packed Load：由於相鄰的多個非 Filter 層共享相同的 Token 索引，它們的 KV Cache 可以被打包（packed）並一次性從 CPU 載入到 GPU，進一步減少 PCIe 頻寬壓力。
非同步傳輸：資料載入與計算可以非同步進行，掩蓋部分資料傳輸延遲。

▲ 圖2：OmniKV Decode 階段框架圖。透過 Filter 層動態選擇，非 Filter 層僅載入和計算稀疏子集。

無損效能，可能更適合 CoT 或多輪對話場景的加速

▲ 圖3：Longbench 的效能表現，在長文字 benchamrk 下，OmniKV 單步推理擁有幾乎無損的效能。

丟棄 Token 的方法可能會丟失對後續推理至關重要的資訊。圖 1b 和 1c 的分析表明，在多步推理（如 CoT）中，不同生成步驟所依賴的關鍵 Token 是動態變化的。

OmniKV 的 Token-Dropping-Free 設計天然地避免了這個問題。因為它保留了所有歷史資訊，模型可以在任何需要的時候，透過動態選擇機制重新關注到之前可能被忽略的 Token。

實驗（圖4）表明，在 2WikiMQA、HotpotQA 和論文提出的 2StageRetr 等多步推理任務中，OmniKV 在各種視訊記憶體預算下都顯著優於 H2O 等丟棄 Token 的方法，展現了其在複雜推理場景下的魯棒性和優越性。

▲ 圖4：多步推理任務效能對比。OmniKV（紅色/綠色/橙色線）在各種預算下均顯著優於基線方法。

推理速度大幅提升

透過 Offloading，OmniKV 能將 Llama-3-8B 在單張 A100 上支援的最大上下文長度從 128K 擴充套件到 450K，且依然實現了 1.7x 的加速比。

推理框架適配與吞吐提升：我們將 OmniKV 成功適配到了業界流行的 LightLLM 推理框架。在多卡張量並行（TP=4）的設定下，與集成了 PagedAttention 的 vLLM 相比，OmniKV + LightLLM 的解碼吞吐量（Throughput）在處理 512K 長序列時提升了 1.7x。