MoBA：用于长上下文大语言模型的块注意力混合机制

完整报告

🚀 重磅推出 MoBA --- 块注意力混合机制

• 可训练的块稀疏注意力：将完整上下文划分为多个块，每个查询 token 自主学习关注最相关的 KV 块，从而高效处理长序列。
• 无参数门控机制：引入一种新颖的无参数 top-k 门控机制，为每个查询 token 选择最相关的块，确保模型仅聚焦于最具信息量的块。
• 全注意力与稀疏注意力之间的无缝切换：MoBA 被设计为全注意力机制的灵活替代方案，可在全注意力和稀疏注意力模式之间实现平滑过渡。

注意：MoBA 需要对现有模型进行持续训练，才能发挥其加速优势。它并非一种即插即用的稀疏注意力解决方案，无法直接应用于预训练模型而无需额外训练。

摘要

扩展有效上下文长度对于推动大型语言模型（LLMs）向通用人工智能（AGI）迈进至关重要。然而，传统注意力机制固有的二次方复杂度增长带来了难以承受的计算开销。现有的方法要么引入强偏置结构，如针对特定任务的 sink 注意力或窗口注意力；要么将注意力机制彻底改造为线性近似，但其在复杂推理任务中的表现仍有待充分探索。

在本工作中，我们提出了一种遵循“少结构”原则的解决方案，允许模型自主决定关注的位置，而非预先设定偏置。我们提出了块注意力混合机制（MoBA），这是一种将专家混合模型（MoE）的思想应用于注意力机制的创新方法。该架构在长上下文任务中表现出色，并具备一个关键优势：能够在全注意力和稀疏注意力之间无缝切换，从而提升效率而不牺牲性能。MoBA 已被部署以支持 Kimi 的长上下文请求，并在 LLM 的高效注意力计算方面展现出显著进展。

我们的代码已开源，详见 MoonshotAI/MoBA。

100万词上下文长度下的评估

环境配置

请注意，当前内核实现依赖 flash-attn==2.6.3 和 torch >= 2.1.0

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

快速入门

我们提供了兼容 Hugging Face Transformers 的 MoBA 实现。

您可以通过 --attn 参数在 moba 和 moba_naive 之间选择不同的注意力后端。

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

实现细节

• moba_naive：基于注意力掩码的朴素实现，旨在帮助理解 MoBA 如何选择对应的块。您可以保存并可视化注意力掩码，以观察块的选择过程。
• moba_efficient：我们面向生产环境的高性能优化实现。与 moba_naive 相比，其速度最高可提升 40 倍（测试条件：序列长度 32K，1 个注意力头，MoBA 块大小 2048，MoBA Topk 3）。我们建议在实际应用中使用此版本。

单元测试

pytest tests/test_moba_attn.py

参考文献

• Llama 实现：huggingface/transformers
• Flash Attention：Dao-AILab/flash-attention

引用

如果您认为 MoBA 对您的工作有所帮助或希望在项目中使用，请引用我们的论文：

@article{lu2025mobamixtureblockattention,
  author = {Enzhe Lu and Zhejun Jiang and Jingyuan Liu and Yulun Du and Tao Jiang and Chao Hong and Shaowei Liu and Weiran He and Enming Yuan and Yuzhi Wang and Zhiqi Huang and Huan Yuan and Suting Xu and Xinran Xu and Guokun Lai and Yanru Chen and Huabin Zheng and Junjie Yan and Jianlin Su and Yuxin Wu and Yutao Zhang and Zhilin Yang and Xinyu Zhou and Mingxing Zhang and Jiezhong Qiu},
  title = {MoBA: Mixture of Block Attention for Long-Context LLMs},
  journal={arXiv preprint arXiv:2502.13189},
  year={2025}
}

MoBA 快速上手指南

MoBA (Mixture of Block Attention) 是一种专为长上下文大语言模型设计的混合块注意力机制。它通过将上下文分块并让模型自主学习关注最相关的 KV 块，实现了高效的长序列处理，同时支持在全量注意力和稀疏注意力之间无缝切换。

重要提示：MoBA 需要对现有模型进行继续训练（Continue Training）才能发挥加速优势，它不是可以直接应用于预训练模型的即插即用方案。

环境准备

在开始之前，请确保您的开发环境满足以下系统要求和依赖版本：

• Python: 推荐 3.10
• PyTorch: torch >= 2.1.0
• Flash Attention: 必须安装 flash-attn==2.6.3 (当前内核实现强依赖此版本)
• 操作系统: Linux (Flash Attention 通常需要在 Linux 环境下编译和运行)

安装步骤

建议使用 Conda 创建独立的虚拟环境以避免依赖冲突。国内用户可使用清华或阿里镜像源加速安装。

1. 创建并激活虚拟环境：

   conda create -n moba python=3.10
   conda activate moba
   

2. 安装依赖与 MoBA： 首先安装指定版本的 flash-attn（可能需要先安装 ninja 和 cuda-toolkit），然后安装 MoBA。

   注：若网络受限，建议配置 pip 国内镜像源（如清华源）。

   # 设置 pip 国内镜像源 (可选，推荐国内开发者使用)
   pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple
   
   # 安装 flash-attn 指定版本 (需确保系统已安装对应的 CUDA 环境)
   pip install flash-attn==2.6.3 --no-build-isolation
   
   # 安装 MoBA
   pip install .
   

基本使用

MoBA 提供了兼容 Hugging Face transformers 的实现。您可以通过命令行参数轻松切换注意力后端。

运行示例

以下命令演示如何使用 MoBA 加载 Llama-3.1-8B 模型并启用 moba 注意力机制：

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

注意力后端说明

通过 --attn 参数可选择不同的实现模式：

• moba (推荐): 对应生产环境的优化版本 (moba_efficient)。相比朴素版本性能提升显著（在特定配置下可达 40 倍加速），适用于实际应用场景。
• moba_naive: 基于注意力掩码的朴素实现。主要用于学习和调试，您可以保存并可视化注意力掩码以观察块选择过程，但不建议用于生产。

# 使用朴素版本进行调试或可视化分析
python3 examples/llama.py --model meta-llama/Llama-3.1-8B --attn moba_naive

验证安装

安装完成后，可运行单元测试确保环境配置正确：

pytest tests/test_moba_attn.py