小米給KV Cache減負80%！MiMo團隊推出混合稀疏注意力架構

小米Mimo大模型團隊投稿
量子位 | 公眾號 QbitAI

小米MiMo大模型團隊，加入AI拜年戰場——

推出HySparse，一種面向Agent時代的混合稀疏注意力架構

HySparse創新使用極少的全注意力（Full Attention）層提供“token選擇+KV Cache”，其余稀疏注意力（Sparse Attention）層直接復用這些信息，實現高效精準的長上下文建模。

在總共49層的80B-A3BMoE模型實驗中，僅保留5層Full Attention仍能保持甚至提升模型能力，同時顯著降低KVCache存儲與計算開銷，實現效果與效率的兼顧，展示出混合稀疏注意力在超長上下文建模中的巨大潛力。

HySparse的設計靈感來源于學術界已有研究工作的經驗和觀察之上。

一部分是顯著token在相鄰層之間相對穩定。

已有工作如TidalDecode等，觀察到連續層的 “重要 token” 會高度重合，因此可以在某層識別重要token并在后續層復用。

HySparse將這一觀察提升用于模型結構設計并直接訓練。

還有部分受啟發于跨層KV Cache共享能顯著省顯存且不顯著傷性能，YOCO、Gemma3n等架構層面工作已經證明了跨層共享KV的可行性。

HySparse將共享直接落在“Full Attention層 → 其后Sparse Attention層” 的hybrid block內共享上。

背景：Sparse Attention的 “兩朵烏云”

多數Sparse Attention方法都遵循同一基本范式：先選擇，再計算。

在真正計算注意力之前，先用一個“選擇器” 決定哪些重要token是值得注意的，只在這些位置上進行注意力計算。

這一范式降低了計算量，但始終繞不開兩個根本問題。

一個問題是重要 Token 的選擇依賴代理（Proxy）

Sparse Attention的核心在于 “選哪些重要token”。

但現實中，這一選擇往往依賴proxy信號：固定模式、啟發式規則、近似估計，或額外的輕量化選擇模塊。

這些proxy本質上都是對真實注意力分布的近似，無法保證能完全精準地識別最重要的token；在長上下文、持續生成的場景下，這種近似誤差還會不可避免地累積。

即便是可訓練的Sparse Attention方法，往往也只是將人工設計的選擇規則替換為可學習的輕量化選擇代理模塊，通過訓練來對齊選擇行為。

這在一定程度上緩解了選擇誤差，但同時顯著增加了訓練復雜度，難以從根本上擺脫這一“代理瓶頸”。

第二個問題是計算量降低，但KV Cache存儲未減

目前，主流的動態Sparse Attention注意力方法主要減少的是計算開銷。

這類做法不再對所有token執行注意力計算；但為了避免在生成過程中不可逆地錯誤丟棄KV Cache（因為token的重要性會隨生成動態變化），推理階段通常仍需保留全量KV Cache。

帶來的結果就是計算量確實下降了，但顯存與帶寬這一主要瓶頸依然存在。

HySparse：把 “選擇” 和 “緩存”交給Full Attention

HySparse做的事是把 “選擇” 和 “緩存” 這兩件對Sparse Attention來說最難的事，交給Full Attention來做

為此，HySparse采用hybrid block結構。

每個hybrid block由1層Full Attention+N層Sparse Attention組成。

Hybrid block內部的Sparse Attention層并不再獨立做token選擇和維護全量KV，而是直接復用前置Full Attention層產生的重要token索引和 KV Cache。

這背后的動機是Full Attention在完成自身計算的同時，已經生成了KV Cache，并且計算出了最準確的token重要性信息，自然可以供后續N個Sparse Attention層直接復用。

這一設計表面上非常簡潔，卻巧妙地同時解決了上述 Sparse Attention 的兩大核心問題，

• 選擇不再依賴proxy；
• Sparse層不引入額外KV Cache開銷。

同時，HySparse稀疏層不是單一路徑，稀疏層內部也做了一次 “全局稀疏 + 局部窗口” 的混合結構。

HySparse 的每一層 Sparse Attention 包含兩條分支：

• 塊級稀疏注意力分支（全局）：在TopK索引對應的共享KV Cache上進行全局Sparse Attention計算；
• 滑動窗口注意力分支（局部）：維護一個很小的本地窗口KV Cache（默認窗口大小為128）以保證局部建模能力。

兩分支輸出通過輕量門控（sigmoid gate）進行融合。

直觀地看，HySparse并不是用Sparse Attention取代Full Attention ，而是將全局信息通路拆解為 “少量昂貴但可靠的全注意力”+“多次廉價而高效的全局稀疏檢索與局部建模”。

實驗結果：整體更優

為了驗證HySparse的有效性，研究團隊將其與兩類主流架構進行對比：

• Full-Attn：所有層均為Full Attention；
• Hybrid SWA：按比例混合Full Attention和Sliding Window Attention；
• HySparse：采用與Hybrid SWA相同的Full Attention比例，但將滑動窗口層升級為“全局稀疏 + 局部窗口”的Sparse Attention。

研究人員分別在7B的Dense模型和80B的MoE模型上進行了實驗。

結果顯示，7B模型共有36層，其中9層是Full Attention；80B MoE模型共有49層，只有5層是Full Attention，這為HySparse帶來了接近10×的KV Cache存儲降低。

在多項通用、數學、代碼和中文評測中，HySparse在7B Dense和80B MoE兩種規模均帶來穩定提升。

特別是在80B MoE的激進混合比例下，HySparse通過全局稀疏檢索通路有效保持長距離信息訪問，維持了全局推理能力，還在部分任務上超越了全注意力基線，充分體現了全局稀疏通路的重要作用。

RULER長文測試表明，HySparse即便將Full Attention層壓到極少，也能穩定保持長距離關鍵信息訪問，充分展現了其混合稀疏結構的優勢。

小米MiMo表示，計劃在更大規模模型上進一步驗證HySparse的極限和潛力，并持續探索降低Full Attention層數量的可能性，讓超長上下文更高效。

arXiv鏈接：
https://arxiv.org/abs/2602.03560