機(jī)器之心報道

機(jī)器之心編輯部

今年 2 月，月之暗面提出了一種名為 MoBA 的注意力機(jī)制，即 Mixture of Block Attention，可以直譯為「塊注意力混合」。

據(jù)介紹，MoBA 是「一種將混合專家（MoE）原理應(yīng)用于注意力機(jī)制的創(chuàng)新方法?！乖摲椒ㄗ裱父俳Y(jié)構(gòu)」原則，并不會引入預(yù)定義的偏見，而是讓模型自主決定關(guān)注哪些位置。

MoBA 在處理長上下文時表現(xiàn)出極強(qiáng)的潛力，它允許 Query 只稀疏地關(guān)注少量 Key-Value 塊，從而大幅降低計算成本。

然而，目前業(yè)界對 MoBA 性能背后的設(shè)計原則仍缺乏深入理解，同時也缺少高效的 GPU 實(shí)現(xiàn)，這限制了其實(shí)際應(yīng)用。

在這篇論文中，來自MIT、NVIDIA機(jī)構(gòu)的研究者首先建立了一個統(tǒng)計模型，用于分析 MoBA 的內(nèi)部機(jī)制。模型顯示，其性能關(guān)鍵取決于路由器是否能夠基于 Query-Key 的相似度，準(zhǔn)確區(qū)分相關(guān)塊與無關(guān)塊。研究者進(jìn)一步推導(dǎo)出一個信噪比，將架構(gòu)參數(shù)與檢索準(zhǔn)確率建立起形式化聯(lián)系。

基于這一分析，本文識別出兩條主要的改進(jìn)路徑：一是采用更小的塊大小，二是在 Key 上應(yīng)用短卷積，使語義相關(guān)信號在塊內(nèi)聚集，從而提升路由準(zhǔn)確性。

然而，盡管小塊尺寸在理論上更優(yōu)，但在現(xiàn)有的 GPU 實(shí)現(xiàn)中，小塊會導(dǎo)致嚴(yán)重的內(nèi)存訪問碎片化和低并行度，速度甚至慢于稠密注意力。

為解決這一矛盾，研究者進(jìn)一步提出了FlashMoBA，一種硬件友好的 CUDA kernel，可在小塊配置下仍然高效地執(zhí)行 MoBA。

結(jié)果顯示優(yōu)化后的 MoBA 在性能上可與密集注意力基線相匹敵。對于小塊場景，F(xiàn)lashMoBA 相比 FlashAttention-2 可實(shí)現(xiàn)最高 14.7 倍加速。

• 論文地址：https://arxiv.org/pdf/2511.11571
• 項目地址：https://github.com/mit-han-lab/flash-moba
• 論文標(biāo)題：OPTIMIZING MIXTURE OF BLOCK ATTENTION

FLASHMOBA：一種面向小塊 MoBA 的優(yōu)化內(nèi)核

理論模型表明，較小的塊尺寸能帶來顯著的質(zhì)量提升，但樸素的 GPU 實(shí)現(xiàn)效率低下。由月之暗面發(fā)布的原始 MoBA 實(shí)現(xiàn)，在配置小塊尺寸時會遭遇性能瓶頸，這些瓶頸抵消了稀疏性帶來的計算節(jié)省，導(dǎo)致執(zhí)行速度比稠密注意力更慢。

研究者推出了 FlashMoBA，這是一種硬件感知的 CUDA 內(nèi)核，旨在使小塊 MoBA 變得實(shí)用且高效。

小塊帶來的性能挑戰(zhàn)

小塊尺寸引入了幾個關(guān)鍵的性能挑戰(zhàn)，要在實(shí)際部署中應(yīng)用必須解決這些問題。

首先，在為每個查詢收集稀疏、不連續(xù)的鍵值塊時，會出現(xiàn)低效的內(nèi)存訪問，導(dǎo)致從 HBM 讀取數(shù)據(jù)時出現(xiàn)非合并內(nèi)存讀取。

FLASHMOBA 內(nèi)核設(shè)計

為了克服這些挑戰(zhàn)，F(xiàn)lashMoBA 采用了三個融合內(nèi)核，以最大限度地減少 HBM 往返次數(shù)，并使計算與 GPU 架構(gòu)相對齊，如圖 1 所示。

• 分塊 Top-K 選擇

Top-k 選擇過程是原始 MoBA 實(shí)現(xiàn)中的主要瓶頸，該實(shí)現(xiàn)顯式生成了完整的分?jǐn)?shù)矩陣并串行處理批次序列。研究者將其替換為Flash TopK（圖 1 中的步驟 1），這是一個由融合內(nèi)核組成的高度優(yōu)化的三階段流水線。

最后，一個高效的后處理步驟將以查詢?yōu)橹行牡乃饕匦赂袷交癁橐枣I塊為中心的變長布局，以便進(jìn)行主注意力傳遞。整個流水線在批次和注意力頭之間完全并行化，消除了原始的性能瓶頸。

• 采用「收集并致密化」策略的前向傳播

為了處理 MoBA 的不規(guī)則稀疏性，前向內(nèi)核使用了一種基于兩級分塊機(jī)制的「收集并致密化」策略，詳見算法 1。

要區(qū)分兩種類型的塊：

這種兩級方法是關(guān)鍵所在，因為在 SRAM 中緩存查詢允許在邏輯鍵塊的所有物理圖塊之間復(fù)用數(shù)據(jù)，從而通過高效的稠密 GEMM（通用矩陣乘法）分?jǐn)偘嘿F的不規(guī)則內(nèi)存訪問成本。

• 帶重計算的反向傳播

反向傳播利用了 FlashAttention-2 的內(nèi)存高效設(shè)計，并實(shí)現(xiàn)為三個內(nèi)核的序列（算法 5）。

主內(nèi)核在鍵維度上并行化計算，每個線程塊處理一個鍵塊。為了處理稀疏性，它鏡像了前向傳播的「收集并致密化」策略，使用變長索引收集查詢子集并將梯度輸出到片上圖塊中。

遵循 FlashAttention-2 的方法，研究者在反向傳播期間重計算注意力分?jǐn)?shù)，以避免將完整的注意力矩陣存儲在內(nèi)存中。雖然鍵和值的梯度直接寫入 HBM，但部分查詢梯度需要跨多個鍵塊進(jìn)行累加，這是通過對高精度全局緩沖區(qū)使用原子加法來高效且安全地處理的。

這種設(shè)計確保了反向傳播在序列長度上保持線性復(fù)雜度，這是相對于標(biāo)準(zhǔn)注意力的二次復(fù)雜度的一個關(guān)鍵改進(jìn)。由于反向傳播通常構(gòu)成優(yōu)化注意力實(shí)現(xiàn)的主要性能瓶頸（通常比前向傳播慢 2-3 倍），因此我們需要反向內(nèi)核的高效率對于實(shí)現(xiàn)長序列的實(shí)際訓(xùn)練至關(guān)重要。

實(shí)驗及結(jié)果

本文從零開始預(yù)訓(xùn)練模型，并進(jìn)行可控實(shí)驗來驗證 MoBA 的設(shè)計原則。實(shí)驗共訓(xùn)練了兩個模型，所有實(shí)驗均在 8× H100 80GB GPU 上完成：

• 340M 參數(shù)模型（hidden size 1024，16 heads，中間層規(guī)模 2816）；
• 1B 參數(shù)模型（hidden size 2048，32 heads，中間層規(guī)模 8192）。

質(zhì)量評估結(jié)果

本文在語言建模、長上下文檢索以及真實(shí)任務(wù)上對 MoBA 的表現(xiàn)進(jìn)行了評估。實(shí)驗結(jié)果表明，改進(jìn)后的模型在多種基準(zhǔn)測試中提高了性能。

這一趨勢在所有基準(zhǔn)和不同模型規(guī)模上都保持一致。對 340M 模型來說，將塊大小從 512 縮小到原來的 1/4 到 128，可帶來如下提升：

• 語言建模準(zhǔn)確率從 44.6% 提升到 45.6%（表 1）；
• RULER 準(zhǔn)確率從 38.8% 提升到 63.9%（表 3）；
• LongBench 綜合得分從 13.2 提升到 15.3（表 5）。

總體來看，小塊尺寸對于 MoBA 達(dá)到與密集注意力相當(dāng)?shù)男阅苁潜匾摹?/p>

Key Convolution 。Key Convolution 在不同任務(wù)中都能帶來性能提升，而且具有任務(wù)偏好特性。對于 340M 模型：

• kconv3 將語言建模準(zhǔn)確率從 45.1% 提升到 45.6%（表 1）；
• kconv5 在 64K 長度檢索任務(wù)中達(dá)到 100% 的檢索率（表 3）；
• 在 LongBench 上，kconv3 得分達(dá)到 15.3%（表 5）。

對于 1B 模型：

• kconv3 將語言建模準(zhǔn)確率提升到 52.7%（表 2）；
• 將 RULER 準(zhǔn)確率提升到 68.2%（表 4）。

這些結(jié)果表明，卷積通過使相關(guān) token 在塊內(nèi)聚集，提升了有效均值差異 ，從而顯著提高路由準(zhǔn)確性。

注：卷積核寬度 W∈{3,5}，分別記作 kconv3 和 kconv5。

稀疏匹配密集注意力機(jī)制。在多個基準(zhǔn)測試和規(guī)模下，MoBA 的表現(xiàn)與密集注意力機(jī)制相當(dāng)甚至更勝一籌。

效率結(jié)果

雖然理論上小塊尺寸能夠帶來更高的模型質(zhì)量，但此前由于 GPU 利用率低下，小塊一直難以在實(shí)際中使用。FlashMoBA 的出現(xiàn)讓這些配置真正變得可行。

端到端性能。圖 3 對比了不同序列長度（8K 至 512K token）下的延遲和內(nèi)存占用。FlashMoBA 在兩項指標(biāo)上都顯著優(yōu)于原始實(shí)現(xiàn)。

在 N=64K 且 B=128 的配置下：FlashMoBA 比原始 MoBA 快 7.4 倍，內(nèi)存占用減少 6.1 倍，原始 MoBA 在 128K 序列就會 OOM（內(nèi)存溢出），而 FlashMoBA 能擴(kuò)展到 512K。

隨著序列越長、塊越小，優(yōu)勢更明顯，因為 FlashMoBA 消除了全局 reindex 的開銷，在長序列條件下可實(shí)現(xiàn)最高 14.7× 快于 FlashAttention-2 的速度。

為了理解 FlashMoBA 的提速來源，圖 4 展示了在 N=64K 下前向傳播的耗時分布。

原始 MoBA 包含 5 個階段：（1）計算質(zhì)心并執(zhí)行 top-k、（2）全局 reindex、（3）在路由后的索引上執(zhí)行注意力、（4）局部因果注意力以及（5）合并結(jié)果。

其中步驟 (1)、(2)、(5) 占據(jù)了超過 70% 的執(zhí)行時間。

FlashMoBA 則使用兩個融合 kernel，這種融合設(shè)計將 64K 序列下的前向傳播時間降至 49 ms，而 FlashAttention-2 在相同設(shè)置下為 99 ms。