一、背景

最近幾天 DeepSeek 的 NSA 以及 Moonshot 的 MoBA 受到廣泛關(guān)注，我也在第一時間看了相關(guān)論文。由于最近 DeepSeek 的火爆，各大平臺上也出現(xiàn)了非常多對這兩篇文章的介紹，這里不再贅述其襲擊，只是簡單陳述一下我的一些思考和看法。

在看這兩篇論文的時候，我們可以發(fā)現(xiàn)其很多觀點和手段都似曾相識，其他在之前的 Quest、SeerAttention（之前也分享過）等論文都有相關(guān)介紹。不過，其在工業(yè)落地方面提供了更多的可能性，比如 Moonshot 已經(jīng)將 MoBA 應(yīng)用到其生產(chǎn)系統(tǒng)中。

二、關(guān)鍵結(jié)論

首先，它們在很大程度上解決了我們在之前分享的 Token 稀疏化文章中提到的問題：

1. 以前的很多方案往往針對特定場景，不夠通用。比如，RaaS 受限于特定的 Long Reasoning 場景。
2. 實現(xiàn)不夠高效，工業(yè)落地存在挑戰(zhàn)。很多方案采用 Token 粒度的稀疏化方式，難以充分發(fā)揮 GPU 的算力；也有些工作，比如 RaaS 雖然也采用了 Block 粒度的稀疏化方式，但是并沒有進(jìn)行高效的實現(xiàn)。

除此之外，這幾篇論文也進(jìn)一步促進(jìn)了一些共識：

• 在長序列場景（Long Prefill 或 Long Decoding），Attention 是高度稀疏化的，也是高度動態(tài)化的。
• 固定 Pattern 的稀疏化方式往往很難保持精度，可學(xué)習(xí) Sparse Pattern 會是通用化且高精度的有效方案。
• Token 粒度的稀疏化很難充分發(fā)揮 GPU 算力，Block 粒度稀疏化是精度和性能（稀疏度、計算量）的良好平衡，基于此的高效 Block Sparse Attention 也成為標(biāo)配，類似 GitHub - mit-han-lab/Block-Sparse-Attention [4]。
• 當(dāng)前常見的 LLM 通常會采用 GQA，也要充分結(jié)合 GQA 的特性來設(shè)計稀疏化方案，不然可能會影響整體的稀疏化程度。
• 在進(jìn)行 Block 選擇時并不需要使用 Block 內(nèi)所有的 KV Cache，選擇一個代表性的“聚類中心”即可，比如取 Avg 或者 Max，MoBA 中使用 Mean Pooling，在 SeerAttention 中也有相應(yīng)消融實驗。
• 不要隨意永久性丟棄 Token，由于 LLM 的自回歸特性，很難推測在后續(xù)的生成中是不是一定不需要某個 Token。這也就是為什么在 NSA 和 MOBA 中并不會節(jié)約 KV Cache 的存儲空間。

三、對訓(xùn)練的影響

最后，這幾個文章都是非常好的工作，對于 Inference 的成本（長序列）節(jié)約非常有幫助；但其對于 Training 成本節(jié)約非常有限。主要是因為當(dāng)前 Training 的主要成本是在 PreTraining 階段，而 PreTraining 階段的序列不會很長，常見的還是 4K。此時稀疏度不會很高，Attention 的占比也并不大，可能不到 10%。如下圖 Table 1 所示，DeepSeek V3 Training 成本中 PreTraining 占到了 95% 以上：

雖然 DeepSeek NSA 中進(jìn)行了 PreTraining 實驗，并且 Loss 甚至低于 Full Attention，但其中的訓(xùn)練速度實際對比的是長序列下 Attention Kernel 的速度，并不是 PreTraining 端到端的速度，如下圖 Figure 6：

而 Moonshot 的 MoBA 中也是聚焦在 SFT 階段：

我們會發(fā)現(xiàn)，Block 粒度的稀疏化執(zhí)行的是 Block Sparse Attention，而我們之前介紹的一系列 Sample Packing 的方案也是 Block Sparse Attention，只不過 Sample Packing 中可以根據(jù)輸入數(shù)據(jù)預(yù)先知道 Block 的位置?；谶@個思路，可以把 PreTraining 中的 Sample Packing 看成是 Block 粒度稀疏化的特例。如下圖所示，每種顏色的 Block 代表一個 Sample：

直觀感覺，在 PreTraining 中采用 Sample Packing 對應(yīng)的 Block Sparse Attention 很有必要也很高效，然而實際的 PreTraining 中卻比較少采用。一方面是 PreTraining 中 Attention 占比不高，另一方面是負(fù)載均衡問題進(jìn)一步降低性能收益；此外，也有多個工作驗證 Sample Packing 對預(yù)訓(xùn)練的效果影響也很小，比如 LLaMA 3 的技術(shù)報告有介紹，忽略 Sample Packing 中的 Sparse Attention Mask，對效果沒什么影響。

當(dāng)然，如果是長序列的 PreTraining 也許上述 Block Sparse Attention 能帶來一些收益，但長序列的 PreTraining 是否有必要也另當(dāng)別論。

四、參考鏈接

1. ??https://github.com/MoonshotAI/MoBA/blob/master/MoBA_Tech_Report.pdf??
2. ??https://arxiv.org/abs/2502.11089??
3. ??https://arxiv.org/pdf/2410.13276??
4. ??https://github.com/mit-han-lab/Block-Sparse-Attention??

本文轉(zhuǎn)載自 ??AI閑談??，作者： AI閑談