DeepSeek NSA & Moonshot MoBA 的見解

amei2000go

發(fā)布于 2025-2-24 11:18

瀏覽

0收藏

一、背景

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

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

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

首先，它們?cè)诤艽蟪潭壬辖鉀Q了我們?cè)谥胺窒淼?Token 稀疏化文章中提到的問題：

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

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

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

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

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

DeepSeek NSA & Moonshot MoBA 的見解-AI.x社區(qū)

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

DeepSeek NSA & Moonshot MoBA 的見解-AI.x社區(qū)

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

DeepSeek NSA & Moonshot MoBA 的見解-AI.x社區(qū)

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

DeepSeek NSA & Moonshot MoBA 的見解-AI.x社區(qū)

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

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