這篇論文的第一作者在實習期間完成了研究。

原標題：剛剛，DeepSeek又發高質量論文，梁文鋒合著！
文章來源：智東西
內容字數：6730字

DeepSeek發布新型稀疏注意力機制NSA：長上下文訓練推理提速利器

近日，DeepSeek團隊發布論文，介紹了一種名為NSA（Novel Sparse Attention）的改進型稀疏注意力機制，該機制在長上下文訓練和推理中展現出顯著的速度提升，尤其在解碼階段最高可達11.6倍。論文第一作者Jingyang Yuan在實習期間完成了這項研究，DeepSeek創始人兼CEO梁文鋒也參與其中。

1. 現有稀疏注意力機制的缺陷與NSA的優勢

傳統的注意力機制在處理長序列時計算復雜度極高，限制了長文本建模的發展。現有的稀疏注意力機制雖然在理論上降低了計算復雜度，但在實際應用中往往未能有效降低延遲，部分原因在于它們只在推理階段或預填充階段應用稀疏性，缺乏對訓練階段的支持，且難以適應現代高效的解碼架構。NSA則旨在解決這些問題，通過在訓練和推理階段都應用稀疏性，并針對現代硬件進行優化，實現全階段加速。

2. NSA的三大核心組件及軟硬件協同優化

NSA的核心思想是通過動態分層稀疏策略，結合粗粒度token壓縮和細粒度token選擇，在提升效率的同時保留模型對全局長上下文感知能力和局部精確性。其三大核心組件包括：

1. 動態分層稀疏策略：根據不同層級的需求動態調整稀疏程度。
2. 粗粒度token壓縮：將鍵值對聚合成塊級表示，捕捉高層語義信息。
3. 精粒度token選擇：根據重要性分數選擇關鍵的細粒度信息。

此外，NSA在Triton上實現了硬件對齊的稀疏注意力內核，并針對共享KV緩存的架構（如GQA和MQA）進行了優化，通過以組為中心的數據加載、共享KV加載和網格循環調度等策略，實現了近乎最優的計算強度平衡。

3. NSA的性能表現：訓練提速6-9倍，推理提速高達11.6倍

DeepSeek使用一個27B參數的LLM模型（其中3B為活躍參數）對NSA進行了評估。結果顯示，在8卡A100計算集群上，NSA的前向傳播和反向傳播速度分別比全注意力快9倍和6倍。在長序列解碼時，NSA相較于全注意力模型速度顯著提升，最高可達11.6倍。在通用基準、長上下文任務和基于指令的推理上，采用NSA機制的模型與全注意力模型相當或表現更優。

4. 總結

NSA作為一種新型的稀疏注意力機制，通過軟硬件協同深度優化，有效解決了現有稀疏注意力機制的缺陷，在長上下文訓練和推理中實現了顯著的加速效果，為下一代大語言模型的發展提供了新的可能性。其在訓練階段的應用以及對現代硬件的優化，使其成為一個極具性價比的解決方案。

聯系作者

文章來源：智東西
作者微信：
作者簡介：智能產業新媒體！智東西專注報道人工智能主導的前沿技術發展，和技術應用帶來的千行百業產業升級。聚焦智能變革，服務產業升級。