DeepSeek NSA & Moonshot MoBA 的見解
一、背景
最近幾天 DeepSeek 的 NSA 以及 Moonshot 的 MoBA 受到廣泛關注，我也在第一時間看了相關論文。由于最近 DeepSeek 的火爆，各大平臺上也出現了非常多對這兩篇文章的介紹，這里不再贅述其襲擊，只是簡單陳述一下我的一些思考和看法。
在看這兩篇論文的時候，我們可以發現其很多觀點和手段都似曾相識，其他在之前的 Quest、SeerAttention（之前也分享過）等論文都有相關介紹。不過，其在工業落地方面提供了更多的可能性，比如 Moonshot 已經將 MoBA 應用到其生產系統中。
二、關鍵結論
首先，它們在很大程度上解決了我們在之前分享的 Token 稀疏化文章中提到的問題：
- 以前的很多方案往往針對特定場景，不夠通用。比如，RaaS 受限于特定的 Long Reasoning 場景。
- 實現不夠高效，工業落地存在挑戰。很多方案采用 Token 粒度的稀疏化方式，難以充分發揮 GPU 的算力；也有些工作，比如 RaaS 雖然也采用了 Block 粒度的稀疏化方式，但是并沒有進行高效的實現。
除此之外，這幾篇論文也進一步促進了一些共識：
- 在長序列場景（Long Prefill 或 Long Decoding），Attention 是高度稀疏化的，也是高度動態化的。
- 固定 Pattern 的稀疏化方式往往很難保持精度，可學習 Sparse Pattern 會是通用化且高精度的有效方案。
- Token 粒度的稀疏化很難充分發揮 GPU 算力，Block 粒度稀疏化是精度和性能（稀疏度、計算量）的良好平衡，基于此的高效 Block Sparse Attention 也成為標配，類似 GitHub - mit-han-lab/Block-Sparse-Attention [4]。
- 當前常見的 LLM 通常會采用 GQA，也要充分結合 GQA 的特性來設計稀疏化方案，不然可能會影響整體的稀疏化程度。
- 在進行 Block 選擇時并不需要使用 Block 內所有的 KV Cache，選擇一個代表性的“聚類中心”即可，比如取 Avg 或者 Max，MoBA 中使用 Mean Pooling，在 SeerAttention 中也有相應消融實驗。
- 不要隨意永久性丟棄 Token，由于 LLM 的自回歸特性，很難推測在后續的生成中是不是一定不需要某個 Token。這也就是為什么在 NSA 和 MOBA 中并不會節約 KV Cache 的存儲空間。
三、對訓練的影響
最后，這幾個文章都是非常好的工作，對于 Inference 的成本（長序列）節約非常有幫助；但其對于 Training 成本節約非常有限。主要是因為當前 Training 的主要成本是在 PreTraining 階段，而 PreTraining 階段的序列不會很長，常見的還是 4K。此時稀疏度不會很高，Attention 的占比也并不大，可能不到 10%。如下圖 Table 1 所示，DeepSeek V3 Training 成本中 PreTraining 占到了 95% 以上：
雖然 DeepSeek NSA 中進行了 PreTraining 實驗，并且 Loss 甚至低于 Full Attention，但其中的訓練速度實際對比的是長序列下 Attention Kernel 的速度，并不是 PreTraining 端到端的速度，如下圖 Figure 6：
而 Moonshot 的 MoBA 中也是聚焦在 SFT 階段：
我們會發現，Block 粒度的稀疏化執行的是 Block Sparse Attention，而我們之前介紹的一系列 Sample Packing 的方案也是 Block Sparse Attention，只不過 Sample Packing 中可以根據輸入數據預先知道 Block 的位置。基于這個思路，可以把 PreTraining 中的 Sample Packing 看成是 Block 粒度稀疏化的特例。如下圖所示，每種顏色的 Block 代表一個 Sample：
直觀感覺，在 PreTraining 中采用 Sample Packing 對應的 Block Sparse Attention 很有必要也很高效，然而實際的 PreTraining 中卻比較少采用。一方面是 PreTraining 中 Attention 占比不高，另一方面是負載均衡問題進一步降低性能收益；此外，也有多個工作驗證 Sample Packing 對預訓練的效果影響也很小，比如 LLaMA 3 的技術報告有介紹，忽略 Sample Packing 中的 Sparse Attention Mask，對效果沒什么影響。
當然，如果是長序列的 PreTraining 也許上述 Block Sparse Attention 能帶來一些收益，但長序列的 PreTraining 是否有必要也另當別論。
四、參考鏈接
- ??https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf??
- ??https://arxiv.org/abs/2502.11089??
- ??https://arxiv.org/pdf/2410.13276??
- ??https://github.com/mit-han-lab/Block-Sparse-Attention??
本文轉載自 ??AI閑談??，作者： AI閑談
