arXiv:2511.11571cs.LGcs.CL2025-11被引 9

提出高效实现MoBA的FlashMoBA,显著提升长文本处理速度。

Optimizing Mixture of Block Attention

  • 通过统计建模揭示路由准确率是性能关键,提出小块尺寸与键卷积优化路径。
  • 理论推荐的小块尺寸在GPU上效率低,但新算法实现14.7倍加速。
  • 适合关注长序列建模、算子优化的研究者和工程团队。

Mixture of Block Attention (MoBA)(Lu et al., 2025)是一种高效的长上下文处理构建模块,通过让查询稀疏地关注少量键值块,大幅降低计算开销。然而其性能设计原则不清晰,且缺乏高效GPU实现,阻碍实际应用。本文首先建立统计模型分析MoBA机制,发现性能关键在于路由器根据查询-键相关性准确区分相关与无关块的能力。我们推导出一个信噪比公式,将架构参数与检索准确率关联。基于此,提出两条改进路径:使用更小的块尺寸,以及在键上应用短卷积以聚集相关信号,从而提升路由精度。尽管理论上更优,小块尺寸在GPU上效率低下。为此,我们提出FlashMoBA——一种面向硬件的CUDA内核,使小块尺寸下的高效执行成为可能。通过从头训练LLM验证,改进后的MoBA模型达到密集注意力基线的性能。FlashMoBA在小块情况下相较FlashAttention-2最高提速14.7倍,使理论优化真正落地。代码已开源:https://github.com/mit-han-lab/flash-moba。

原文摘要 · Abstract (English)

Mixture of Block Attention (MoBA) (Lu et al., 2025) is a promising building block for efficiently processing long contexts in LLMs by enabling queries to sparsely attend to a small subset of key-value blocks, drastically reducing computational cost. However, the design principles governing MoBA's performance are poorly understood, and it lacks an efficient GPU implementation, hindering its practical adoption. In this paper, we first develop a statistical model to analyze MoBA's underlying mechanics. Our model reveals that performance critically depends on the router's ability to accurately distinguish relevant from irrelevant blocks based on query-key affinities. We derive a signal-to-noise ratio that formally connects architectural parameters to this retrieval accuracy. Guided by our analysis, we identify two key pathways for improvement: using smaller block sizes and applying a short convolution on keys to cluster relevant signals, which enhances routing accuracy. While theoretically better, small block sizes are inefficient on GPUs. To bridge this gap, we introduce FlashMoBA, a hardware-aware CUDA kernel that enables efficient MoBA execution even with the small block sizes our theory recommends. We validate our insights by training LLMs from scratch, showing that our improved MoBA models match the performance of dense attention baselines. FlashMoBA achieves up to 14.7x speedup over FlashAttention-2 for small blocks, making our theoretically-grounded improvements practical. Code is available at: https://github.com/mit-han-lab/flash-moba.

MoBA长文本GPU优化注意力

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。