Command Palette
Search for a command to run...
优化块注意力混合
优化块注意力混合
Guangxuan Xiao Junxian Guo Kasra Mazaheri Song Han
Abstract
块注意力混合(Mixture of Block Attention, MoBA)(Lu 等,2025)是一种极具前景的构建模块,能够通过使查询仅稀疏地关注一小部分键值块,高效处理大长度上下文,从而显著降低计算开销。然而,当前对 MoBA 性能背后设计原则的理解仍十分有限,且缺乏高效的 GPU 实现,限制了其实际应用。本文中,我们首先构建了一个统计模型,用以分析 MoBA 的内在机制。该模型揭示,性能的关键在于路由机制能否基于查询与键之间的亲和度,准确区分相关块与无关块。我们推导出一个信噪比指标,形式化地将架构参数与这一检索准确性联系起来。基于该分析,我们识别出两条关键的优化路径:采用更小的块尺寸,以及在键上施加短卷积以聚类相关信号,从而提升路由准确性。尽管理论上更优,但小块尺寸在 GPU 上效率较低。为弥合这一差距,我们提出了 FlashMoBA,这是一种面向硬件的 CUDA 内核,可在我们理论推荐的小块尺寸下实现高效的 MoBA 执行。通过从头训练大语言模型,我们验证了上述洞察的有效性,结果表明,改进后的 MoBA 模型在性能上可媲美密集注意力基线。FlashMoBA 在小块尺寸下的性能相较 FlashAttention-2 最高提升达 14.7 倍,使我们基于理论指导的改进真正具备实用性。