MoBA: Block-Level Sparse Attention at 95% Sparsity, Near-Lossless and Already Serving Kimi

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

cs.LG, cs.AI, cs.CL

2025-02-18

Moonshot ports mixture-of-experts into attention: each query picks top-k blocks of context, hitting 95% sparsity with near-zero loss and a 6.5x prefill speedup at 1M tokens. Already serving Kimi.

What problem this solves

Attention scales quadratically with context length. Once a model has to ingest million-token inputs, which is the direction Kimi, Claude and Gemini are all pushing, the compute bill for full attention gets ugly.

The two existing escape routes each snag. Static sparse attention (sliding window, attention sinks) hard-codes where the model should look; it saves compute but the pattern is fixed and often breaks when the task changes. The other route replaces attention entirely with a linear approximation (Mamba, RWKV), which brings complexity down but whose behaviour on complex reasoning is still unsettled.

MoBA takes a third path: keep the standard Transformer and full softmax attention, but let the model learn where to look instead of prescribing it. "Less structure" is the slogan.

Method

The idea is to port mixture-of-experts (MoE) into attention. Slice a context of length N into n blocks of B = N/n tokens. For each query token, a gate scores every block by relevance:

si = <q, meanpool(K of block i)>

the inner product of the query with the mean-pooled keys of that block. Take the top-k blocks and run standard softmax attention only over them.

A few design choices matter:

Route by block, not by token. Block granularity lets MoBA reuse MoE and FlashAttention kernels, so it actually runs.
The current block is mandatory. Every token always attends to its own block (with causal masking), like a "shared expert" in MoE. This stops the model from sparsifying into oblivion and preserves local context.
No new parameters. MoBA has identical parameters to full attention, so the only variable across experiments is the attention module itself. The gate reuses the existing q and K.
Mixable with full attention. Layer-wise (keep the last few layers full) or stage-wise (90% of tokens through MoBA, 10% through full). No notable loss spikes on switching. This matters for SFT later.

Sliding window and attention sinks are both special cases of MoBA under particular gates, so the authors argue MoBA is strictly more expressive.

Results

A caveat up front: in evaluation MoBA switches back to full attention for generation, so the 6.5x speedup is a prefill number, not end-to-end.

Setting	Metric	MoBA	Full attention
8K context, 81.25% sparse	validation loss gap	within 1e-3 of full	(baseline)
32K context, 95.31% sparse	trailing loss	slightly higher, gap narrows with scale	(baseline)
Llama-8B-1M	MMLU	0.4903	0.4904
Llama-8B-1M	RULER @128K	0.7818	0.7849

In the scaling-law runs (568M to 2.1B params), at 8K context and 81% attention dropped, MoBA stays within 1e-3 of full attention. Push to 32K and 95.31% sparsity and the trailing-token loss runs slightly higher, but the gap shrinks as the model grows.

The largest run continuously pretrains and fine-tunes Llama 3.1 8B to 1M tokens (block 4096, top-k=12, last 3 layers full). It matches its full-attention twin across 16 benchmarks, and needle-in-haystack holds at 1M tokens.

On efficiency, prefilling 1M tokens is up to 6.5x faster. Push to 10M tokens and attention compute drops to a sixteenth (achieved by extending tensor parallelism to the query-head level, broadcasting K/V across distributed heads to dodge the memory wall).

Why it matters

It is one of the few long-context sparse attention methods proven in production, already serving Kimi's long-context requests. Because it is a drop-in replacement, it slots into existing full-attention models without retraining the whole architecture from scratch.

For anyone building long-context applications, near-lossless sub-quadratic attention that handles millions to tens of millions of tokens is exactly the bottleneck. The layer-wise hybrid means you do not have to commit fully to sparsity.

A dose of scepticism: this is a Moonshot technical report. Every number comes from their own setup, there is no independent replication, and there is no head-to-head against other dynamic-sparse methods (Native Sparse Attention, Quest). The baseline is mainly their own full-attention twin.

Limitations

Weak on SFT. Prompt tokens are excluded from the loss, leaving sparse gradients and weak backprop; it needs a layer-wise hybrid (full attention on the last few layers) to recover. So it is not a drop-in for any training pipeline.
MoBA only during prefill in evaluation. Generation falls back to full attention, so the 6.5x reflects one-shot processing of long inputs, not end-to-end inference.
Block granularity is sensitive. At a fixed 75% sparsity, the coarsest cut (2 of 8 blocks) is 1e-2 worse than finer ones; it needs tuning.
Trailing loss still runs slightly high. Shrinking with scale, but trailing tokens are exactly the ones that most need to see the whole context.

Terms

MoE (mixture of experts): splits a layer into several "expert" sub-networks and uses a gate to activate only a few per step. MoBA applies this to attention, where each "expert" is a block of context.
Attention sink: the tendency of models to dump attention onto the first few tokens, often used as a fixed compute-saving pattern.
Sliding window attention: each token only looks at a fixed window of neighbours; the most common static sparse scheme.
Prefill vs generation: ingesting a long input (prefill) and emitting tokens one by one (generation) are different stages with different bottlenecks.
Trailing loss: loss on the final tokens of a sequence; a sharper signal than overall loss for whether a model can "see the whole context then wrap up."
Tensor parallelism: splits a layer's weights or activations across cards; MoBA extends it to the query-head dimension.