MIT英伟达团队革新注意力机制，破解LLM性能难题

对于大语言模型来说，长文本处理一直是难以突破的性能短板。随着序列变长，自注意力机制的计算量会以平方速度膨胀，使得模型的成本快速上升、扩展困难。

早在今年 2 月，月之暗面（Moonshot AI）就提出了全新的“块注意力混合”（Mixture of Block Attention, MoBA） 架构，试图突破这一难题。

MoBA 借鉴了混合专家系统（MoE）的思路，将长文本切分为多个数据块，并训练一个“路由器”来动态、稀疏地只关注最相关的部分。通过这种方式，MoBA 将大量无效的注意力计算剔除，从而把整体复杂度从二次方压缩到近似线性，为长上下文处理带来了真正可扩展的架构方案。

然而，尽管这一架构在理论上极为完美，但其实际应用与优化却面临两大难题：一是其成功背后的设计原理尚不明确，二是它缺少一个为现代硬件深度优化的 GPU 实现。尤其是在处理理论上更优的小尺寸数据块时，其朴素实现所带来的额外开销（Overhead）甚至会抵消稀疏化带来的计算优势，这使得其理论上的高效性难以完全转化为实际生产力。

面对这一挑战，来自麻省理工学院（MIT）韩松教授的 H.A.N. 实验室与芯片巨头英伟达（NVIDIA）的研究人员联手，推出了名为 FlashMoBA 的优化方案。

图 | 项目论文（来源：arXiv）

这项研究不仅揭示了 MoBA 成功的关键，更通过一个硬件感知的定制 CUDA（Compute Unified Device Architecture）内核，释放了其潜能。实验结果显示：FlashMoBA 在处理小数据块时，相较于当前业界领先的 FlashAttention-2，实现了高达 14.7 倍的速度提升，并且能够处理长达 512K 的超长序列而不会出现内存溢出（Out of Memory, OOM），为长上下文处理提供了解决方案。

要理解这个方案，需要明白 MoBA 的核心挑战在于其路由机制，好比让它在一个巨大的草垛中寻找一根针。模型需要从成千上万个文本块中，为每个查询精确地挑出少数几个最相关的部分。原始 MoBA 通过计算每个文本块内所有信息的“平均值”（即质心）来做决策，但这很容易让关键的信号被海量的噪声淹没。

为此，MIT 与英伟达的联合团队通过建立一个统计模型，首次将这个问题量化为一个清晰的“信噪比”（SNR）公式。该公式揭示了一个核心洞见：路由的准确性取决于注意力机制的“丰富度”（头维度 d）与文本块“粒度”（块大小 B）的比值。简而言之，要想让路由器看得更准，最直接的办法就是在模型容量不变的情况下，将文本块切分得更小。此外，研究还发现通过短卷积操作鼓励相关信息在块内“抱团”（聚类），可以作为性能倍增器，进一步放大关键信号。

然而，在 GPU 上处理大量小数据块，就像让一艘货轮去递送成千上万封单独的信件，效率极其低下。这会导致三次性能“灾难”：海量、不连续的内存读取；因需要给海量小块排序评分而产生的巨大管理开销；以及因每个任务太小而导致的 GPU 大部分时间处于闲置状态。

为了解决这些问题，FlashMoBA应运而生。它不是对 MoBA 做几处代码优化，而是一套 针对现代 GPU 从零设计的 CUDA 内核。这个优化的目的很明确：让小块 MoBA也能在 GPU 上跑得又快又省。其设计原则是尽可能把运算融合到同一内核中，并尽量减少与 HBM 之间来回搬运数据的成本。

FlashMoBA 有着两个主要核心创新点：

1.FlashTopK

这是一个全流程整合的高效流水线，能够同时完成质心计算和 top-k 选择，并且完全不需要在显存中生成巨大评分矩阵。换句话说，它从源头上避免了显存爆炸，也消除了传统 top-k 的主要性能瓶颈。

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

这是一个非常聪明的两阶段策略。可以把它想象成一位熟练的图书管理员：先把读者需要的、分散在各个书架的书一次性推到手推车上（相当于放进 GPU 的高速缓存），然后再统一处理。这样做不仅避免了频繁、缓慢的 HBM 访问，还能把原本稀疏、零散的计算重新整理成 GPU 最擅长的密集矩阵操作，大幅提升效率。

这一系列软硬件协同优化带来的成果十分显著。在端到端性能测试中，当序列长度为 64K 时，FlashMoBA 的速度是原始 MoBA 实现的 7.4 倍，内存使用量则减少了 6.1 倍。更关键的是，原始 MoBA 在处理 128K 序列时便会耗尽内存，而 FlashMoBA 则能轻松扩展至 512K。在与业界公认的高效注意力实现 FlashAttention-2 的直接对比中，FlashMoBA 在长序列场景下实现了高达 14.7 倍的加速。

（来源：论文）

效率的提升并未以牺牲模型质量为代价。研究团队从头训练了多个模型进行验证，结果有力地证明，理论上的改进完全转化为了一致的质量提升。实验证实，采用更小的块尺寸，能够显著提升模型在语言建模和长上下文检索任务上的准确率。

经过优化的 MoBA 模型，在多个基准测试中其性能与更耗费资源的密集注意力模型相当，甚至在长上下文场景中实现了超越。这是因为它有效减轻了“注意力稀释”问题——即在长文本中，标准注意力会因关注点过于分散而“失焦”，而 MoBA 通过稀疏路由，能始终将计算资源集中在最关键的信息上。

（来源：论文）

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

2.项目地址： https://github.com/mit-han-lab/flash-moba