Source-linked AI summary
MoBA: Mixture of Block Attention for Long-Context LLMs
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang, Xinyu Zhou, Mingxing Zhang, Jiezhong Qiu
TL;DR
Long-context LLMs face quadratic attention costs, while existing sparse methods can impose predefined, task-specific structures. MoBA applies Mixture-of-Experts routing to block-based attention, and experiments report performance comparable to full attention with substantially reduced computation and support for scaling to very long contexts.
Problem
Extending LLM sequence length is difficult because traditional attention has quadratic computational complexity, while existing sparse approaches may rely on task-specific predefined structures.
Method
MoBA partitions context into blocks and dynamically routes each query token to historically relevant key-value blocks using Mixture-of-Experts principles applied to Transformer attention.
Results
MoBA achieves performance comparable to full attention with attention sparsity up to 75%, and at 10M tokens reduces attention computation time by 16x versus standard Flash Attention.
Takeaways & Limitations
MoBA supports transitions between full and sparse attention while improving the efficiency and scalability of long-context attention computation.
Takeaways & Limitations
MoBA must preserve autoregressive causality through current-block routing and causal masking to prevent information leakage from future tokens.
Abstract
from arXiv · showhide
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.
1 Introduction
Long-context processing is increasingly important for advanced LLM capabilities, but vanilla attention scales quadratically with sequence length. MoBA addresses this by dynamically routing queries to relevant context blocks while preserving the Transformer framework and supporting efficient long-sequence processing.
- Motivation: Long-sequence processing is important for LLM applications involving complex reasoning, decision-making, and extended prompts.The paper connects long-context capability to progress toward AGI and cites growing use of long inputs and long chain-of-thought outputs.
- Challenge: Vanilla attention creates a quadratic computational-growth challenge when LLM sequence lengths are extended.This motivates research into more efficient attention mechanisms.
- Existing approaches: Predefined sink or sliding-window structures can exploit attention sparsity but may be task-specific and limit generalizability.Dynamic sparse methods instead select token subsets at inference time, representing an alternative direction.
- Research question: The paper seeks an adaptable attention architecture that lets the model determine where to attend without predefined structural biases.The desired design should retain the original Transformer framework while allowing transitions between full and sparse attention.
- MoBA contribution: MoBA applies Mixture-of-Experts principles to attention by partitioning context into blocks and dynamically routing queries to relevant key-value blocks.This block-sparse design targets efficient processing of longer and more complex prompts without a proportional resource increase.
- Paper scope: The paper presents MoBA’s block partitioning, routing strategy, computational efficiency, and long-sequence experimental results.These components frame MoBA as an approach to efficient attention computation for lengthy inputs.
2 Method
MoBA applies Mixture-of-Experts principles to attention by partitioning context into blocks and dynamically selecting relevant key-value blocks for each query. Its causal routing, FlashAttention-based implementation, and optional hybridization with full attention support efficient long-context processing.
- MoBA Architecture: MoBA partitions the full context into blocks and uses a gating mechanism to select a subset of key-value blocks for each query.The block size is B = N/n, where N is context length and n is the number of blocks.
- MoBA Architecture: The gate scores each block by the inner product between the query and the block’s mean-pooled keys, then applies top-k selection.The selected blocks form the sparse query-to-block attention pattern.
- MoBA Architecture: A running example routes two queries to different top-two block pairs, illustrating that block selection is dynamic across queries.The first query selects the first and second blocks, while the second selects the third and fourth.
- Causality: MoBA preserves autoregressive causality by excluding future blocks and always routing each query to its current block with causal masking.Current-block routing prevents information leakage from future tokens while retaining local context.
- Design Choices: MoBA can replace full attention without changing parameter count, allowing attention layers to transition between full attention and MoBA during training.The method also explores fine-grained segmentation and views sliding-window and sink attention as special cases of block selection.
- Implementation: The implementation partitions key-value matrices, computes gated assignments under a causal mask, reorganizes queries and blocks, and computes separate attentions with FlashAttention before combining outputs.The workflow uses variable-length FlashAttention and online softmax to combine current-block and MoBA attention results.
3 Experiments
Experiments show that MoBA can match full attention across scaling and long-context evaluations while using substantially sparser attention. Fine-grained blocks and hybrid training further improve its performance–efficiency trade-off.
- Long Context Scalability: 95.31% sparsity at 32K context leaves MoBA with a marginally higher last-block LM loss, but the gap narrows across five experiments.Trailing LM loss evaluates the final tokens of sequences reaching the maximum length.
- Ablation Study on Fine-Grained Block Segmentation: 1e-2 separates the coarsest 2-of-8 block setting from finer-grained configurations at consistent 75% sparsity.The ablation uses a 1.5B-parameter model with 32K context and compares 8, 16, 32, 64, and 128 blocks.
- MoBA/Full Hybrid Training: A MoBA/full hybrid recipe reaches nearly the same position-wise loss as full attention without significant loss spikes during the attention switch.The hybrid trains with MoBA for 90% of tokens and full attention for the remaining 10%.
- Large Language Modeling Evaluation: 95.31% sparsity with the last three layers kept as full attention enables MoBA to nearly match full attention on RULER at 1M-token-scale evaluation.RULER scores are 0.7818 for MoBA and 0.7849 for full attention; generation uses full attention.
- Efficiency and Scalability: 6.5x speedup during 1M-token prefilling demonstrates MoBA’s sub-quadratic attention computation relative to full attention.The efficiency comparison isolates the attention layer because other layers have identical FLOPs.
4 Related Work
Efficient attention research spans static and dynamic sparse patterns, training-free inference optimizations, and alternatives that depart from traditional Transformer attention. These approaches offer different trade-offs in computational efficiency, adaptability, and compatibility with pretrained models.
- Efficient attention methods aim to reduce self-attention’s quadratic complexity while maintaining model performance.
- Static Sparse Patterns: Static sparse patterns impose predefined structures, including strided, fixed, sink-based, and sliding-window attention patterns.
- Dynamic Sparse Patterns: Dynamic sparse methods adaptively select tokens or clusters using techniques such as locality-sensitive hashing, K-means, and nearest-neighbor retrieval.
- Training-free Sparse Attention: Training-free sparse attention targets prefill, decode, or both inference stages by profiling prompts or selecting relevant tokens.
- Longheads is closely related to MoBA because its top-1 gating network selects the most relevant key-value block for each query.
- Beyond Traditional Attention Architecture: Alternative architectures based on CNNs, RNNs, SSMs, or linear attention depart from conventional attention but generally require training from scratch and cannot reuse pretrained Transformer models.
5 Conclusion
MoBA is presented as an MoE-inspired attention architecture for efficient, scalable long-context processing. It partitions contexts into blocks, dynamically routes queries to relevant key-value blocks, and is reported to preserve performance while improving efficiency and supporting existing models.
- MoBA applies Mixture of Experts principles to attention by partitioning context into blocks and dynamically routing queries to relevant key-value blocks.
- MoBA achieves performance comparable to full attention while significantly improving computational efficiency across long-context evaluations.
- MoBA scales to long contexts while maintaining low language-model losses and high performance on various benchmarks.
- MoBA can be integrated with existing models without substantial training cost, supporting continual pre-training for long-context capabilities.
A.1 Long Context Scalability
The long-context scalability analysis evaluates loss across position-based segments and examines scaling laws over two context ranges. It reports that trailing tokens account for most of the performance discrepancy between full-context and sparse-attention baselines.
- Sequences are segmented by actual positions to counter the short-context bias in natural data distributions and balance evaluation across context lengths.
- The 30K-32K segment measures losses for documents exceeding 30K context length and masks positions from 30K to 32K.
- Trailing tokens account for the majority of the performance discrepancy between the full-context baseline and newly proposed sparse-attention architectures.
- Figure 8 reports scaling laws separately for positions 0-16k and 16-32k, while Table 3 presents loss scaling by position.