扩展有效上下文长度对于推动大语言模型迈向通用人工智能至关重要。然而，传统注意力机制固有的二次计算复杂度带来了高昂的开销。现有方法要么引入强偏置结构（如汇聚注意力或窗口注意力），这些结构具有任务特异性；要么将注意力机制彻底修改为线性近似，其在复杂推理任务中的性能尚未得到充分探索。本文提出了一种遵循“少结构”原则的解决方案，允许模型自主决定关注位置，而非引入预定义偏置。我们引入了混合块注意力（MoBA），这是一种将混合专家（MoE）原理应用于注意力机制的创新方法。该新型架构在长上下文任务上表现出优越性能，同时具备一个关键优势：能够在全注意力和稀疏注意力之间无缝切换，从而在提升效率的同时避免性能下降的风险。MoBA 已部署用于支持 Kimi 的长上下文请求，并在大语言模型的高效注意力计算方面展现了显著进展。我们的代码可在 https://github.com/MoonshotAI/MoBA 获取。
Scaling the effective context length is essential for advancing large language models (LLMs) toward artificial general intelligence (AGI). However, the quadratic increase in computational complexity inherent in traditional attention mechanisms presents a prohibitive overhead. Existing approaches either impose strongly biased structures, such as sink or window attention which are task-specific, or radically modify the attention mechanism into linear approximations, whose performance in complex reasoning tasks remains inadequately explored. In this work, we propose a solution that adheres to the ``less structure'' principle, allowing the model to determine where to attend autonomously, rather than introducing predefined biases. We introduce Mixture of Block Attention (MoBA), an innovative approach that applies the principles of Mixture of Experts (MoE) to the attention mechanism. This novel architecture demonstrates superior performance on long-context tasks while offering a key advantage: the ability to seamlessly transition between full and sparse attention, enhancing efficiency without the risk of compromising performance. MoBA has already been deployed to support Kimi's long-context requests and demonstrates significant advancements in efficient attention computation for LLMs. Our code is available at https://github.com/MoonshotAI/MoBA.
核心贡献 · Key contributions
提出 MoBA，将混合专家机制应用于注意力，实现动态块选择。 Proposes MoBA, applying Mixture-of-Experts to attention for dynamic block selection.
实现次二次复杂度，在 1000 万词元时加速比达 16 倍。 Achieves sub-quadratic complexity with up to 16x speedup at 10M tokens.
支持全注意力和稀疏注意力之间无缝切换，且无性能损失。 Enables seamless transition between full and sparse attention without performance loss.
展示与全注意力相当的缩放定律，稀疏度高达 95%。 Demonstrates scaling laws comparable to full attention with up to 95% sparsity.
在长上下文基准（如 RULER）和百万词元大海捞针任务上验证。 Validates on long-context benchmarks (e.g., RULER) and 1M-token Needle-in-Haystack.
已部署于 Kimi，支持高效长上下文处理。 Deployed in Kimi for efficient long-context processing.
局限 · Limitations
监督微调中因损失掩码导致稀疏梯度，存在性能差距。 Performance gap in supervised fine-tuning due to sparse gradients from loss masking.
需仔细调整块大小和 top-k 以优化稀疏性与性能的权衡。 Requires careful tuning of block size and top-k for optimal sparsity-performance trade-off.
复杂推理任务评估有限；深度推理的有效性尚未充分探索。 Limited evaluation on complex reasoning tasks; effectiveness for deep reasoning not fully explored.
当前实现聚焦预填充阶段；生成阶段仍使用全注意力。 Current implementation focuses on prefill stage; generation still uses full attention.
扩展到极长上下文（如超过 1000 万词元）可能面临内存和路由开销。 Scalability to extremely long contexts (e.g., >10M tokens) may face memory and routing overhead.
论文章节 · Sections（共 12）
摘要Abstract
1 引言1 Introduction
2 方法2 Method
2.1 预备知识：Transformer 中的标准注意力2.1 Preliminaries: Standard Attention in Transformer
2.2 MoBA 架构2.2 MoBA Architecture
2.3 实现2.3 Implementation
3.1 缩放定律实验与消融研究3.1 Scaling Law Experiments and Ablation Studies
3.2 MoBA 与全注意力的混合3.2 Hybrid of MoBA and Full Attention