MIT英偉達(dá)團(tuán)隊(duì)革新注意力機(jī)制，破解LLM性能難題

對(duì)于大語(yǔ)言模型來(lái)說(shuō)，長(zhǎng)文本處理一直是難以突破的性能短板。隨著序列變長(zhǎng)，自注意力機(jī)制的計(jì)算量會(huì)以平方速度膨脹，使得模型的成本快速上升、擴(kuò)展困難。

早在今年 2 月，月之暗面（Moonshot AI）就提出了全新的“塊注意力混合”（Mixture of Block Attention, MoBA） 架構(gòu)，試圖突破這一難題。

MoBA 借鑒了混合專(zhuān)家系統(tǒng)（MoE）的思路，將長(zhǎng)文本切分為多個(gè)數(shù)據(jù)塊，并訓(xùn)練一個(gè)“路由器”來(lái)動(dòng)態(tài)、稀疏地只關(guān)注最相關(guān)的部分。通過(guò)這種方式，MoBA 將大量無(wú)效的注意力計(jì)算剔除，從而把整體復(fù)雜度從二次方壓縮到近似線性，為長(zhǎng)上下文處理帶來(lái)了真正可擴(kuò)展的架構(gòu)方案。

然而，盡管這一架構(gòu)在理論上極為完美，但其實(shí)際應(yīng)用與優(yōu)化卻面臨兩大難題：一是其成功背后的設(shè)計(jì)原理尚不明確，二是它缺少一個(gè)為現(xiàn)代硬件深度優(yōu)化的 GPU 實(shí)現(xiàn)。尤其是在處理理論上更優(yōu)的小尺寸數(shù)據(jù)塊時(shí)，其樸素實(shí)現(xiàn)所帶來(lái)的額外開(kāi)銷(xiāo)（Overhead）甚至?xí)窒∈杌瘞?lái)的計(jì)算優(yōu)勢(shì)，這使得其理論上的高效性難以完全轉(zhuǎn)化為實(shí)際生產(chǎn)力。

面對(duì)這一挑戰(zhàn)，來(lái)自麻省理工學(xué)院（MIT）韓松教授的 H.A.N. 實(shí)驗(yàn)室與芯片巨頭英偉達(dá)（NVIDIA）的研究人員聯(lián)手，推出了名為 FlashMoBA 的優(yōu)化方案。

圖 | 項(xiàng)目論文（來(lái)源：arXiv）

這項(xiàng)研究不僅揭示了 MoBA 成功的關(guān)鍵，更通過(guò)一個(gè)硬件感知的定制 CUDA（Compute Unified Device Architecture）內(nèi)核，釋放了其潛能。實(shí)驗(yàn)結(jié)果顯示：FlashMoBA 在處理小數(shù)據(jù)塊時(shí)，相較于當(dāng)前業(yè)界領(lǐng)先的 FlashAttention-2，實(shí)現(xiàn)了高達(dá) 14.7 倍的速度提升，并且能夠處理長(zhǎng)達(dá) 512K 的超長(zhǎng)序列而不會(huì)出現(xiàn)內(nèi)存溢出（Out of Memory, OOM），為長(zhǎng)上下文處理提供了解決方案。

要理解這個(gè)方案，需要明白 MoBA 的核心挑戰(zhàn)在于其路由機(jī)制，好比讓它在一個(gè)巨大的草垛中尋找一根針。模型需要從成千上萬(wàn)個(gè)文本塊中，為每個(gè)查詢(xún)精確地挑出少數(shù)幾個(gè)最相關(guān)的部分。原始 MoBA 通過(guò)計(jì)算每個(gè)文本塊內(nèi)所有信息的“平均值”（即質(zhì)心）來(lái)做決策，但這很容易讓關(guān)鍵的信號(hào)被海量的噪聲淹沒(méi)。

為此，MIT 與英偉達(dá)的聯(lián)合團(tuán)隊(duì)通過(guò)建立一個(gè)統(tǒng)計(jì)模型，首次將這個(gè)問(wèn)題量化為一個(gè)清晰的“信噪比”（SNR）公式。該公式揭示了一個(gè)核心洞見(jiàn)：路由的準(zhǔn)確性取決于注意力機(jī)制的“豐富度”（頭維度 d）與文本塊“粒度”（塊大小 B）的比值。簡(jiǎn)而言之，要想讓路由器看得更準(zhǔn)，最直接的辦法就是在模型容量不變的情況下，將文本塊切分得更小。此外，研究還發(fā)現(xiàn)通過(guò)短卷積操作鼓勵(lì)相關(guān)信息在塊內(nèi)“抱團(tuán)”（聚類(lèi)），可以作為性能倍增器，進(jìn)一步放大關(guān)鍵信號(hào)。

然而，在 GPU 上處理大量小數(shù)據(jù)塊，就像讓一艘貨輪去遞送成千上萬(wàn)封單獨(dú)的信件，效率極其低下。這會(huì)導(dǎo)致三次性能“災(zāi)難”：海量、不連續(xù)的內(nèi)存讀取；因需要給海量小塊排序評(píng)分而產(chǎn)生的巨大管理開(kāi)銷(xiāo)；以及因每個(gè)任務(wù)太小而導(dǎo)致的 GPU 大部分時(shí)間處于閑置狀態(tài)。

為了解決這些問(wèn)題，F(xiàn)lashMoBA應(yīng)運(yùn)而生。它不是對(duì) MoBA 做幾處代碼優(yōu)化，而是一套 針對(duì)現(xiàn)代 GPU 從零設(shè)計(jì)的 CUDA 內(nèi)核。這個(gè)優(yōu)化的目的很明確：讓小塊 MoBA也能在 GPU 上跑得又快又省。其設(shè)計(jì)原則是盡可能把運(yùn)算融合到同一內(nèi)核中，并盡量減少與 HBM 之間來(lái)回搬運(yùn)數(shù)據(jù)的成本。

FlashMoBA 有著兩個(gè)主要核心創(chuàng)新點(diǎn)：

1.FlashTopK

這是一個(gè)全流程整合的高效流水線，能夠同時(shí)完成質(zhì)心計(jì)算和 top-k 選擇，并且完全不需要在顯存中生成巨大評(píng)分矩陣。換句話說(shuō)，它從源頭上避免了顯存爆炸，也消除了傳統(tǒng) top-k 的主要性能瓶頸。

2.收集-致密化（Gather-and-Densify）

這是一個(gè)非常聰明的兩階段策略。可以把它想象成一位熟練的圖書(shū)管理員：先把讀者需要的、分散在各個(gè)書(shū)架的書(shū)一次性推到手推車(chē)上（相當(dāng)于放進(jìn) GPU 的高速緩存），然后再統(tǒng)一處理。這樣做不僅避免了頻繁、緩慢的 HBM 訪問(wèn)，還能把原本稀疏、零散的計(jì)算重新整理成 GPU 最擅長(zhǎng)的密集矩陣操作，大幅提升效率。

這一系列軟硬件協(xié)同優(yōu)化帶來(lái)的成果十分顯著。在端到端性能測(cè)試中，當(dāng)序列長(zhǎng)度為 64K 時(shí)，F(xiàn)lashMoBA 的速度是原始 MoBA 實(shí)現(xiàn)的 7.4 倍，內(nèi)存使用量則減少了 6.1 倍。更關(guān)鍵的是，原始 MoBA 在處理 128K 序列時(shí)便會(huì)耗盡內(nèi)存，而 FlashMoBA 則能輕松擴(kuò)展至 512K。在與業(yè)界公認(rèn)的高效注意力實(shí)現(xiàn) FlashAttention-2 的直接對(duì)比中，F(xiàn)lashMoBA 在長(zhǎng)序列場(chǎng)景下實(shí)現(xiàn)了高達(dá) 14.7 倍的加速。

（來(lái)源：論文）

效率的提升并未以犧牲模型質(zhì)量為代價(jià)。研究團(tuán)隊(duì)從頭訓(xùn)練了多個(gè)模型進(jìn)行驗(yàn)證，結(jié)果有力地證明，理論上的改進(jìn)完全轉(zhuǎn)化為了一致的質(zhì)量提升。實(shí)驗(yàn)證實(shí)，采用更小的塊尺寸，能夠顯著提升模型在語(yǔ)言建模和長(zhǎng)上下文檢索任務(wù)上的準(zhǔn)確率。

經(jīng)過(guò)優(yōu)化的 MoBA 模型，在多個(gè)基準(zhǔn)測(cè)試中其性能與更耗費(fèi)資源的密集注意力模型相當(dāng)，甚至在長(zhǎng)上下文場(chǎng)景中實(shí)現(xiàn)了超越。這是因?yàn)樗行p輕了“注意力稀釋”問(wèn)題——即在長(zhǎng)文本中，標(biāo)準(zhǔn)注意力會(huì)因關(guān)注點(diǎn)過(guò)于分散而“失焦”，而 MoBA 通過(guò)稀疏路由，能始終將計(jì)算資源集中在最關(guān)鍵的信息上。

（來(lái)源：論文）

1.論文地址：https://arxiv.org/pdf/2511.11571

2.項(xiàng)目地址： https://github.com/mit-han-lab/flash-moba