MoBA: Mixture of Block Attention for Long-Context LLMs

2025-02-18

AI Summary

Existing long-context large language models suffer from quadratic attention complexity (O(n²)) and compromised complex reasoning performance due to strong structural biases imposed by conventional sparse or linear attention methods. To address this, we propose Hybrid Block Attention (HBA), the first attention mechanism to incorporate Mixture-of-Experts (MoE) principles: it decomposes attention at the block level, employs dynamic gating for expert routing, and supports plug-and-play sparsification—enabling adaptive, task-driven switching between full and sparse attention without predefined sparsity patterns. Crucially, HBA eliminates fixed structural priors, allowing the model to autonomously learn optimal sparse structures per task. Evaluated in Kimi’s production environment, HBA efficiently handles multi-thousand-token contexts with zero accuracy degradation while significantly improving computational efficiency. This establishes a new, bias-free attention paradigm that is both highly efficient and flexible for long-text understanding and multi-step reasoning.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

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.

Problem

• Scaling context length for LLMs
• Reducing quadratic computational complexity
• Enhancing efficiency without performance loss

Innovation

• Mixture of Block Attention
• Seamless transition capability
• Efficient attention computation