MoBA implements a trainable block-sparse attention mechanism for long-context LLMs, with a parameter-less gating approach and options to switch between full and sparse attention. The project includes a transformers-friendly implementation and a production-ready variant, with unit tests and conda/pip setup instructions.
Collecting history — the radar snapshots this repo daily. The trend line appears after 3 days of data (1 so far).
What it is
MoBA: Mixture of Block Attention for Long-Context LLMs. It offers a trainable block sparse attention approach where the full context is divided into blocks and each query token learns to attend to relevant KV blocks. It also introduces a parameter-less gating mechanism to select blocks and supports switching between full and sparse attention modes.
How it works
The method applies a block-based attention scheme where blocks of keys/values are attended to by queries, guided by a gating mechanism to choose the most relevant blocks. It provides two implementations: moba_naive (attention-mask-based) and moba_efficient (production-ready, optimized for performance). An example claim notes up to 40x speedup for certain configurations when using the efficient variant.
Getting started
Environment notes:
conda create -n moba python=3.10
conda activate moba
pip install .
Quick start example:
python3 examples/llama.py --model meta-llama/Llama-3.1-8B --attn moba
Key options:
- --attn between moba and moba_naive (choice of attention backend)
- moba_naive: naive implementation based on attention masks
- moba_efficient: production-ready, optimized for performance (up to 40x speedup in tested setup)
## Recent releases
- Release notes: none listed.
## Traction
Stars: 2159
Forks: 157
Open issues: 13
## Behind the repo
Linked to MoonshotAI as the organization hosting MoBA. The README references arXiv publication and external code references (huggingface transformers, flash-attention).
## Caveats
License: MIT
Created: 2025-02-17
Last push: 2025-04-03
Notes:
- Requires continue training of existing models to realize acceleration benefits; not a drop-in sparse attention solution for pretrained models without additional training.
- Environment requires flash-attn==2.6.3 and torch >= 2.1.0