9 月 29 日,DeepSeek 最新發布的 DeepSeek-V3.2-Exp 模型引入了自主研發的 DeepSeek Sparse Attention(DSA)稀疏注意力機制,在幾乎不影響模型性能的前提下,實現了長文本訓練和推理效率的大幅提升。本文旨在深入解析 DSA 的技術原理,並重點探討中昊芯英「剎那®」TPU 平台如何憑藉其片上緩存與高度並行矩陣計算單元,在 Lightning Indexer 鍵緩存管理、突破「內存牆」瓶頸、提升長文本處理效率及降低推理成本等方面取得的顯著收益與適配成果。
1.DSA稀疏注意力機制的技術突破
1.1 背景:如何突破傳統注意力機制的瓶頸
標準自注意力(Self-Attention)機制是 Transformer 架構的核心,但其計算和內存複雜度均與序列長度 L 的平方 O(L²) 成正比。當處理長文本時,這種二次方增長的複雜度會迅速成為性能瓶頸,導致顯存耗盡和計算延遲過高,限制了模型處理長序列的能力。為克服這一限制,學術界與工業界提出了多種稀疏注意力方案(即僅對部分 token 進行注意力計算),如滑動窗口注意力、塊稀疏注意力、可學習稀疏模式、低秩近似等。稀疏注意力可以顯著減少計算、內存開銷,並提升吞吐或降低成本。DeepSeek-V3.2-Exp的核心武器DeepSeek 稀疏注意力(DSA),首次實現了細粒度稀疏注意力機制,在幾乎不影響模型輸出效果的前提下,實現了長文本訓練和推理效率的大幅提升。
論文地址:https://github.com/deepseek-ai/DeepSeek-V3.2-Exp/blob/main/DeepSeek_V3_2.pdf
1.2 DeepSeek稀疏注意力(DSA)核心架構
DeepSeek-V3.2-Exp 與上一版本DeepSeek-V3.1-Terminus相比,核心創新在於引入了DeepSeek 稀疏注意力(DSA)。DSA 通過篩選與當前任務高度相關的文本,而非對全部歷史 token 進行全量注意力計算,從而顯著提升運算效率。
DSA 主要包含兩項關鍵技術:閃電索引器(Lightning Indexer)和細粒度稀疏注意力(Fine-grained Sparse Attention)。整體實現基於 MLA 架構,主要流程可參考 Fig.1:
Fig.1 DeepSeek-V3.2-Exp 基於MLA的注意力架構圖,其中綠色部分顯示了DSA如何根據索引器選擇top-k鍵值條目
1.3 DSA兩大核心組件
1.閃電索引器(Lightning Indexer)
閃電索引器負責快速、高效地為每個查詢(Query)從海量候選的鍵(Key)中,識別出最可能相關的 Top-k 個鍵。閃電索引器的核心目標是以極低的計算開銷完成相關鍵的「海選」。其實現方式如下:
- 低維投影:將原始高維度的 Query 和 Key 向量,通過一個獨立的、可學習的線性層投影到極低的維度(例如128維)。這使得後續的相似度計算變得異常高效。
- 高效相似度計算:使用低維投影后的向量 qI 和 kI 計算索引分數,這也是索引器的核心作用,也即為每個查詢 token(query token)計算「與前文每個 token 的相關性得分」,即索引得分 Iₜ,ₛ,公式如下:
論文指出:選擇 ReLU 激活函數的主要考量是其計算上的高吞吐量(throughput),因為與 Softmax 等需要全局歸一化的函數相比,ReLU 僅需進行一次簡單的閾值操作,計算成本低。
2.細粒度稀疏注意力(Fine-grained Sparse Attention)
基於索引器輸出的分數,token選擇機制僅保留 Top-k 索引分數對應的鍵值對(KV),再通過注意力機制計算最終輸出 uₜ(僅基於這些篩選後的「關鍵鍵值對」計算注意力)。具體工作流程為:
1). 為每個查詢token ht計算索引分數{𝐼𝑡,𝑠}
2). 選擇 Top-k 索引分數對應的鍵值條目{𝑐𝑠}
3). 在稀疏選擇的鍵值條目上應用注意力機制,公式如下:
複雜度降低:通過這一機制,核心注意力的計算複雜度從 O(L²) 成功降至 O(L*k)(其中 k 遠小於 L),實現了計算量和內存訪問的巨大節省。在 DeepSeek-V3.2-Exp 的訓練中,k 值設為 2048;也就是説即使處理 128K 長度的文本,每個查詢 token 也只需與 2048 個最相似的 token 計算注意力。
這種兩階段設計,既保證了篩選過程的高效率,又確保了最終注意力計算的高精度。
1.4 在MLA架構下的工程實現
DSA架構與DeepSeek-V3.1-Terminus 兼容,DSA 採用 MLA(多頭潛在注意力)框架的 MQA(多查詢注意力)模式(而非 MHA 模式)實現。MQA 模式下每個 KV 條目可被多個查詢共享,這種設計的關鍵優勢在於:
- 計算共享:每個潛在向量在所有查詢頭之間共享
- 內存效率:顯著減少鍵值緩存的內存佔用
- 訓練穩定性:支持從已有檢查點的平滑繼續訓練
2.「剎那®」TPU平台上的DSA效能優勢
中昊芯英「剎那®」TPU 的硬件架構與 DSA(DeepSeek Sparse Attention)稀疏計算模式的兩階段設計(即 Lightning Indexer 的篩選階段與稀疏 Attention 的計算階段)形成了高度協同,具體表現為以下兩個層面:
2.1 在Lightning Indexer計算階段的高效協同
- 片上緩存與數據局部性:Lightning Indexer 需要為每個 token 維護一個小維度(如 128 維)的鍵緩存。這個緩存尺寸能夠被完全存放在「剎那®」TPU 的高速片上靜態隨機存取存儲器(SRAM)中。這使得索引器所需的核心計算(查詢與鍵的點積運算)可完全在芯片內部完成,實現極高的數據局部性,大幅減少向高延遲、高功耗的 DRAM 的數據搬運,成為提升能效的關鍵。
- 並行矩陣計算單元:「剎那®」TPU 內部集成了大規模、高度並行的矩陣計算單元(類似於 GPU 中的 Tensor Core)。這些計算單元能夠以流水線方式,對海量的查詢-鍵對進行並行點積運算,在極短時間內完成對所有候選鍵的評分和 Top-k 篩選(如前 2048 個 token)。為後續的精確注意力計算完成了高效的「粗篩」。
2.2 在稀疏Attention計算階段的高效協同
計算效率的革新:DSA 將核心計算複雜度從 O(L²) 降至 O(L*k),直接減少了「剎那®」TPU 計算單元處理的總運算量,這使得 TPU 能夠以極短的時間完成超長序列的注意力計算。
- 確保計算單元高利用率:更關鍵的是,DSA 將計算任務從密集全注意力,轉化為一次數據量可控的稀疏 Gather 操作和一次規模顯著減小的密集矩陣乘法。這種規則化的計算模式使得 TPU 的指令調度與數據預取單元能夠高效工作,確保核心計算單元持續處於飽和狀態,避免「空轉」,從而在處理長序列時獲得數倍的性能提升。
- 極致降低訪存開銷:傳統注意力機制在處理每個新生成的 token 時,都需要將整個序列的 Key-Value 緩存從 DRAM 頻繁加載至片上緩存。隨着序列長度(L)的增加,這種操作對顯存帶寬的需求呈平方級增長,形成嚴重的「內存牆」瓶頸。DSA 的稀疏計算模式從根本上改變了這一數據訪問範式。它不再處理所有歷史 token,而是通過 Lightning Indexer 等組件,智能地篩選出少量最關鍵 token(如Top-2048)。這意味着,每次計算只需從 DRAM 中選擇性加載一小部分 KV 緩存,而非全部,從而將訪存量從 O(L²) 降 O(L*k),其中 k 為常數且 k≪L。
3.「剎那®」TPU平台上的DSA適配實踐
DeepSeek Sparse Attention(DSA) 推理流程包含兩大核心組件:Lightning Indexer 和 Sparse Multi-Latent Attention (MLA)。
- Lightning Indexer:對每個 query token,首先通過低秩投影得到 128 維的稀疏 key 表示(而 MLA 的 KV 通常為更高維,如 512),並對所有歷史 token 的 key 進行重要性打分。隨後,為每個 query 選出最相關的 Top-2048 個 key token(即 Top-k 稀疏選擇)。
- Sparse MLA:僅對被 Indexer 選中的 Top-2048 個 key token 執行稀疏注意力計算,這可極大降低計算和訪存壓力。
基於上述架構和 vLLM 社區提供的流程圖(Fig. 2),DSA 稀疏注意力的推理(prefill + decode)流程可分為以下四個核心步驟:
Fig. 2 Illustration of DeepSeek Sparse Attention (DSA) Mechanism.(圖片來源:vLLM)
Step 1:低秩投影與特徵預處理(projection for query, key, and query_head_weights)
核心實現:MLA._init_、MLA.forward、Indexer._init_
實現要點:
輸入 hidden_states 先經低秩線性變換(wq_a、wq_b),得到壓縮後的 query(q_compressed)。
query 經過 RMSNorm 歸一化(q_norm)。
query 拆分為帶/不帶 RoPE 的兩部分,帶位置信息部分用 apply_rotary_emb 加入旋轉位置編碼。
生成 query_head_weights(weights_proj),為後續稀疏選擇提供打分依據。
Step 2:稀疏打分計算(compute mqa logits)
核心實現:Indexer.forward
實現要點:
用壓縮後的query、key及query_head_weights計算稀疏打分。
該打分僅用於重要性篩選,不直接參與最終注意力輸出。
Step 3:Top-k 關鍵token篩選(select Top-k)
核心實現:Indexer.forward
實現要點:
若context_len≤2048,保留所有token連接
若 context_len > 2048,每個 query 僅保留 Top-k(如 2048)個最重要的 key,其餘跳過計算。
Step 4:稀疏注意力計算與KV緩存訪問(run sparse MLA)
核心實現:MLA.forward
實現要點:flash_mla_sparse_*(TPU原生指令優化)。
僅對上一步篩選出的 Top-k token 對執行注意力運算(通過 index_mask 控制)。
結合 KV cache 實現高效推理,支持 prefill 與 decode 兩種模式,decode 階段不會全量加載 KV cache,而是隻訪問 Top-k 相關的 KV。
關鍵收穫
算法核心:輕量預打分 → Top-k 稀疏選擇 → Sparse MLA 高效運算
複雜度變化:理論上從 O(L²) 降至 O(L*k)
實際收益:顯著降低運算量、訪存量與推理成本,在長上下文場景下性價比提升可達 50% 以上。在「剎那®」TPU 芯片上的測試顯示,DeepSeek-V3.2 與 DSA 結合帶來了以下實質性收益:
prefill階段
decode階段
4. 總結與展望
DeepSeek-V3.2-Exp 的 DSA 架構,在保持模型性能的基礎上,通過「閃電索引器+Top-k 篩選」的稀疏注意力機制,極大提升了長上下文的計算效率,將長文本推理成本降低至原來的 50% 以下,有效解決了「長文本處理貴」的核心痛點。基於這一技術突破,DSA 架構在以下場景中展現出重要的應用價值:
- 長文檔處理:針對法律合同分析、學術論文解析及大型代碼庫審查等需處理長上下文的場景,DSA 機制通過細粒度 token 選擇策略,顯著降低內存佔用與計算開銷。
- 代碼編程:對於大型代碼庫的理解和生成任務,DSA 能夠高效捕捉代碼中的長距離依賴關係,提升代碼任務中對跨文件函數調用和複雜邏輯鏈的解析能力。
- 多輪對話系統:在需要維護長對話歷史的智能助手應用中,DSA 能有效管理不斷增長的上下文長度和保持響應質量,並同時控制推理延遲。
在大型語言模型(LLM)向「更長上下文、更準確、更便宜」演進的技術浪潮中,稀疏注意力機制通過算法與硬件協同優化,正成為突破傳統 Transformer 計算瓶頸的關鍵路徑。DeepSeek-V3.2-Exp 的實踐表明,其採用的 DSA 機制將注意力計算複雜度從 O(L²) 降至 O(L*k) 級別,在長序列任務中實現顯著加速,同時保持模型性能。中昊芯英「剎那®」TPU 芯片的硬件特性——片上緩存結構、高度並行的矩陣計算單元等對稀疏計算的天然親和性及架構級支持,展現出其在承載 DeepSeek-V3.2-Exp 等前沿稀疏注意力模型方面的潛在優勢,也為未來實現更高效、低成本的長文本推理提供了可行的技術路徑。