Lightning Indexer(閃電索引器)是DeepSeek稀疏注意力(DSA)架構中的關鍵組件,於2025年9月在DeepSeek-V3.2-Exp模型中引入。它是一個輕量級的索引模組,負責在超長序列中快速掃描並為每個查詢token計算相關性分數,僅篩選出Top-K個最相關的歷史token進行後續完整的注意力計算,從而顯著降低長序列推理的計算開銷。
隨著2025年12月DeepSeek-V3.2的發布,該組件成為模型的標準配置。華為昇騰團隊針對該運算元進行了專門的最佳化設計與實現,包括運算元融合、流水並行等,並將其代碼開源。後續,該技術也被套用於智譜AI的GLM-5等模型中。
基本介紹
- 中文名:閃電索引器
- 外文名:Lightning Indexer
- 上市時間:2025年12月
技術原理與架構,工作流程,第一步:輕量級打分,第二步:細粒度Token選擇,設計特點與最佳化,訓練方法,實現與最佳化,性能表現,
技術原理與架構
Lightning Indexer是DeepSeek稀疏注意力架構中的關鍵組件,負責在超長序列中以極低的計算開銷快速篩選出對當前查詢token最相關的一小部分上下文tokens。Lightning Indexer扮演著模型的“快速過濾器”角色,其核心目標是以極小的計算代價從超長序列中定位關鍵信息,為後續完整的注意力計算減負提速。該組件首次在DeepSeek-V3.2-Exp實驗性模型中引入,DeepSeek-V3.2使用了與DeepSeek-V3.2-Exp完全相同的架構,主要的架構調整是添加了稀疏注意力機制。
Lightning Indexer的工作原理分為兩個階段。首先由一個輕量級的打分器快速掃描整個歷史上下文序列,為每個查詢token計算與所有歷史token的簡化相關性分數。該分數通過查詢向量和鍵向量之間的縮放點積計算,並使用ReLU函式激活,可以運行在低精度下以降低開銷。其次,一個細粒度的Token選擇器根據Lightning Indexer的打分,僅保留得分最高的Top-K個token並構建一個稀疏注意力掩碼,在DeepSeek-V3.2中,K=2048。整個DSA機制的訓練也分為兩個階段。第一階段先凍結主模型只訓練Lightning Indexer,訓練目標是讓Indexer的輸出分布對齊主注意力的分布。第二階段放開所有參數,讓模型適應稀疏注意力模式。
引入Lightning Indexer機制後,在128K長度上下文下,模型長序列推理的開銷降低約50%,且模型精度幾乎無損失。長序列推理的首token輸出延遲可維持在2秒以內,每產生一個新token的逐token輸出延遲約為30毫秒。其高效性得益於定製化的融合運算元實現,該運算元將相關性打分、ReLU、ReduceSum和Top-K篩選等步驟合併。通過分塊處理、片上數據復用以及Cube與Vector計算核的流水並行等深度最佳化手段,顯著提升了端到端的篩選速度。該融合運算元已在華為昇騰CANN平台上開源實現。
工作流程
Lightning Indexer是DeepSeek稀疏注意力架構的關鍵組件,扮演快速過濾器的角色。其工作流程分為兩個階段,旨在以極低計算開銷從超長序列中快速篩選出與當前查詢最相關的少量歷史Token。
第一步:輕量級打分
Lightning Indexer使用多頭潛在注意力中的壓縮Token表示,通過查詢向量與鍵向量的縮放點積計算相似度分數,並經過ReLU激活函式產生相關性分數。該打分過程可運行在FP8低精度下,以降低算力開銷。
第二步:細粒度Token選擇
根據Lightning Indexer的打分結果,僅保留得分最高的Top-K個Token,在DeepSeek-V3.2中K=2048。基於此構建稀疏注意力掩碼,後續完整的注意力計算僅對這些被選中的Token進行。
設計特點與最佳化
Lightning Indexer僅針對查詢進行操作,因為模型只需決定每個新查詢應考慮哪些過去的Token,而鍵已被壓縮並存儲在KV快取中,無需重複評分或壓縮。其高效性得益於被實現為一個融合運算元,將打分、ReLU、TopK等操作合併。在硬體層面,該運算元通過昇騰NPU上的Cube矩陣運算核與Vector向量運算核協同及流水並行等最佳化來提升性能。
訓練方法
DeepSeek稀疏注意力(DSA,包含Lightning Indexer)的訓練分為密集預熱(Dense Warm-up)和稀疏訓練(Sparse Training)兩個階段。
第一階段為密集預熱,此階段凍結主模型參數,僅訓練Lightning Indexer,訓練目標是讓Indexer的輸出分布與原始稠密注意力的分布對齊,使用KL散度作為損失函式,此階段訓練步數較少,約為1000步,共使用2.1B tokens。
第二階段為稀疏訓練,此階段解凍所有模型參數,讓整個模型適應稀疏注意力模式,繼續使用KL散度對齊Indexer與主注意力的分布,進行大規模訓練,約15000步,共使用943.7B tokens。
實現與最佳化
Lightning Indexer的高效性得益於其針對特定硬體平台的深度最佳化實現。在華為昇騰NPU上,該運算元獲得了原生支持與深度最佳化。其核心被實現為一個融合核心,將相關性打分、非線性激活、規約求和及TopK篩選等多個操作合併,顯著減少了核心啟動開銷並提升了計算流水線效率。
針對昇騰硬體的特性,開發團隊進行了多項關鍵最佳化,包括採用運算元分塊設計並結合片上數據復用策略以最佳化記憶體訪問模式,以及充分利用昇騰NPU中Cube矩陣運算核與Vector向量運算核的協同能力,通過精細的流水並行調度來掩蓋不同計算單元間的等待時間。尤其在處理長序列時,計算密集的TopK操作耗時可以被巧妙地用來掩蓋其他相對輕量操作的耗時,從而最大化計算流水收益。
在易用性方面,昇騰已面向開發者開源了相關推理代碼和融合運算元實現。開發者可以利用昇騰的PyPTO或TileLang等高級編程工具,僅需編寫數百行代碼即可完成Lightning Indexer這類動態Shape運算元的編程與整網部署。
性能表現
引入Lightning Indexer後,128K長上下文推理的開銷降低約50%,且模型精度幾乎沒有損失。在128K長度上下文下,首token輸出延遲(TTFT)可維持在2秒以內,逐token輸出延遲(TPOT)約為30毫秒。
經過昇騰平台最佳化,Lightning Indexer運算元在128K長序列推理場景下的端到端延遲大幅降低、吞吐率成倍提升。
