Title: MoBA: Mixture of Block Attention for Long-Context LLMs
Venue: NeurIPS 2025
Accepted: yes
Keywords: LLM

Sparse Attention

Long Context
Abstract: Scaling the effective context length is essential for advancing large language models (LLMs) toward artificial general intelligence (AGI). However, the quadratic increase in computational complexity inherent in traditional attention mechanisms presents a prohibitive overhead. Existing approaches either impose strongly biased structures, such as sink or window attention which are task-specific, or radically modify the attention mechanism into linear approximations, whose performance in complex reasoning tasks remains inadequately explored.

In this work, we propose a solution that adheres to the ``less structure'' principle, allowing the model to determine where to attend autonomously, rather than introducing predefined biases. We introduce Mixture of Block Attention (MoBA), an innovative approach that applies the principles of Mixture of Experts (MoE) to the attention mechanism. This novel architecture demonstrates superior performance on long-context tasks while offering a key advantage: the ability to seamlessly transition between full and sparse attention, enhancing efficiency without the risk of compromising performance. MoBA has already been deployed to handle actual production workloads with long-context requirements, demonstrating significant advancements in efficient attention computation for LLMs. Our code is available at https://github.com/MoonshotAI/MoBA.


Peer reviews:
comments: SUMMARY:
This paper proposes MOBA, a method inspired by the Mixture of Experts (MoE) paradigm. The key idea is to divide the context into blocks and select the top-k most relevant blocks for attention computation. The approach effectively accelerates the prefill phase in long-context scenarios, though it does not provide significant speedups during decoding. The authors also released a flash_attn_varlen-based implementation to support reproducibility.

STRENGTHS AND WEAKNESSES:
**Pros:**

The paper is clearly written with a well-structured and logical flow, making it easy to read.

The core idea is neat and is shown to be empirically effective.

While the method appears simple, it is actually non-trivial. The authors identify the issue of unstable gradients introduced by the router and mitigate it with a simple hybrid strategy, highlighting the thoughtfulness of the design.

Code is released.

**Cons:**

All attention heads within a group are restricted to attend to the same selected blocks, which may degrade performance. I encourage the authors to add an ablation study to evaluate the impact of this design choice.

The routing mechanism requires compressing all keys within a block into a single vector, which may lead to information loss. For instance, in a synthetic experiment with keys and values as: k₁:v₁, k₂:v₂, ..., kₙ:vₙ — as n increases, can the moba still perform NIAH (for any possible key)? I encourage the authors to include such an experiment and analyze its dependence on block size. If a clear improvement is demonstrated, I would consider increasing my score.

It would be valuable to include results on the GSM-Infinite benchmark.

QUESTIONS:
See cons.

LIMITATIONS:
na
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-06-20

comments: SUMMARY:
This paper introduces MoBA, a sparse attention mechanism for long-context LLMs. It partitions the input into blocks and lets each attention head dynamically select top‑k relevant blocks using mean-pooled key vectors, while ensuring causal masking. The method is fully compatible with standard Transformer architectures and shows strong results up to 1M-token contexts.

STRENGTHS AND WEAKNESSES:
**Strengths:**
1. The proposed method is novel and elegant. Its per-head, parameter-free routing is simple yet effective, avoiding complex structures while retaining flexibility.

2. Experiments are comprehensive and well-executed. MoBA performs comparably to full attention while offering large speedups, and ablations support the design choices.

3. Strong scaling to long contexts. Even at a million-token length, it maintains stable loss and efficiency, which is rare for sparse attention methods.

**Weaknesses:**
- The use of mean pooling for block representation may dilute important token-level signals and lead to suboptimal routing. This is a minor suggestion, however, and does not affect my overall positive assessment of the paper.

QUESTIONS:
See weaknesses.

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-06-24

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work presents Mixture of Block Attention (MoBA), an approach that applies the top-k gating mechanism in Mixture of Experts (MoE) to attention. The targeted scenario is long context Large Language Models (LLMs), where the attention mechanism requires each query token to attend to all key and values tokens, causing a quadratic time complexity. MoBA partitions the context into blocks and lets the model to learn which blocks should a query token attend to via the top-k gating mechanism, along with specific designs to preserve causality in auto-regressive LLMs. Integration with FlashAttention is implemented to achieve high efficiency. The experiment results validate MoBA’s effectiveness on preserving good long context ability while configured with considerable sparsity.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n\n(S1) Timely and important topic. The proposed technique is also well motivated and well designed.\n\n(S2) The in-depth discussion about additional design choices, including the granularity of block partitioning and the layer-wise integration with full attention, is a joy to read. This shows that MoBA is able to generalized to a broader context.\n\n(S3) The experimental results are promising.\n\n---\n\n**Weaknesses:**\n\n(W1) Limited discussion about related works on sparse LLM inference. In particular, QUEST [12] has a quite similar idea as MoBA,. In my humble opinion, MoBA seems to be more generalizable and can be integrated with training. However, such differences should be discussed/compared in details.\n\n(W2) When assessing the efficiency, only the attention mechanism is evaluated. Although the other layers like FFN are unaltered, it would be helpful to see how the proposed technique contributes to the end-to-end efficiency. Additionally, it would be nice to see a figure demonstrating the pareto front of MoBA in terms of efficiency and long-context capability.\n\n(W3) Minor: The models trained from scratch using MoBA are relatively small (from 545M to 2.1B), and those trained with SFT are not so-called “large” (8B).\n\nQUESTIONS:\nPlease address the weaknesses:\n- discuss/compare with QUEST;\n- provide end-to-end efficiency;\n- provide pareto front in terms of efficiency and long-context capability;\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces MoBA, a novel sparse attention architecture inspired by the MoE mechanism. MoBA dynamically selects the most relevant key blocks for each query token, enabling the model to maintain sparsity without losing critical information. The authors implement MoBA using FlashAttention, which significantly improves its computational efficiency in practice. Experimental results further demonstrate that MoBA achieves performance comparable to full attention on both general language tasks and long-context benchmarks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nVery solid paper. The paper is clearly written. The method is designed in a highly intuitive and natural way, without relying on overly complex or engineering-heavy mechanisms. The experiments demonstrate the strong applicability of the proposed method.\n\nWeaknesses:\n\nThe experiments lack comparison with other dynamic sparse attention baseline methods. From the paper, I cannot tell whether MoBA outperforms the related works(performance, efficiency).\n\nIt also lacks ablation studies on the configuration of hyperparameters, e.g., the configuration of hybrid attention. As a result, one cannot tell how to choose the hyperparameters on their own settings, and whether the hyperparameters are sensitive.\n\nQUESTIONS:\n1. I understand that the whole recipe for post-training a long-context LLM is expensive, but I have a question that I cannot infer from existing experiments: Why do you choose hybrid attention with the last 3 layers using full attention? Is this configuration only applicable to the Llama family, or all model families?  If my concerns about the experiments are solved, I am willing to raise my score.\n\n2. Just curious, the primary purpose of this method is to improve the efficiency of attention during training and generation, but it does not reduce the memory footprint of the KV cache, right?\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}