Research map / Architectures and efficiency / Attention and transformers

MoBA: Mixture of Block Attention for Long-Context LLMs

NeurIPS 2025 · Architectures and efficiency › Attention and transformers · approach: Sparse attention for long-context LLMs · 8 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes MoBA, a trainable block-sparse attention that routes each query to top-k KV blocks via MoE-style gating, and can switch between full and sparse attention; continually pre-trains Llama-3.1-8B to 1M context.

Problem

Quadratic cost of standard attention limits long-context LLMs; existing sparse methods use predefined structure (sink, sliding window) or linear approximations that are weak on complex reasoning and hard to adapt to pretrained Transformers.

Approach

Apply the Mixture-of-Experts top-k gating to attention: partition KV into blocks, score each block by query–mean-pooled-key inner product, route each query to top-k blocks plus its own current block (causal), so the model learns where to attend with no predefined bias and the same parameters allow switching between full and sparse attention.

Key results

• Llama-8B-1M-MoBA scores 0.7818 vs 0.7849 for full attention on RULER @128K at up to 62.5% sparsity, and is comparable across 16 benchmarks.
• MoBA achieves up to 6.5x speedup in attention when prefilling 1M tokens versus full attention.
• At 10M tokens MoBA reaches a 16x reduction in attention computation time versus FlashAttention.

Measured effects reported in the paper

Surprising findings

• Trailing tokens account for most of the performance gap between full and sparse attention, motivating trailing-token LM loss as a long-context metric.
• Sparse MoBA attention in SFT yields suboptimal loss, which the authors attribute to sparse gradients from loss masking; a layer-wise hybrid with full attention in the last layers significantly reduces SFT loss.

Limitations

• Llama-8B experiments use MoBA only for prefill; generation switches to full attention, so decode-time savings are not claimed.
• SFT with sparse MoBA attention underperforms full attention without the layer-wise hybrid workaround.
• Efficiency results are reported on attention layers only, excluding other layers.

Experimental setup

Datasets: Long-context pretraining corpus (unspecified) SFT data (unspecified)

Benchmarks: AGIEval BBH CEval GSM8K HellaSWAG Loogle Competition Math MBPP MBPP Sanitized MMLU MMLU Pro OpenAI HumanEval

Metrics: accuracy pass@1 LM loss trailing LM loss speedup ratio attention computation time

Baselines: Full attention (Llama-8B-1M-Full) FlashAttention Quest Longheads

Builds on: Mixture of Experts FlashAttention Quest sliding window attention attention sink

Closest papers by meaning

• Training-Free Long-Context Scaling of Large Language Models (ICML 2024 · similarity 0.89)
• Mixture of In-Context Experts Enhance LLMs' Long Context Awareness (NeurIPS 2024 · similarity 0.87)
• Beyond Linear Approximations: A Novel Pruning Approach for Attention Matrix (ICLR 2025 · similarity 0.87)
• LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models (ICLR 2024 · similarity 0.87)
• MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention (NeurIPS 2024 · similarity 0.87)
• Low-Rank Approximation for Sparse Attention in Multi-Modal LLMs (CVPR 2024 · similarity 0.87)