Elamur Research⌕
☰
home / papers / 2025 / MoBA: Mixture of Block Attention for Long-Context LLMs

MoBA: Mixture of Block Attention for Long-Context LLMs

Paper · 2025 · Enzhe Lu, Zhejun Jiang, Jingyuan Liu et al.

Summary

Standard softmax attention grows quadratically with sequence length, which makes long-context LLMs expensive to train and run. Earlier fixes either hard-code sparsity patterns such as sink or sliding-window attention, which tend to suit only some tasks, or switch to linear attention models like Mamba, RWKV and RetNet, which need costly conversion or retraining and have little evidence on complex reasoning. The authors, from Moonshot AI and partners, propose MoBA (Mixture of Block Attention). It applies the Mixture of Experts (MoE) routing idea to attention, following a "less structure" principle. The context is split into fixed-size key/value blocks, and a parameter-free top-k gate picks which blocks each query attends to. MoBA uses the same parameters as full attention, so individual layers can be switched between sparse and full modes, including during training. The authors report that MoBA performs on par with full attention while costing much less compute. This makes it a practical continual pre-training route to longer context, and it already serves long-context requests in Kimi.

Key points

  • Gating: each block's relevance score is the dot product of the query with the mean-pooled keys of that block. Top-k selection then decides which blocks the query attends to. Fine-grained (smaller) blocks are also explored, following MoE practice.
  • Causality: two rules preserve it.
    • Future blocks are masked out of routing.
    • The query's own block is always selected and handled with a causal mask, similar to shared experts in MoE.
  • Hybrid mode: layers can be set to full attention or MoBA at initialization, and this choice can change during training. MoBA is implemented together with Flash Attention.
  • Related methods: the authors frame Quest as MoBA with smaller blocks and min/max pooling, and Longheads as MoBA with top-1 gating.
  • Resource: code is available at github.com/MoonshotAI/MoBA. The provided excerpt omits the experimental section, so specific benchmark figures are not included here.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Tsinghua UniversityBeijing · China137 works · 83,262 citations🎓University of Illinois Urbana-ChampaignUrbana · United States108 works · 67,432 citations🎓Renmin University of ChinaBeijing · China35 works · 14,757 citations🎓Beihang UniversityBeijing · China27 works · 8,089 citations🎓Nanjing UniversityNanjing · China28 works · 7,072 citations

Also on the byline: Megvii · Moonshot AI.

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 9,086 worksArtificial Intelligence · 8,889 worksArts and Humanities · 9,296 worksLanguage and Linguistics · 9,206 worksComputer Vision and Pattern Recognition · 1,282 works

🧭 Topics 1Research subjects and languages the work is about.Language models & pretraining · 1,343 works

Cite

Cite as 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 and Jiezhong Qiu. 2025. MoBA: Mixture of Block Attention for Long-Context LLMs. arXiv:2502.13189.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.48550/arxiv.2502.13189 · arXiv 2502.13189

@misc{lu2025moba,
  title = {{MoBA: Mixture of Block Attention for Long-Context LLMs}},
  author = {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 and Jiezhong Qiu},
  year = {2025},
  doi = {10.48550/arxiv.2502.13189},
  eprint = {2502.13189},
  archivePrefix = {arXiv},
  url = {https://doi.org/10.48550/arxiv.2502.13189}
}

References

65 references, 19 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Ainslie, Joshua, Tao Lei, et al. “Colt5: Faster long-range transformers with conditional computation”. In: arXiv preprint arXiv:2303.09752 (2023).
  2. Ainslie, Joshua, Santiago Ontanon, et al. “ETC: Encoding long and structured inputs in transformers”. In: arXiv preprint arXiv:2004.08483 (2020).
  3. An, Chenxin et al. “Why Does the Effective Context Length of LLMs Fall Short?” In: arXiv preprint arXiv:2410.18745 (2024). — [arXiv:2410.18745](https://arxiv.org/abs/2410.18745) find
  4. Anthropic. Introducing 100K Context Windows. https://www.anthropic.com/news/100k-context-windows.2023. Beltagy, Iz, Matthew E Peters, and Arman Cohan. “Longformer: The long-document transformer”. In: arXiv preprint arXiv:2004.05150 (2020). ↗
  5. Bertsch, Amanda et al. “Unlimiformer: Long-range transformers with unlimited length input”. In: Advances in Neural Information Processing Systems 36 (2024). find
  6. Bick, Aviv et al. “Transformers to ssms: Distilling quadratic knowledge to subquadratic models”. In: Advances in Neural Information Processing Systems 37 (2025), pp. 31788–31812. find
  7. Chen, Shouyuan et al. “Extending Context Window of Large Language Models via Positional Interpolation”. In: arXiv preprint arXiv:2401.06066 (2023). find
  8. Child, Rewon et al. “Generating long sequences with sparse transformers”. In: arXiv preprint arXiv:1904.10509 (2019). find
  9. Choromanski, Krzysztof et al. “Rethinking attention with performers”. In: arXiv preprint arXiv:2009.14794 (2020). find
  10. Dai, Damai et al. “Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models”. In: arXiv preprint arXiv:2401.06066 (2024). find
  11. Dao, Tri, Dan Fu, et al. “Flashattention: Fast and memory-efficient exact attention with io-awareness”. In: Advances in Neural Information Processing Systems 35 (2022), pp. 16344–16359. find
  12. Dao, Tri and Albert Gu. “Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality”. In: arXiv preprint arXiv:2405.21060 (2024). find
  13. Ding, Jiayu et al. “Longnet: Scaling transformers to 1,000,000,000 tokens”. In: arXiv preprint arXiv:2307.02486 (2023). find
  14. Fedus, William, Barret Zoph, and Noam Shazeer. “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity”. In: Journal of Machine Learning Research 23.120 (2022), pp. 1–39. find
  15. Fu, Tianyu et al. “Moa: Mixture of sparse attention for automatic large language model compression”. In: arXiv preprint arXiv:2406.14909 (2024). find
  16. Gao, Yizhao et al. “SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs”. In: arXiv preprint arXiv:2410.13276 (2024). find
  17. Ge, Suyu et al. “Model tells you what to discard: Adaptive kv cache compression for llms”. In: arXiv preprint arXiv:2310.01801 (2023).
  18. Gu, Albert and Tri Dao. “Mamba: Linear-time sequence modeling with selective state spaces”. In: arXiv preprint arXiv:2312.00752 (2023). find
  19. Guan, Melody Y et al. “Deliberative alignment: Reasoning enables safer language models”. In: arXiv preprint arXiv:2412.16339 (2024).
  20. Guo, Daya et al. “Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning”. In: arXiv preprint arXiv:2501.12948 (2025).
  21. Guo, Mandy et al. “LongT5: Efficient text-to-text transformer for long sequences”. In: arXiv preprint arXiv:2112.07916 (2021).
  22. Guo, Qipeng et al. “Star-transformer”. In: arXiv preprint arXiv:1902.09113 (2019). find
  23. Gupta, Ankit and Jonathan Berant. “Gmat: Global memory augmentation for transformers”. In: arXiv preprint arXiv:2006.03274 (2020). — [arXiv:2006.03274](https://arxiv.org/abs/2006.03274) find
  24. Ho, Jonathan et al. “Axial attention in multidimensional transformers”. In: arXiv preprint arXiv:1912.12180 (2019). — [arXiv:1912.12180](https://arxiv.org/abs/1912.12180) find
  25. Hoffmann, Jordan et al. “Training compute-optimal large language models”. In: arXiv preprint arXiv:2203.15556 (2022).
  26. Jiang, Huiqiang et al. “Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention”. In: arXiv preprint arXiv:2407.02490 (2024).
  27. Katharopoulos, Angelos et al. “Transformers are rnns: Fast autoregressive transformers with linear attention”. In: International conference on machine learning. PMLR. 2020, pp. 5156–5165. find
  28. Kitaev, Nikita, Łukasz Kaiser, and Anselm Levskaya. “Reformer: The efficient transformer”. In: arXiv preprint arXiv:2001.04451 (2020).
  29. Lepikhin, Dmitry et al. “Gshard: Scaling giant models with conditional computation and automatic sharding”. In: arXiv preprint arXiv:2006.16668 (2020). find
  30. Li, Aonian et al. “Minimax-01: Scaling foundation models with lightning attention”. In: arXiv preprint arXiv:2501.08313 (2025). find
  31. Liu, Di et al. “Retrievalattention: Accelerating long-context llm inference via vector retrieval”. In: arXiv preprint arXiv:2409.10516 (2024). find
  32. Liu, Hao and Pieter Abbeel. “Blockwise Parallel Transformer for Large Context Models”. In: arXiv preprint arXiv:2305.19370 (2023). find
  33. Lu, Yi et al. “LongHeads: Multi-Head Attention is Secretly a Long Context Processor”. In: arXiv preprint arXiv:2402.10685 (2024). find
  34. Mercat, Jean et al. “Linearizing Large Language Models”. In: arXiv preprint arXiv:2405.06640 (2024). find
  35. Milakov, Maxim and Natalia Gimelshein. “Online normalizer calculation for softmax”. In: arXiv preprint arXiv:1805.02867 (2018). find
  36. MoonshotAI. Kimi Chat. https://kimi.moonshot.cn/. 2023. ↗
  37. Oren, Matanel et al. “Transformers are multi-state rnns”. In: arXiv preprint arXiv:2401.06104 (2024). find
  38. Peng, Bo, Eric Alcaide, et al. “Rwkv: Reinventing rnns for the transformer era”. In: arXiv preprint arXiv:2305.13048 (2023). find
  39. Peng, Bo, Daniel Goldstein, et al. “Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence”. In: arXiv preprint arXiv:2404.05892 (2024). find
  40. Poli, Michael et al. “Hyena hierarchy: Towards larger convolutional language models”. In: International Conference on Machine Learning. PMLR. 2023, pp. 28043–28078. find
  41. Qiu, Jiezhong et al. “Blockwise self-attention for long document understanding”. In: arXiv preprint arXiv:1911.02972 (2019). find
  42. Rajbhandari, Samyam et al. “Deepspeed-moe: Advancing mixture-of-experts inference and training to power next-generation ai scale”. In: International conference on machine learning. PMLR. 2022, pp. 18332–18346.
  43. Reid, Machel et al. “Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context”. In: arXiv preprint arXiv:2403.05530 (2024).
  44. Roy, Aurko et al. “Efficient content-based sparse attention with routing transformers”. In: Transactions of the Association for Computational Linguistics 9 (2021), pp. 53–68.
  45. Shazeer, Noam et al. “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer”. In: arXiv preprint arXiv:1701.06538 (2017). find
  46. Shoeybi, Mohammad et al. “Megatron-lm: Training multi-billion parameter language models using model parallelism”. In: arXiv preprint arXiv:1909.08053 (2019). find
  47. Sun, Yutao et al. “Retentive network: A successor to transformer for large language models”. In: arXiv preprint arXiv:2307.08621 (2023). find
  48. Tang, Jiaming et al. “Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference”. In: arXiv preprint arXiv:2406.10774 (2024).
  49. Tay, Yi, Dara Bahri, et al. “Sparse sinkhorn attention”. In: International Conference on Machine Learning. PMLR. 2020, pp. 9438–9447. find
  50. Tay, Yi, Mostafa Dehghani, et al. “Efficient transformers: A survey. arXiv”. In: arXiv preprint arXiv:2009.06732 (2020).
  51. Team, Kimi et al. “Kimi k1. 5: Scaling Reinforcement Learning with LLMs”. In: arXiv preprint arXiv:2501.12599 (2025).
  52. Wang, Junxiong et al. “The mamba in the llama: Distilling and accelerating hybrid models”. In: arXiv preprint arXiv:2408.15237 (2024). find
  53. Wang, Sinong et al. “Linformer: Self-attention with linear complexity”. In: arXiv preprint arXiv:2006.04768 (2020). find
  54. Waswani, A et al. “Attention is all you need”. In: NIPS. 2017.
  55. Watson, Jake F et al. “Human hippocampal CA3 uses specific functional connectivity rules for efficient associative memory”. In: Cell 188.2 (2025), pp. 501–514. find
  56. Wu, Yuhuai et al. “Memorizing transformers”. In: arXiv preprint arXiv:2203.08913 (2022). find
  57. Xiao, Guangxuan et al. “Efficient streaming language models with attention sinks”. In: arXiv preprint arXiv:2309.17453 (2023). find
  58. Xiong, Wenhan et al. “Effective Long-Context Scaling of Foundation Models”. In: arXiv preprint arXiv:2309.16039 (2023).
  59. Yang, An et al. “Qwen2. 5 Technical Report”. In: arXiv preprint arXiv:2412.15115 (2024). find
  60. Zaheer, Manzil et al. “Big bird: Transformers for longer sequences”. In: Advances in neural information processing systems 33 (2020), pp. 17283–17297.
  61. Zhang, Michael et al. “LoLCATs: On Low-Rank Linearizing of Large Language Models”. In: arXiv preprint arXiv:2410.10254 (2024). — [arXiv:2410.10254](https://arxiv.org/abs/2410.10254) find
  62. Zhang, Xuan et al. “Simlayerkv: A simple framework for layer-level KV cache reduction”. In: arXiv preprint arXiv:2410.13846 (2024). — [arXiv:2410.13846](https://arxiv.org/abs/2410.13846) find
  63. Zhang, Zhenyu et al. “H2o: Heavy-hitter oracle for efficient generative inference of large language models”. In: Advances in Neural Information Processing Systems 36 (2024). find
  64. Zheng, Zangwei et al. Open-Sora: Democratizing Efficient Video Production for All. Mar. 2024. URL : https://github.com/hpcaitech/Open-Sora. ↗
  65. Zoph, Barret et al. “St-moe: Designing stable and transferable sparse expert models”. In: arXiv preprint arXiv:2202.08906 (2022). find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/lu-2025-moba-mixture-block-attention/