MoBA: Mixture of Block Attention for Long-Context LLMs
Authors: name: Enzhe Lu, name: Zhejun Jiang, name: Jingyuan Liu, name: Yulun Du, name: Tao Jiang, name: Chao Hong, name: Shaowei Liu, name: Weiran He, name: Enming Yuan, name: Yuzhi Wang, name: Zhiqi Huang, name: Huan Yuan, name: Suting Xu, name: Xinran Xu, name: Guokun Lai, name: Yanru Chen, name: Huabin Zheng, name: Junjie Yan, name: Jianlin Su, name: Yuxin Wu, name: Neo Y. Zhang, name: Zhilin Yang, name: Xinyu Zhou, name: Mingxing Zhang, name: Jiezhong Qiu
Organization: Moonshot AI

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.