1501
2025年2月18日，DeepSeek發(fā)布了最新的論文《Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention》，這一創(chuàng)新成果迅速引發(fā)了科技圈的廣泛關(guān)注。
問題與挑戰(zhàn)
論文的核心內(nèi)容圍繞一種全新的稀疏注意力機(jī)制——NSA（Native Sparse Attention）展開。隨著大語言模型（LLM）的快速發(fā)展，處理長文本的能力已成為衡量模型性能的重要指標(biāo)之一。然而，傳統(tǒng)LLM采用的“全注意力”機(jī)制在處理長文本時(shí)計(jì)算成本過高，因?yàn)樾枰?jì)算每個(gè)詞與其他所有詞之間的關(guān)系，從而導(dǎo)致高延遲和高能耗。
盡管稀疏注意力機(jī)制作為一種減少計(jì)算量的方法被廣泛應(yīng)用，但現(xiàn)有的稀疏注意力機(jī)制在訓(xùn)練端效果不佳，且難以同時(shí)提升推理效率和模型性能。
創(chuàng)新路徑deep
NSA機(jī)制正是為了解決這一問題而誕生的。它具有硬件對(duì)齊和本地可訓(xùn)練的特性，通過優(yōu)化設(shè)計(jì)，顯著提升了大模型訓(xùn)練和推理的效率，同時(shí)降低了對(duì)硬件的要求。
圖：NSA架構(gòu)概述。左：該框架通過三個(gè)并行的注意力分支處理輸入序列：對(duì)于給定的查詢，前面的鍵和值被處理為粗粒度模式的壓縮注意力、重要令牌塊的選定注意力和局部上下文的滑動(dòng)注意力。右圖：每個(gè)分支產(chǎn)生的不同注意力模式的可視化。綠色區(qū)域表示需要計(jì)算注意力得分的區(qū)域，而白色區(qū)域表示可以跳過的區(qū)域。
NSA機(jī)制提出了“三合一”的注意力處理路徑，具體包括：
粗粒度的token壓縮（Compression）：保留粗顆粒度的token模塊，將長文本分成多個(gè)塊，并用機(jī)器學(xué)習(xí)模型提取每個(gè)塊的摘要，從而減少計(jì)算量。
細(xì)粒度的token選擇（Selection）：通過引入query token（qt），篩選出與當(dāng)前計(jì)算token相關(guān)性最高的模塊，并只計(jì)算這些塊的詳細(xì)信息，進(jìn)一步提高計(jì)算效率。
滑動(dòng)窗口（Sliding Window）：獲取局部最近的完整token序列，補(bǔ)充上下文信息，確保模型的準(zhǔn)確性和完整性。
性能優(yōu)化
NSA注意力機(jī)制還通過硬件優(yōu)化來進(jìn)一步提升性能。它通過將數(shù)據(jù)按連續(xù)塊加載到GPU的內(nèi)存中，減少隨機(jī)訪問的開銷，實(shí)現(xiàn)分塊內(nèi)存訪問。在解碼階段，令多個(gè)注意力頭共享相同的鍵值緩存，降低內(nèi)存帶寬需求。而在訓(xùn)練階段，直接引入稀疏性，確保模型學(xué)會(huì)如何有效利用稀疏結(jié)構(gòu)，實(shí)現(xiàn)端到端訓(xùn)練。
實(shí)測結(jié)果顯示，在解碼速度上，NSA提升了11.6倍，正向和反向傳播階段分別提速9倍和6倍。在基準(zhǔn)測試中，NSA的表現(xiàn)與完全注意力機(jī)制相當(dāng)甚至更好，尤其是在長上下文任務(wù)和推理任務(wù)中。
圖：全注意力模型和我們的NSA之間的性能和效率比較。左：盡管數(shù)量稀少，但NSA在一般基準(zhǔn)測試、長上下文任務(wù)和推理評(píng)估中平均超過了全注意力基線。右：對(duì)于64k長度的序列處理，與Full Attention相比，NSA在所有階段（解碼、前向傳播和后向傳播）都實(shí)現(xiàn)了顯著的計(jì)算速度提升。
此外，NSA機(jī)制還通過硬件對(duì)齊系統(tǒng)和訓(xùn)練感知設(shè)計(jì)，支持高效部署和端到端訓(xùn)練。論文中提到的Triton框架支持多種GPU平臺(tái)（如CUDA、ROCm、CANN），為國產(chǎn)化算力提供了可能性。
NSA的內(nèi)核設(shè)計(jì)。內(nèi)核按GQA組加載查詢（網(wǎng)格循環(huán)），獲取相應(yīng)的稀疏KV塊（內(nèi)環(huán)），并在SRAM上執(zhí)行注意力計(jì)算。綠色塊表示SRAM上的數(shù)據(jù)，藍(lán)色塊表示HBM上的數(shù)據(jù)。
意義與影響
NSA注意力機(jī)制的提出，解決了稀疏注意力機(jī)制在實(shí)際應(yīng)用中的兩大難題：推理效率和訓(xùn)練支持。這一成果證明了稀疏注意力不僅可以用于推理階段，還可以用于訓(xùn)練階段，為下一代低功耗、高吞吐的LLM部署提供了關(guān)鍵技術(shù)基礎(chǔ)。同時(shí)，NSA降低了對(duì)GPU顯存和卡間通信能力的要求，可能緩解國產(chǎn)大模型在GPU芯片上的依賴，推動(dòng)AI算力需求的發(fā)展。
基于Triton的NSA內(nèi)核與基于Triton FlashAttention-2內(nèi)核的比較。DeepSeek的實(shí)現(xiàn)顯著降低了所有上下文長度的延遲，隨著輸入長度的增加，這種改善變得更加明顯。
總結(jié)
DeepSeek的這篇最新論文提出了一種創(chuàng)新的稀疏注意力機(jī)制NSA，通過動(dòng)態(tài)分層稀疏策略、粗粒度的token壓縮和細(xì)粒度的token選擇，以及硬件優(yōu)化等方法，顯著提升了LLM處理長文本時(shí)的推理速度和訓(xùn)練效率。該論文不僅展示了DeepSeek在注意力機(jī)制上的創(chuàng)新實(shí)力，也為大模型的高效訓(xùn)練和推理提供了新的解決方案，具有重要的學(xué)術(shù)和產(chǎn)業(yè)價(jià)值。
