MoBA：LLM长文本救星！月之暗面开源新一代注意力机制：处理1000万token能快16倍，已在Kimi上进行验证

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！

AI 在线答疑 -> 智能检索历史文章和开源项目 -> 尽在微信公众号 -> 搜一搜：蚝油菜花

---

「你的LLM还在为长文本「爆内存」？Moonshot突破性方案：千万token推理速度提升16倍！」

大家好，我是蚝油菜花。当同行还在为以下问题头疼时——

• 32k上下文跑批处理，显存占用直接OOM
• 百万token文档检索，响应速度堪比树懒
• 全注意力机制计算量爆炸，稀疏方案又损失精度…

MoBA的横空出世正在改写游戏规则！这个由Moonshot AI开源的下一代注意力机制，通过块稀疏架构+无参数门控实现「鱼与熊掌兼得」：

• ✅ 处理1000万token速度提升16倍
• ✅ 内存占用降低58%且零精度损失
• ✅ 无需微调直接替换现有Attention层

已在Kimi智能助手验证的MoBA，究竟如何突破Transformer瓶颈？接下来我们将深入解析其三大核心黑科技。

快速阅读

MoBA 是一种创新的注意力机制，旨在提高大型语言模型（LLMs）处理长上下文任务的效率。

1. 核心功能：支持块稀疏注意力、无参数门控机制以及全注意力与稀疏注意力的无缝切换。
2. 技术原理：基于细粒度块划分和因果性设计，结合 FlashAttention 和 MoE 的优化技术。

MoBA 是什么

MoBA（Mixture of Block Attention）是由 Moonshot AI 提出的一种新型注意力机制，专为解决大型语言模型（LLMs）在处理长上下文任务时的计算复杂度问题而设计。通过将上下文划分为多个块（block），并引入无参数的 top-k 门控机制，MoBA 能让每个查询 token 动态选择最相关的键值（KV）块进行注意力计算，从而显著降低计算成本，同时保持与传统全注意力机制相当的性能。

这种机制的核心优势在于“少结构”原则，避免了预定义的偏见，使模型能够自主决定关注点。实验表明，MoBA 在处理 100 万 token 的长文本时，速度比传统全注意力机制快 6.5 倍；而在处理 1000 万 token 的超长文本时，速度提升可达 16 倍。

Moonshot AI 已经在 Kimi 平台上验证了 MoBA 的实际效果，并开源了相关代码。

MoBA 的主要功能

• 块稀疏注意力：将上下文划分为多个块，每个查询 token 动态选择最相关的键值（KV）块进行注意力计算，实现高效的长序列处理。
• 无参数门控机制：通过新颖的 top-k 门控机制，为每个查询 token 动态选择最相关的块，确保模型只关注最有信息量的部分。
• 全注意力与稀疏注意力的无缝切换：设计为全注意力的灵活替代品，能在全注意力和稀疏注意力模式之间无缝切换，提高效率而不影响性能。
• 高性能实现：结合 FlashAttention 和 MoE（混合专家模型）的优化技术，显著降低计算复杂度。
• 与现有模型的兼容性：可以轻松集成到现有的 Transformer 模型中，无需大量训练调整。

MoBA 的技术原理

• 因果性设计：为了保持自回归语言模型的因果关系，MoBA 确保查询 token 不能关注未来的块，在当前块中应用因果掩码，避免信息泄露，同时保留局部上下文信息。
• 细粒度块划分与扩展性：支持细粒度的块划分，类似于 MoE 中的专家划分策略，提升了性能，使 MoBA 能扩展到极长的上下文（如 1000 万 token），在长上下文任务中表现出色。

如何运行 MoBA

1. 环境搭建

MoBA 的当前内核实现依赖于 flash-attn==2.6.3 和 torch >= 2.1.0。以下是环境搭建步骤：

conda create -n moba python=3.10
conda activate moba
pip install .

2. 快速启动

MoBA 提供了一个与 transformers 兼容的实现，用户可以通过以下命令快速启动：

python3 examples/llama.py --model meta-llama/Llama-3.1-8B --attn moba

此命令会使用 MoBA 注意力机制运行指定的 Llama 模型。

3. 单元测试

运行以下命令以执行单元测试：

pytest tests/test_moba_attn.py

资源

• GitHub 仓库：https://github.com/MoonshotAI/MoBA

---

❤️ 如果你也关注 AI 的发展现状，且对 AI 应用开发感兴趣，我会每日分享大模型与 AI 领域的开源项目和应用，提供运行实例和实用教程，帮助你快速上手AI技术！

AI 在线答疑 -> 智能检索历史文章和开源项目 -> 尽在微信公众号 -> 搜一搜：蚝油菜花