MoBA: Mixture of Block Attention for Long-Context LLMs

Jiezhong Qiu (Tsinghua University) · Xinyu Zhou (Peking University) · Jingyuan Liu (Nanjing University) · Zhilin Yang (Tsinghua University, Tsinghua University) · Yuxin Wu (Renmin University of China) · Enzhe Lu (Moonshot AI) · Zhejun Jiang (Moonshot AI) · Yulun Du (Moonshot AI) · Tao Jiang (Moonshot AI) · Chao Hong (Moonshot AI) · Shaowei Liu (Moonshot AI) · Weiran He (Moonshot AI) · Enming Yuan (Tsinghua University, Tsinghua University) · Yuzhi Wang (Moonshot AI) · Zhiqi Huang (Moonshot AI) · Huan Yuan (Beihang University) · Suting Xu (Moonshot AI) · Xinran Xu (Moonshot AI) · Guokun Lai (Moonshot AI Ltd) · Yanru Chen (Tsinghua University, Tsinghua University) · Huabin Zheng (Moonshot AI) · Junjie Yan (Megvii Technology Inc.) · Jianlin Su (Shenzhen Zhuiyi Technology Co., Ltd.) · Yutao Zhang · Mingxing Zhang (Tsinghua University)
artificial general intelligenceattention mechanismsbiased structurescomplex reasoningcomputational complexitycontext lengthefficient attention computationless structure principlelinear approximationsmixture of block attentionmixture of expertssink attentionsparse attentionwindow attention

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 handle actual production workloads with long-context requirements, demonstrating significant advancements in efficient attention computation for LLMs. Our code is available at https://github.com/MoonshotAI/MoBA.