MoBA: Mixture of Block Attention for Long-Context LLMs
Authors: 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
Yutao Zhang
Zhilin Yang
Xinyu Zhou
Mingxing Zhang
Jiezhong Qiu
Venue: NeurIPS 2025 spotlight ()
Keywords: LLM
Sparse Attention
Long Context
Primary area: deep_learning
TL;DR: MoBA is a dynamic sparse attention mechanism for long-context LLMs. It uses the Mixture of Experts (MoE) principle to allow models to autonomously decide where to focus attention in training, without predefined biases.

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.

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
Yutao Zhang' Zhilin Yang
Xinyu Zhou²
Mingxing Zhang
Jiezhong Qiu
Moonshot AI
2 Tsinghua University
3 Hangzhou Institute of Medicine, CAS
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 window attention which are task-specific, radically modify the attention mechanism into linear approximations, whose performance 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 intro- ducing 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 htttps://github wwt. homm/MoonshotA//MoBA..
Introduction
The pursuit of artificial general intelligence (AGI) has driven the development of large language models (LLMs) unprecedented scales, with the promise of handling complex tasks that mimic human cognition. pivotal capability for achieving AGI is the ability to process, understand, and generate long sequences, which is essential for many applications, from historical data analysis to complex reasoning and decision-making processes. This growing demand for extended context processing can be seen only in the popularity of long input prompt understanding. showcased 3 but also in recent explorations of long chain-of- (5). and OpenAI o1/o3 [6].
by models like Kimi [. Claude (2) and Gemini
thought (CoT) output capabilities Kimi kL.5 (4, DeepSeek-RI
However, extending the sequence length in LLMs is non-trivial due to the quadratic growth in computational complexity associated with the vanilla attention [7]. This challenge has spurred wave of research aimed at improving efficiency without sacrificing performance. One prominent direction capitalizes on the inherent sparsity of attention scores. This sparsity arises both mathematically
"Co-corresponding authors: Xinyu Zhou (ebouxinyyemmonnhoot. en), Mingxing Zhang (zhang
mingring@mail bongrinngmail..einnuuu tsinghua edu. en) and Jiezhong Qiu gggujiizzoon@@hiae cas. cn).
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

from the softmax operation, where various sparse attention patterns have be studied 8 and biologically (9) where sparse connectivity is observed in brain regions related to memory storage.
Existing approaches often leverage predefined structural constraints, such as sink-based [TO] or sliding window attention [1T] exploit this sparsity. While these methods can be effective, they tend to be highly task-specific, potentially hindering the model's overall generalizability. Alternatively, range of dynamic sparse attention mechanisms, exemplified by Quest [12], Minference [8, and RetrievalAttention [3] select subsets of tokens at inference time. Although such methods can reduce computation for long sequences, they do not substantially alleviate the intensive training costs of long context models, making it challenging to scale LLMs efficiently to contexts on the order of millions of tokens. Another promising alternative way has recently emerged in the form of linear attention models, such as Mamba [4), RWKV [LS] [6], and RetNet [17]. These approaches replace canonical softmax-based attention with linear approximations, thereby reducing the computational overhead for long- sequencee processing. However, due to the substantial differences between linear and conventional attention, adapting existing Transformer models typically incurs high conversion or requires training entirely new models from scratch [22]. More importantly, evidence of their effectiveness in complex reasoning tasks remains limited.
costs [8 99 20 21
Consequently, a critical research question arises: How can we design a robust and adaptable attention architecture that retains the original Transformer framework while adhering principle, allowing the model to determine where to attend without relying on predefined biases? Ideally. such an architecture would transition seamlessly between full and sparse attention modes, thus maximizing compatibility with existing pre-trained models and enabling both efficient inference and accelerated training without compromising performance.
"less structure"
Thus, we introduce Mixture of Block Attention (MoBA), a novel architecture that builds upon the innovative principles of Mixture of Experts (MoE) 23 and applies them to the attention mechanism. MoE has been used primarily in the feedforward network (FFN) layers of Transformers [24] 25 226,, but MoBA pioneers its application long context attention, allowing dynamic selection of historically relevant KV blocks for each query token. This approach not only enhances the efficiency of LLMs but also enables them to handle longer and more complex prompts without a proportional increase in resource consumption. MoBA addresses the computational inefficiency of traditional attention mechanisms by partitioning the context into blocks and employing gating mechanism to selectively route query tokens to the most relevant blocks. This block sparse attention significantly reduces the computational costs, paving the way for more efficient processing of long sequences. The model's ability to dynamically select the most informative blocks of keys leads to improved performance and efficiency, particularly beneficial for tasks involving extensive contextual information.
In this paper, we detail the architecture of MoBA, firstly its block partitioning and routing strategy, and secondly its efficient implementation. We further present experimental results that demonstrate MoBA's superior performance in tasks requiring the processing of long sequences. Our work contributes novel approach to efficient attention computation, pushing the boundaries of what is achievable with LLMS in handling complex and lengthy inputs.
2 Method
In this work, we introduce a novel architecture, termed Mixture of Block Attention (MoBA), which extends the capabilities the Transformer model by dynamically selecting historical KV blocks for attention computation. MoBA is inspired by techniques of Mixture of Experts (MoF) and sparse attention. The former technique has been predominantly applied to the feedforward network (FFN) layers within the Transformer architecture, while the latter has been widely adopted in scaling Transformers to handle long contexts. Our method is innovative in applying the MoE principle to the attention mechanism itself, allowing for more efficient and effective processing of long sequences.
Preliminaries: Standard Attention in Transformer. We first revisit the standard Attention (7] in Transformers. For simplicity, we revisit the case where a single query token E R. xd attends to the N key and value tokens, denoting K, V E Nxd, respectively. The standard attention is computed as Aitn(g, K. V)= Softmax kg V. where d denotes the dimension of a single attention head. We focus on the single-head scenario for clarity. The extension to multi-head attention involves concatenating the outputs from multiple such single-head attention operations.

---

Querco
Router
Nowr Peeling
Rlack
Slock
Toyet Guling
Inder
Selected Block
Alto score
ASST A0TT
Attention Outpur
(a)
(b)
Figure 1: Illustration of mixture block attention (MoBA). (a) running example of MoBA; (b)
Integration of MoBA into Flash Attention.
Algorithm MoBA (Mixture of Block Attention) Implementation
Require: Query, key and value matrices Q, K, V E Nxxxx,, MoBA hyperparameters (block size B and top-k); h and d denote # attention heads and head dimension. Also denote n = N/B to be # blocks.
1: ,K, V.} = split_ blocks(K, V, B), where K,, V, E Rxkxd, [n] pool(K, B) Brxkodd
2: K
3: S QK E xhxm
4: M create satteeeccsaaaaaeee mask(N,
5: G topk(S + M,
6: Q' K", V get_self_atm, self__att_blooc((( Q, V)
7: Q" t K'n Vm = index_select moba attn block Q K, G)
O' flash attention antention_vvalen((( Q" K V Tausal=True)
9: 0 m flash attention Irrention_varlen(000 Q K T causal=False)
11: return srinnn_wwtt___ntie wwwooonnnnootmmalaaea O Om)
10: O combine_
2.1 MoBA Architecture
Different from standard attention where each query tokens attend to the entire context, MoBA enables each query token to only attend to a subset keys and values:
MoBA K, V) Softmax (qK(n)) V(T
(1)
where / C [M] is the set of selected keys and values. The key innovation in MoBA is the block partitioning and selection strategy. We divide the full context of length N into n blocks, where each block represents subset of subsequent tokens. Without loss of generality, we assume that the length N is divisible by the number blocks n. We further denote B
be the block size and I, = [i - 1) x B + 1, x B] to be the range of the i-th block. By applying the top-k gating from MoE, we enable each query to selectively focus on a subset of tokens from different blocks, rather than the entire context, namely 1 Ug. Here the model employs gating mechanism, denoted by gi, to select the most relevant blocks for each query token. The MoBA gate first computes the affinity score 8; measuring the relevance between query q and the i-th block, and applies top-k
3

---

gating among all blocks. More formally, the gate value for the i-th block 9i is computed by
8; E Topk ((7)lā E [w]],k)
9 = o otherwise
where Topk(,, Ie) denotes the set containing i highest scores among the affinity scores calculated for each block. In this work, the score is computed by the inner product between and the mean pooling of K [4:] along the sequence dimension s; = (q. (oop_pooooo_]]]
It is important to maintain causality in autoregressive language models, as they generate text by next-token prediction based on previous tokens. This sequential generation process ensures that token cannot influence tokens that come before it, thus preserving the causal relationship. MoBA preserves causality through two specific designs:
No Attention to Future Blocks. MoBA ensures that a query token cannot be routed to any future blocks. By limiting the attention scope to current and past blocks, MoBA adheres the autoregressive nature of language modeling. More formally, denoting pos(g) as the position index of the query 9, we set 8; oo and 9; 0 for any blocks such that pos(g) < < x B.
Causality:
Current Block Attention and Causal Masking. We define the "current block" as the block that contains the query token itself. The routing to the current block could also violate causality, since mean pooling across the entire block can inadvertenlly include future information. To address this, we enforce that each token must be routed to its respective current block and apply causal mask during the current block attention. This strategy not only avoids any leakage of information from subsequent tokens but also encourages attention to the local context. More formally, we set 9. for the block i where the position of the query token pos(g) is within the interval 7,. From the perspective of ofxture-of-Experes (MoE), the current block attention in MoBA is akin to the role of shared experts in modern MoE architectures [27 28), where static routing rules are added when expert selection.
Next, we discuss some additional key design choices of MoBA, such as its block segmentation strategy and the hybrid of MoBA and full attention.
Fine-Grained Block Segmentation. The positive impact of fine line-grained expert segmentation in improving mode performance has been vell-documented in the cexturre-f-Experrr (MoE) litera- ture 27 28 In this work, we explore the potential advantage of applying a similar fine-grained segmentation technique to MoBA. MoBA, inspired by MoE, operates segmentation along the context- length dimension rather than the FFN intermediate hidden dimension. Therefore our investigation aims to determine if MoBA can also benefit when we partition the context into blocks with a finer grain. More experimental results can be found in Section 3.1
Hybrid of MoBA and Full Attention. MoBA is designed to be substitute for full attention, maintaining the same model parameters without any addition or subtraction. This feature inspires us conduct smooth transitions between full attention and MoBA. Specifically, at the initialization stage, each attention layer has the option to select full attention or MoBA, and this choice can be dynamically altered during training if necessary. A similar idea of transitioning full attention to sliding window attention has been studied in [29] More results can be found in Section 3.2
Implementation. We provide high-perffrmance implementation of MoBA, by incorporating optimization techniques from FlashAttention 30 and MoE [31] Figure demonstrates the high efficiency of MoBA, while we defer the detailed experiments on efficiency and scalability to Sec- tion 3.4 Our implementation consists of five major steps: (1) Determine the assignment of query tokens to KV blocks according to the gating network and causal mask; (2) Arrange the ordering of query tokens based on their assigned KV blocks; (3) Compute attention outputs for each KV block and the query tokens assigned to This step can be optimized by FlashAttention with varying lengths: (4) Re-arrange the attention outputs back to their original ordering; (5) Combine the cor- responding attention outputs using online Softmax (i.e., tiling), query token may attend to its current block and multiple historical KV blocks. The detailed algorithmic workflow is formalized in Algorithm| [ and visualized in Figure [b] illustrating how MoBA can be implemented based on MoE and FlashAttention. First, the KV matrices are partitioned into blocks (Line 1-2). Next, the gating score is computed, which measures the relevance between query tokens and KV blocks (Lines 3-7). A top-k operator is applied on the gating score (together with causal mask), resulting in a sparse query-to-K -V-blekkk mapping matrix G to represent the assignment of queries to KV blocks (Line 8). Then, query tokens are arranged based on the query-to- KV-block mapping, and block-wit attention

---

Comor
Figure 2: Efficiency of MoBA. (a) Sequence length speedup evaluation: Computation time scaling of MoBA v.s. Flash Attention with increasing sequence lengths (8K-IM). (b) Fixed Sparsity Ratio scaling: Computation time scaling comparison between MoBA and Flash Attention across increasing sequence lengths (8K10MM) maintaining a constant sparsity ratio of 95.31 % (fixed 64 MoBA blocks with variance block size and fixed top-k=3),
1-17
50
160
Figure 3: Scaling Law (the fitted scaling law curve can be found in Table Appendix). (a) LM loss on validation set (seqlen=8K); (b) trailing LM loss on validation set (seqien=32K, last 2K tokens).
outputs are computed (Line 9-12). Notably, attention to historical blocks (Line 11 and 14) and the current block attention (Line 10 and 13) are computed separately, as additional causality needs to be maintained in the current block attention. Finally, the attention outputs are rearranged back their
original ordering and combined with online softmax (Line 16) 32 33
Overall, MoBA allows the model to adaptively and dynamically focus on the most informative blocks of the context. This is particularty beneficial for tasks involving long documents or sequences, where attending the entire context may be unnecessary and computationally expensive. MoBA's 's ability to selectively attend to relevant blocks enables more nuanced and efficient processing of information.
3 Experiments
3.1 Scaling Law Experiments and Ablation Studies
Scalability w.r.t. LM To assess the effectiveness of MoBA, we perform scaling law experiments by comparing the validation loss of language models trained using either full attention or MoBA. Following the Chinchilla scaling law (34], train five language models of varying sizes with sufficient number of training tokens to ensure that each model achieves its training optimum. Detailed configurations of the scaling law experiments can be found in Table 2 of Appendix. Both MoBA and full attention models are trained with a sequence length of 8K. For MoBA models, we set the block size to 512 and select the top-3 blocks for attention, resulting in a sparse attention pattern with sparsity up to 512x3 81 (2 In particular, MoBA serves as an alternative to full attention, meaning that does not introduce new parameters or remove existing ones. This design simplifies our
8192
Since top-k=3, thus each query token attend most 2 history block and the current block.
5

---

Figure 5: Hybrid of MoBA and full attention. (a) position-wise LM loss for MoBA, full attention, and MoBA/full hybrid training: (b) SFT LM loss w.r.t full attention layers in layer-wise hybrid; (c) SFT trailing LM loss (scqlen=32K, last 2K) w.r.t # full attention layers in layer-wise hybrid.
comparison process. as the only difference across all experiments lies in the attention modules, while all other hyperparameters, including the learning rate and batch size, remain constant. As shown Figure Ba, the validation loss curves for MoBA and full attention display very similar scaling trends. Specifically, the validation loss differences between these two attention mechanisms remain consistent within range of 1e-3. This suggests that MoBA achieves scaling performance that is comparable to full attention, despite its sparse attention pattern with sparsity up to 75%.
Long Context Scalability. However, LM loss may be skewed by the data length distribution 355] which is typically dominated by short sequences. To fully assess the long-context capability of MoBA, we assess the LM loss of trailing tokens (trailing LM loss, in short), which computes the LM loss of the last few tokens in the sequence. We count this loss only for sequences that reach the maximum sequence length to avoid biases that may arise from very short sequences. A detailed discussion on the scale law regarding trailing LM loss can be found in Appendix A.. These metrics provide insights into the model's ability to generate the final portion of sequence, which can be particularly informative for tasks involving long context understanding. Therefore, we adopt a modified experimental setting by increasing the maximum sequence length from 8k to 32k. This adjustment leads to an even sparser attention pattern for MoBA, achieving a sparsity level of to 32768 = 95.31 %. As shown Figure 3b although MoBA exhibits a marginally higher trailing LM loss compared to full attention in all five experiments, the loss gap is progressively narrowing. This experiment implies the long-context scalability of MoBA.
1 512x3
Ablation Study on Fine-Grained Block Segmentation. We carry out series of experiments using a 1.5B parameter model with a 32K context length. The hyperparameters of block size and top-k are adjusted maintain consistent level of attention sparsity. Specif- ically, we divide the 32K context into 8, 16. 32, 64, and 128 blocks, and correspondingly select 2, 4, 8, 16, and 32 blocks, ensuring an tention sparsity of 75%. As shown in Figure 4 MoBA's performance significantly affected by block granularity. Specifically, there a performance difference of 1e-2 between the coarsest- -rarsedt-grained setting (selecting 2 blocks from 8) and the settings with finer granularity (se- lecting 32 blocks from 128). These findings suggest that fine-gnt tre-ggri segmentation appears to be a general technique for enhancing the performance of models within the MoE family, including MoBA.
Figure 4: OnneGGrinned Block Segmentation.
Pareto Frontier in terms of Efficiency and Long-context Capability. We investigate the pareto frontier by studying the relationship between attention sparsity and LM loss. We pre-train an 800M parameter model 32K context length. We divide the 32k context length into different number of blocks (ranging from 16 to 256), and let MoBA select a fixed number of topk=3 blocks. As shown in Table of Appendix, a sparsity around 3/32 turns out to be a good choice to balance efficiency (sparsity) and long-contex: capability.

---

3.2 Hybrid of MoBA and Full Attention
As discussed in Section 2, we design MoBA to be a flexible substitute for full attention, so that it can easily switch from/to full attention with minimal overhead and achieve comparable long-context performance. In this section, we first show seamless transition between full attention and MoBA can be a solution for efficient long-contexi pre-training. Then we discuss the layer-wise hybrid strategy, mainly for the performance of supervised fine-tuning (SFT).
MoBA/Full ull Hybrid Training. We train three models, each with 1.5B parameters, on 30B tokens with a context length of 32K tokens. For the hyperparameters of MoBA, the block size is set to 2048, and the top-k parameter is set 3. The detailed training recipes are follows: (1) MoBA/full hybrid: This model is trained using a two-stage recipe. In the first stage, MoBA is used to train on 90% of the tokens. In the second stage, the model switches to full attention for the remaining 10% of the tokens;
(2) Full attention: This model is trained using full attention throughout the entire training; (3) MoBA: This model is trained exclusively using MoBA.
We evaluate their long-context performance via position-wise language model (LM) loss, which is a fine-grained metric to evaluate Im loss at each position within sequence. Unlike the vanilla LM loss, which is computed by averaging the LM loss across all positions, the position-wise LM loss breaks down the loss for each position separately. Similar metrics have been suggested by [3] (3), who noticed that position-wise LM loss follows a power-law trend relative to context length. As shown in Figure 5a the MoBA-only recipe results in higher position-wise losses for trailing tokens. Importantly, our MoBA/full hybrid recipe reaches a loss nearly identical to that of full attention, which highlights the effectiveness of the hybrid training recipe in balancing training efficiency with model performance. More interestingly, we have not observed significant loss spikes during the switch between MoBA and full attention, again demonstrating the ficxibility and robustness of MoBA.
Layer-wise Hybrid. This flexibility of MoBA encourages us to delve into a more sophisticated strategy the layer-wise hybrid of MoBA and full attention. We investigate this strategy with a particular focus on its application during the supervised fine-tuning (SFT). The motivation for investigating this strategy stems from our observation that MoBA sometimes results in suboptimal performance during SFT, as shown in Figure [5b] We speculate that this may be attributed to the loss masking in SFT prompt tokens are typically excluded from the loss calculation which can pose a sparse gradient challenge for sparse attention methods like MoBA (because it may hinder the backpropagation of gradients, which are calculated only from unmasked tokens, throughout the entire context). To address this issue, we propose a hybrid approach - switching the last several Transformer layers from MoBA to full attention, while the remaining layers continue to employ MoBA. As shown in Figure 5b and Figure Bc, this strategy can significantly reduce SFT loss.
3.3 Large Language Modeling Evaluation
We assess MoBA across variety of real-world downstream tasks, evaluating its performance in comparison full attention models. For ease of verification, our experiments begin with the Llama 3.1 8B Base Model, which is used as the starting point for long-context pre-training. This model, termed
is initially trained with a context length of 128K tokens, and we gradually increase the context length to 256K, 512K, and 1M tokens during the continual pre-training. To ease this transition, we use position interpolation method [37] at the start of the 256K continual pre-training stage. This technique enables us to extend the effective context length from 128K tokens to 1M tokens. After completing the 1M continuous pre-training, MoBA is activated for 100B tokens. We set the block size to 4096 and the K parameter to 12, leading to an attention sparsity of up to 1 1096x12 = 95.31%. To preserve some full attention capabilities, we adopt the layer-wise hybrid LM
strategy the last three layers remain as full attention, while the other 29 full attention layers are switched to MoBA. For supervised fine-tuning. we follow a similar strategy that gradually increases the context length from 32K to M.. The baseline full attention models (termed Llama- -B-IM-Full) also follow a similar training strategy as shown in Figure 8] with the only difference being the use of full attention throughout the process. This approach allows us to directly compare the performance of MoBA with that of full attention models under equivalent training conditions.
In particular, across all evaluation tasks, we conduct two distinct inference configurations for MoBA. In the first configuration (termed as MoBA--Prefll-MoBAAAD throughout both the prefill and generation phases, maintaining the same settings as during training
we apply MoBA consistently
7

---

MoBA-Prefill
MoBA-Prefill Full-Decode
Benchmark
Full
AGIEval [0-shot]
0.5268
0.5144
0.5146
BBH [3-shot]
0.6584
0.6573
0.6589
CEval [S-shot]
0.6298
0,6273
0.6165
GSM8K [5-shot]
0.7324
0.7278
0.7
HellaSWAG [0-shot]
0.8262
0.8262
0.8279
Loogle [0-shot]
0.3333
0.4209
0,4016
Competition Math [0-shot]
0.4454
0,4254
0.4324
MBPP [3-shot]
0,5220
0,5380
0,5320
MBPP Sanitized [0-shot]
0.6498
0.6926
0.6615
MMLJU [0-shot]
0.4968
0.4903
0.4904
MMLJU Pro [S-shot][CoT]]
0.4343
0.4295
0.4328
OpenAl HumanEval [O-sho[[[asss@]]
0.6890
0.6951
0.7012
SimpleQA [0-shot]
0,0460
0,0465
0.0492
TriviaQA [0-shot]
0,5610
0.5673
0.5667
LongBench @ 32K [0-shot]
0.4807
0.4828
0.4821
RULER @ 128K [0-shot]
0.7690
0.7671
0.8031
Table 1: Performance comparison of MoBA-Preeil-MoBA -ecode and BoBA--refill-Full- MoBA has two inference configurations: Decode. Details of the RULER benchmark can be found in Table 5 of Appendix.
and IM-Full Llama-8B-IMe
(a)
Figure 6: Performance of PoBA- Decode the Needle in the Haystack benchmark (upto 1M context length).
and (b) LLama-8B-IM M-
where the last three layers use full attention. For the second configuration (termed as MoBA- Prefil-FFullDeccode) we employ MoBA exclusively during the prefill phase while switching to full attention during generation. As shown in Table MoBA, across both its configurations, exhibits performance levels that are highly comparable to those of the Llama- -ama--BB-a 88--8MMM MM-Full model. It is particularly noteworthy that in the longest benchmark, RULER, where MoBA operates at a sparsity level of up to 1 4096x11 62.5%, MoBA with both two inference configurations nearly matches the 188
performance of full attention, with scores 0.7690/0.7671 compared 0.8031. For context lengths of up to 1M tokens, we evaluate the model using the traditional Needle in the Haystack benchmark. As shown in Figure 6 MoBA demonstrates satisfactory performance, while PEBA--Peilll-Full-DD is slightly better than BreA-PPeeillMMoB-Deeee
3.4 Efficiency and Scalability
The above experimental results show that MoBA achieves comparable performance not only regarding language model losses but also in real-world tasks. To further investigate its efficiency, we compare the forward pass time (prefill time) of the attention layer in two models trained in Section 3..3
Alama-8B-1MM 8B-1M-MoBA and Llama-8B- 8--1---Blll. We mainly focus on the attention layer, all other layers (e.g., FFN) have identical FLOPs in both models. As shown in Figure 2a MoBA is more efficient than full attention across all context lengths, demonstrating a sub-quadratic computational complexity. In particular, it achieves speedup ratio of up to 6.5x when prefilling 1M tokens. Furthermore, the end-to-end forward time of Llama-8B- MM-MoBA and Llama-888 1a-naa-------- 1--1Ml can be found in of Appendix.

---

We also explore the length scalability of MoBA by gradually increasing the context length to 10 million tokens. To maintain a constant attention sparsity, we keep the top-k value and number of MoBA Block fixed while proportionally increasing the block size. To reach the 10M context length, we expanded tensor parallelism [38 toward the query head level, Specifically, we broadcast key and value tensors across distributed query heads, addressing GPU memory limitations while preserving computational efficiency. As shown in Figure 2b) MoBA demonstrates superior efficiency compared to standard Flash Attention when scaling to longer sequences. Specifically, at 10M tokens moba achieves a speedup ratio of 16x reduction in attention computation time. The inset graph in the Figure 2b] focusing shorter sequences (32K to 512K), shows that even though both methods perform comparably at smaller scales, MoBA's computational advantage becomes increasingly evident as sequences grow longer, highlighting its particular strength in processing extremely long sequences.
Overall, the high efficiency of MoBA can be attributed to two key innovations: (1) the block sparse attention mechanism, and (2) the optimized implementation combining MoE and FlashAttention, as described in Section 2. These techniques effectively address the quadratic complexity limitation of full attention, reducing it to a more economical sub-quadratic scale.
4 Related Work
The development of efficient attention [39] mechanisms has been critical area of research in the field of natural language processing, particularly with the rise of Large Language Models (LLMs). We review related work in this section:
have been dedicated to the design of static attention patterns in LLMs. Their choices of static attention patterns can encompass strided and fixed attention, window attention, global token attention, random attention, dilated attention, block sparse attention, or any combinations of them. In the realm of multimodal models, static sparse attention mechanisms have also been developed, such as axial attention (48 for 2D images and spatial-tempores attention 449 for 3D videos.
Static Sparse Patterns:
Significant efforts
Dynamic Sparse Patterns: Different from static patterns, dynamic sparse attention techniques adap- tively determine which tokens to attend. Reformer [50] and Routing Transformer [51] respectively
employ locality-sensitive hashing (LSH) and K-means to cluster tokens, and attend to clusters rather
than the full context. Memorizing Transformers [22] and Unlimiformer
[53] dynamically attend to
tokens selected by the k-nearest-neighhor (kNN) algorithms. CoLT5 [54] designs a routing modules to select the most important queries and keys. Sparse Sinkhorn Attention [55] learns to permute
blocks from the input sequence, allowing dynamic block sparse attention computation. Notably, DeepSeek NSA (56) is most recent work that studies dynamic sparse attention, whose selected attention branch is similar to MoBA's block selection mechanism. We believe that the emergence of dynamic sparse attention underscores its significance for the future of long-context LLMs
Training-free Sparse Attention. There are also strategies designed to incorporate sparse attention to model inference. During model inference, the complete prompt can be utilized for attention profiling, which allows for the exploration of more intricate sparse attention patterns. For instance, MoA [57]
Minference [8] and SeerAttention [58] have investigated sparse attention configurations such as A-shape, vertical-slash, and dynamic block sparsity. Besides, considerable work has been dedicated to compressing and pruning the KV-cache to achieve a balance between the quality and speed. Notable
efforts in this area include [59] 10 60
61 12 62)
Beyond Traditional Attention Architecture: Another line of research investigates novel architec-
tures that deviate from the conventional attention mechanism. Studies in this domain have explored architectures that are inspired by Convolutional Neural Networks (CNNs), Recurrent Neural Networks
(RNNs), State Space Models (SSMs), or Linear Attention, such as 63 64
68
5 Conclusion
In this paper, we introduce Mixture of Block Attention (MoBA), a novel attention architecture inspired by the principles of Mixture of Experts (MoE) that aims to enhance the efficiency and scalability of large language models (LLMs) for long-context tasks. MoBA addresses the computational challenges

---

associated with traditional attention mechanisms by partitioning the context into blocks and employing dynamic gating mechanism to selectively route query tokens to the most relevant KV blocks. This approach not only reduces computational complexity but also maintains model performance. Moreover, it allows for seamless transitions between full and sparse attention. Through extensive experiments, we demonstrated that MoBA achieves performance comparable to full attention while significantly improving computational efficiency. Our results show that MoBA can scale effectively to long contexts, maintaining low LM losses and high performance on various benchmarks. Additionally, MoBA's flexibility allows to be integrated with existing models without substantial training cost, making it a practical continual pre- training solution for enhancing long-cor congccotext capabilities in LLMs. In summary, MoBA represents a significant advancement in efficient attention, offering a balanced approach between performance and efficiency.
Limitations. Firstly, the router of MoBA is parameter-free. This design choice may limit its ability to capture more nuanced long-context patterns. Future work could explore the development of learnable router, which allows more sophisticated and adaptive block selection mechanisms. Secondly, this work focuses text inputs. Extending MoBA to other modalities, such as images, audio, multimodal inputs, could significantly broaden its applicability. Finally, investigating MoBA's potential for improving generalization in complex reasoning tasks, such as long chain-of-thought (CoT) reasoning, could provide further insights into its capabilities.
Broader Impact. The development of MoBA has the potential to significantly enhance the efficiency and scalability of large language models (LLMs), enabling more powerful and resource- resource-efficient processing of long-context tasks. This advancement could positively impact various applications, such as natural language understanding, complex reasoning. and decision-making processes, leading to more capable Al systems. However, is important to consider potential negative societal impacts, such as the misuse of enhanced language models for generating disinformation or the unintended consequences of deploying models with biases. To mitigate these risks, future work should focus on developing robust safeguards and ensuring transparency model deployment and usage.
10

---

Acknowledgments and Disclosure of Funding
Mingxing Zhang supported by Beijing Natural Science Foundation (L.252014. Jiezhong Qiu is supported by NSFC 62306290.
References
[1] MoonshotAI. Kimi chat. https::/kiiii.moonnsset cn/ 2023.
[2] Anthropic. Introducing 100k context windows.
cookktontttt-viii
2023.
[3] Machel Reid, Nikolay Savinov, Denis Teplyashin, Dmitry Lepikhin, Timothy Lillicrap, Jean- baptiste Alayrac, Radu Soricut, Angeliki Lazaridou, Orhan Firat, Julian Schrittwieser, et al. Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context, arXiv preprint wxxxx:2403300558 2024.
[4] Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al. Kimi kl. Scaling reinforcement learning with Ilms. arXiv preprint arXiv:2301.12599 2025.
[5] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501. 2025.
[6] Melody Y Guan, Manas Joglckar, Eric Wallace, Saachi Jain, Boaz Barak, Alec Heylar, Rachel Dias, Andrea Vallone, Hongyu Ren, Jason Wei, al. Deliberative alignment: Reasoning enables safer language models. arXiv preprint arXiv:2412. 16339, 2024.
[7] A Waswani, N Shazeer, N Parmar, J Uszkoreit, L Jones, A Gomez, L Kaiser, and Polosukhin. Attention is all you need. In NIPS, 2017.
[8] Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H Abdi, Dongsheng Li, Chin- Yew Lin, et al. Minference 1.0: Accelerating pre-filling for long-context Iims via dynamic sparse attention. arXiv preprint arXiv:2407.024900 2024.
[9] Jake F Watson, Victor Tavakoli, Johann G Danzl, Matthias Tomschik, Karl Rössler, and Peter Jonas. Human hip- pocampal ca3 uses specific functional connectivity rules for efficient associative memory. Cell, 188(2):501-514 2025.
Rebecca J Morse-Mora, Andrea Navas-Olive, Mojtaba R
[10] Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint 2xXi:23309.14533 2023.
[11] Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long dong-document transformer. arXiv preprint arXiv:2004.05150, 2020.
[12] Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikci, and Song Han. Quest: Query-aware sparsity for efficient long-context Ilm inference. arXiv preprint arXiv:2406. 10774, 2024.
[13] Di Liu, Meng Chen, Baotong Lu, Huiqiang Jiang, Zhenhua Han, Qianxi Zhang, Qi Chen, Chengruidong Zhang, Bailu Ding. Kai Zhang, al. Retrievalaateenion:: Accelerating long- context Ilm inference via vector retrieval. arXiv preprint arXiv:2409..00116, 2024.
[14] Tri Dao and Albert Gu. Transformers ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.21066/ 2024.
[15] Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv; Reinventing mos for the transformer arXiv preprint 23X5v:2305.13045 2023.

---

[16] Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Xingjian Du, Teddy Ferdinan, Haowen Hou, al. Eagle and finch: Rwkv with matrix-valued states and dynamic recurrence. arXiv preprint arXiiv2444.055000
2024.
[17] Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXivv2307..8861,, 2023.
[18] Jean Mercat, Igor Vasiljevic, Sedrick Keh, Kushal Arora, Achal Dave, Adrien Gaidon, and Thomas Kollar. Linearizing large language models. arXiv preprint Siii::200506644)) 2024.
[19] Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the Ilama: Distilling and accelerating hybrid models. arXiv preprint arXiv:2408. /5237, 2024.
[20] Aviv Bick, Kevin Li, Eric Xing, Zico Kolter, and Albert Gu. Transformers ssms: Distilling quadratic knowledge to subquadratic models. Advances in Neural Information Processing Systems, 37:31788-3191 2025.
[21] Michael Zhang, Simran Arora, Rahul Chalamala, Alan Wu, Benjamin Spector, Aaryan Singhal, Krithik Ramesh, and Christopher Ré. Lolcats: On low-rank linearizing of large language models. arXiv preprint arXiv:2410.10254, 2024.
[22] Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, et al. Minimax-01: Scaling foundation models with lightning attention. urXiv preprini arXiv:2501.083/1 2025.
[23] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously large neural networks: The sparsely- gated mixture-of-expee layer. arXiv preprint arXiv: / 011.700110065338 2017.
[24] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang. Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. arXiv preprint arXiv:2006. 16668, 2020.
[25] William Fedus, Barret Zoph. and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity, Journal of Machine Learning Research, 23(120):1-3, 2022.
[26] Barret Zoph, Irwan Bello, Sameer Kumar, Nan Du, Yanping Huang, Jeff Dean, Noam Shazeer, and William Fedus. St-moe: Designing stable and transferable sparse expert models. arXiv preprint arXiv:2202.08900 2022.
[27] Damai Dai, Chengqi Deng, Chenggang Zhao, RX Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y Wu, al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-expere language models. arXiv preprint 200i 166066, 2024.
[28] An Yang, Baosong Yang. Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.1 15115, 2024.
[29] Xuan Zhang, Cunxiao Du, Chao Du, Tianyu Pang, Wei Gao, and Min Lin. Simlayerkv: A simple framework for layer-level kv cache reduction. arXiv preprint arXiv:2410.13844 2024.
[30] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficiint exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
[31] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. Deepspeed-moe: Advancing mixture-of- experts inference and training to power next-generation ai scale. In International conference on machine learning. pages 18332 18346. PMLR, 2022.
[32] Maxim Milakov and Natalia Gimelshein. Online normalizer calculation for softmax. arXiv preprint 1805002600 2018
12

---

[33] Hao Liu and Pieter Abbeel. Blockwise parallel transformer for large context models. arXiv preprint TrXiv:2305.193700 2023.
[34] Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbi, Aidan Clark, al. Training compute-optimal large language models. arXiv preprint arXiv:2203. 2022.
[35] Chenxin An, Jun Zhang, Ming Zhong. Lei Li, Shansan Gong, Yao Luo, Jingjing Xu, and Lingpeng Kong. Why does the effective context length of Ilss fall short? arXiv preprint arXiv:2410.18745, 2024.
[36] Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, Madian Khabsa, Han Fang. Yashar Mehdad, Sharan Narang. Kshitiz Malik, Angela Fan, Shruti Bhosale, Sergey Edunov, Mike Lewis, Sinong Wang, and Hao Ma. Effective congecontext scaling of foundation models. arXiv preprint arXiv:2309, 16039, 2023.
[37] Shouyuan Chen, Sherman Wong, Liangjian Chen, and Yuandong Tian. Extending context window of large language models via positional interpolation. arXiv preprint arXiv:2401.0600 06066, 2023.
[38] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron- Training billi-n parameter language models using model parallelism. arXiv preprint
08053, 2019.
[39] Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. arxiv. arXiv preprint arXiv:2009.06732. 2020.
[40] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 1050, 2019.
[41] Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. Star-transformer. arXiv preprint arXiv: 2019.
[42] Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding. arXiv preprint arXiv:1911.02972. 2019.
[43] Ankit Gupta and Jonathan Berant, Gmat: Global memory augmentation for transformers. arXiv preprint arXiv: 2006.030043 2020.
[44] Joshua Ainslie, Santiago Ontanon, Chris Alberti, Vaclav Cvicek, Zachary Fisher, Philip Pham. Anirudh Ravula, Sumit Sanghai, Qifan Wang, and Li Yang. Ete: Encoding long and structured inputs in transformers. arXiv preprint arXiv:2004.0844 00040440........ 2020.
[45] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santi- ago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang. al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17299 2020.
[46] Mandy Guo, Joshua Ainslic, David Uthus, Santiago Ontanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. Longt5: Efficient text-to-text transformer for long sequences. arXiv preprint arXiv:2112.07911 2021.
[47] Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, and Furu Wei. Longnet: Scaling transformers to 1,000,000,000 tokens. arXiv preprint arXiv:2307.02488 2023.
[48] Jonathan Ho, Nal Kalchbrenner, Dirk Weissenborn, and Tim Salimans. Axial attention in multidimensional transformers. arXiv preprint arXiv:1912.0 12180, 2019.
[49] Zangwei Zheng, Xiangyu Peng, Tianji Yang, Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all, March 2024.
13

---

[50] Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint 000::200010441,,, 2020.
[51] Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, 2021.
[52] Yuhuai Wu, Markus N Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing trans- formers. arXiv preprint arXiv:2203.089133 2022.
[53] Amanda Bertsch, Uri Alon, Graham Neubig, and Matthew Gormley. Unlimiformer: Long-range transformers with unlimited length input. Advances in Neural Information Processing Systems, 36, 2024.
[54] Joshua Ainslie, Tao Lei, Michiel de Jong. Santiago Ontañón, Siddhartha Brahma, Yury Zemlyan- skiy, David Uthus, Mandy Guo, James Lee-Thorp. Yi Tay, et al. Colt5: Faster long-range transformers with conditional computation. arXiv preprint arXiv:2303.0977 230.00953. 2023.
[55] Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, and Da- Da-Cheng Juan. Sparse sinkhorn attention. PMLR, 2020.
In International Conference on Machine Learning, pages
[56] Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv preprini arXiv:2502. 1089, 2025.
[57] Tianyu Fu, Haofeng Huang, Xuefei Ning, Genghan Zhang, Boju Chen, Tianqi Wu, Hongyi Wang, Zixiao Huang, Shiyao Li, Shengen Yan, et al. Moa: Mixture of sparse attention for automatic large language model compression. arXiv preprint arXiv:2406. 14909, 2024.
[58) Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang, and Mao Yang. Seerattention: Learning intrinsic sparse attention in your llms. arXiv preprint arXiv:2410. 13276, 2024.
[59] Zhenyu Zhang. Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, et al. H2o: Heavy-hitter oracle for efficient generative inference of large language models. Advances in Neural Information Processing Systems, 36, 2024.
[60] Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers are multi-state mns. arXiv preprint arXiv:2401.06114 2024.
[61] Suyu Ge, Yunan Zhang, Liyuan Liu, Minjia Zhang, Jiawei Han, and Jianfeng Gao. Model tells you what to discard: Adaptive kv cache compression for llms. arXiv preprint 1rX.v:2310.01801, 2023.
[62] Yi Lu, Xin Zhou, Wei He, Jun Zhao, Tao Ji, Tao Gui, Qi Zhang, and Xuanjing Huang. Longheads: Multi-heat attention is secretly long context processor. arXiv preprint 1086:, 2024.
[63] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International conference on machine learming, pages 5156-5165. PMLR, 2020.
[64] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pages 28043-28078. PMLR, 2023.
[65] Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
[66] Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint 200X::2000000768,, 2020.
14

---

[67] Albert Gu and Tri Dao. Mamba: Linear-tine time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752, 2023.
[68] Songlin Yang, Bailin Wang. Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule over sequence length. arXiv preprint wrXiv:2406.064844 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We show Mixture of Block Attention (MoBA), an innovative approach that applies the principles Mixture of Experts (MoE) to the attention mechanism, achieves significant advancements in efficient attention computation for LLMs.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contri- butions made in the paper and important assumptions and limitations. A No or NA answer this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected generalize to other settings.
It is fine include aspirational goals motivation long it is clear that these goals not attained by the paper.
Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
We discuss the main limitations in Section 5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, c.g., the approach was only tested on few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-ter system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Question:
Answer: [NA]
16

---

Justification:
This paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Answer: [Yes]
Justification: We list the pseudocode of MoBA in Algorithm 1 and include the code in the Supplementary Material. Implementation details and experimental settings are described in Section 2. and Section respectively.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also provided via detailed instructions for how to replicate the results, access to hosted model (e.g., in the case of a large language model), relcasing a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submissions provide some reasonable avenue for reproducibility, which may depend the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g.. to registered users), but should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes]
17

---

Justification:
We include our code in Supplementary Material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines battps://nipp pttttss//nips..ccpubbii cc/plicc///
quiddes//ooeeSuunissii..ooo ooooliicciiiiiioiiiioooooo for more details.
While we encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https //nips.cm httpsscccppbcccpceeeeeeeeeeeeece Sewuddee/oodeSuiiieeeeeeeeeeeeee wwt.immmsssoooooooooooo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide the training and test details in Section and Appendix Section
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance the experiments?
Question:
Answer: [No]
Justification: Our experiments follow the standard protocol in long-context LLM tasks, and thus did not report statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main
claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors)
It should clear whether the error bar the standard deviation or the standard error the mean.
18

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification: We confirm that all necessary computational resource information provided in Section and Section A.3 including the GPU types, cluster network architecture, and training durations for our various experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Question:
Answer: [Yes]
Justification: We confirm the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consideration a
due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss the broader impacts of our work in Section 5
Guidelines:
The answer NA means that there is societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
19

---

The conference expects that many papers will be foundational research and not tied to particular applications. let alone deployments. However, if there direct path any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses addition attacks, mecha- nisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks, because the main contribution of this work is novel architecture, therefore do not release data models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We properly cite the original paper that produced the code packages datasets that are used this paper.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode licenses for some datasets. Their licensing guide can help determine the license of a dataset.
com/dataeets has curated
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's creators.
20

---

13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Our new assets include the code this work, which will be released under The MIT License (MIT).
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodl as part of their sub- missions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LIMM usage
an important, original, or non- standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question: Does the paper describe the usage of LLMs
21

---

Answer:
[NA]
Justification:
This research does not involve LLMs as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve
LLMs as any important, original, or non-standard components.
Please refer to our LLM policy what should or should not be described.
cc/Confereeces//2222 /20cce/ LLM)) for
22

---

A Technical Appendices and Supplementary Material
A.1 Scaling Law
For our validation experiments in the Section 3.1 and Section 3.2, we utilized a distributed computing infrastructure consisting of server nodes, each equipped with NVIDIA H800 GPUs (64 GPUs in total). Within each node, the GPUs were interconnected via NVLink for high-bandwidth intra- node communication, while Remote Direct Memory Access (RDMA) was employed for efficient inter-node communication across the cluster. The learning rate schedule followed the optimal decay pattern established during the training phase. Each experiment was terminated upon completion of the expected number of training tokens.
and
No-Emb Model Param
Head
Layer
Hidden
Training Token
Block size
TopK
545M 822M
14
14
1792
10.8B
512
3
16
16
2048
15.3B
512
3
1.1B
18
2304
20.6B
512
1.5B
20
20
2560
27.4B
512
3
2.1B
22
22
2816
36.9B
512
3
Table 2: Configuration of Scaling Law Experiments
L(C)
MoBA
Full
LM loss (seglen=8K)
2.625 x C q-0.063
2.622 x C' 0.063
Trailing LM loss (seqlen=32K, last 2K)
1.546 C-0.108
1.464 C 0.007
Table 3: Fitted scaling law curve.
A.2 Long Context Scalability
To address the bias in natural data distribution that favors short contexts, we strategically segmented the overall sequences into discrete segments based on their actual positions. For example, the segment spanning positions 30K-32K exclusively reflects losses associated with documents exceeding 30K context lengths and also masks the positions from 30K to 32K. This approach ensures a more balanced and representative evaluation across different context lengths. our exploration of long- context scalability, made pivotal discovery: the trailing tokens account for the majority of the performance discrepancy between the full context baseline and the newly proposed sparse attention architectures. Consequently, we streamlined the long-context scaling process by focusing on trailing token scaling. This not only simplifies the computational requirements but also significantly enhances the efficiency and effectiveness of investigating long-context scenarios. This finding holds substantial implications for the development of more efficient and scalable attention mechanisms the future.
A.3 Large Language Modeling Evaluation
For our validation experiments in the Section 3.3 we utilized distributed computing infrastructure consisting of 128 GPU server nodes, each equipped with 8 NVIDIA H800 GPUs (1024 GPUs in total). Within each node, the GPUs were interconnected via NVLink for intra-node communication. while Remote Direct Memory Access (RDMA) was employed for inter-node communication. We trained our model following the recipe outlined in Figure 8 requiring over 30 days to complete training recipe of the Full model with a context length of 1M tokens. The MoBA variant achieved significantly reduced training time by leveraging the same pretrained checkpoint established prior to the IM context length training phase.
23

---

(a) Scaling law (0-2k)
(b) Scaling law (2-4k)
(c) Scaling law (4-6k)
(d) Scaling law (6-8k)
(e) Scaling law (8-104 (8-k0k)
(f) Scaling law (10-128
(g) Scaling law (12-14k)
(h) Scaling law (14-16k)
Figure 7: Scaling laws for positions 0-16k
(i) Scaling law (16-186 (16-18k) 18k)
i Scaling law (18-20) 20k)
(k) Scaling law (20-22k)
Scaling law (22-24k)
(m) Scaling law (24-26k)
(n) Scaling law (26-28k)
(o) Scaling law (28-30k)
(p) Scaling law (30-32k)
Figure 7: Scaling laws for positions 16- 16-32k
24

---

LM Loss Position Range
MoBA
Full
3.075 C00.77
3.068 C 0.078
2K 4K
2.415 x C C000884
2.411 X C 0.00
4K 6K
2.085 C C-0.081
2.077 C
6K 8K
1.899 X C Q00..92
1.894 X C' D.092
8K 10K
1.789 C C
1.774 C
0,089
10K 12K
1.721 X C C-0.092 0.002
1.697 X C - 0.087
12K 14K
1.670 X C-0.069
1.645 C
14K 16K
1.630 C-00889 0.089
1.600
D.087
16K 18K
1.607 X C 0.000
1,567 X C DDD.087
18K 20K
1.586 C 0000091
1.542 C D.008
20K 22K
1.571 X C Q-000933
1.519 X C C00008
22K 24K
1.566 C
1.513 C
24K 26K
1.565 X C Q0.099
1.502 X C 0.086
26K 28K
1.562 C Q00.
1.493 C
28K 30K
1.547 X Q-0.097
1.471 X 0.091
30K 32K
1.546 X C-0.108
1.464 X C 0.097
Table 4: Loss scaling with different positions
MoBA-Prefill MoBA- Decode
MoBA-Prefill Full- Dell-Decode
Ruler Tasks
Full
Avg.
0.769
0.767
0.803
niah multikey
0.982
0.982
1.000
niah multikey_2 2
0,762
0.754
0.992
niah fiah_multikey..
0.944
0.948
0.994
niah_multique multmuuutrrrry niah taihmmuutivleee niah single niah single 2
0.819
0.812
0.765
0.558
0.559
0.571
1.000
1.000
1.000
0.992
0.984
1.000
niah single 3 ruler_ewe
0.966
0.958
0.982
0.009
0.010
0.010
ruler_fwe
0.882
0.881
0.905
aale__a_hhotpp
0,510
0,518
0.542
ruler ailer_a__sop qa_squad
0.604
0.602
0.753
ruler_ vt
0.968
0.965
0.926
Table 5: Ruler Test Result @ 128KK
Context Length
Full Prefill Time (sec)
MoBA Prefill Time (sec)
Speed-up
256 K
5.67
3.55
1.60x
512 K
19,59
8.40
2.33
M
73.97
22.15
3.34x
Table 6: End-to-end Efficiency of Llam--SB-IM-MoBA and Llamma-IM-FE (Note that Llama-8B- IM-MoBA is a hybrid model with the last 3 layers to be full attention).
#Blocks
8
32
128
512
LM
2.48
2.48
2.47
2.53
2.60
Table 7: The relationship between attention sparsity and LM loss.
25

---

250K Continua
512K Continue
1M Continue
LLama3.1 8B
instruct Model
IM SFT
2564 SFT
S2K SFT
Figure 8: The continual pre-training and SFT recipes.
26


Reviews:
[{"review_id": "UYf8bgubv7", "reviewer": "Reviewer_2qxm", "summary": "This 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.", "questions": "1. 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?", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\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.", "was_revised": true, "final_justification": "Since the authors did not provide experimental results to address my concerns, but only stated that they would include them in a future manuscript, I will maintain my score. However, as I am not particularly familiar with this field, please primarily refer to the opinions of the other reviewers.", "rebuttal": "Thank you for your positive feedback and insightful questions! Below we clarify the three main concerns.\n\n**(W1) Regarding MoBA vs other dynamic sparse attention baseline methods:**\n\nAlthough we do not include comparison with other dynamic sparse attention (e.g., DeepSeek's  NSA) in our manuscript, we would like to refer to an independent benchmark by Tilderesearch (\"Sparsity is Cool\", Dhruv Pai et al.[1])  which rigorously compares MoBA with DeepSeek-NSA on efficiency, perplexity, and attention-pattern. We will cite and discuss this report in the next version of our manuscript.\n\n\n[1] Dhruv Pai,  Timor Averbuch,  Mason Wang,  Ben Keigwin. \"Sparsity is Cool\". Tilderesearch Tech blog.\n\n**(W2 and Q1) Regarding hyperparameters, e.g., the configuration of hybrid attention:**\n\nWe briefly discuss the motivation of switching the last layers to full attention in section 3.2 (Line 222-232).  In summary, we hypothesize that replacing the last few MoBA layers with full-attention during SFT alleviates the sparse-gradient problem (especiall for long prefill tasks) caused by loss masking, enabling gradients from the unmasked tokens to propagate along the entire sequence and thus achieving better SFT loss.\n\nTo further answer your question, we ablate another layer hybrid strategies during SFT --- switching the first, the middle, and the last layer from MoBA to full attention, and we observe a performance drop in SFT loss (comparing to switching the last 3 layers to full attention). We will add more ablation studies in the next version of our manuscript.\n\n**(Q2) Regarding KV-Cache:**\n\nMoBA does not reduce KV-cache; it only reduces the compute cost of the attention matrix.  The memory footprint remains identical to full attention."}, {"review_id": "SGppeggkje", "reviewer": "Reviewer_a1Px", "summary": "This 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.", "questions": "Please 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;", "limitations": "yes", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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).", "was_revised": true, "final_justification": "My concerns have been addressed. I suggest adding the discussion and results into the paper.", "rebuttal": "Thank you for your positive feedback and insightful questions! Below we clarify the four main concerns.\n\n**(W1 and Q1) Regarding discussion about related works on sparse LLM inference (in particular QUEST [12]):**\n\nQUEST pioneered inference-time dynamic sparse attention and also employs top-k block selection.  The key difference is that QUEST uses more sophisticated block representation which combining max and min pooling, whereas MoBA currently uses simple mean pooling.  \n\nTo further answer your question, we ablate adding max- and min-pooled blocks on top of mean pooling and observed a small but consistent LM loss improvement, confirming that richer and learnable block representations are promising. We will add more details in the next version of our manuscript.\n\n**(W2 and Q2) Regarding end-to-end efficiency:**\n\nThe end-to-end efficiency of Llama-8B-1M-MoBA and Llama-8B-1M-Full can be found in the following table:\n\n| Context | Full Prefill Time (s) | MoBA Prefill Time (s) | Speed-up |\n|---------|--------------------------------|------------------------------------|----------|\n| 256 K   | 5.67                           | 3.55                               | 1.60×    |\n| 512 K   | 19.59                          | 8.40                               | 2.33×    |\n| 1 M     | 73.97                          | 22.15                              | 3.34×    |\n\nNote that Llama-8B-1M-MoBA is a hybrid model with the last 3 layers to be full attention.\n\n**(W3) Regarding scaling law:**\n\nWe appreciate the valuable suggestion. Limited by computing resources, the largest model we trained for scaling law experiments is 2.1B, and we agree that training larger model can help better understanding towards MoBA.\n\n\n**(Q3) Regarding pareto front in terms of efficiency and long-context capability:**\n\nWe investigate the pareto frontier by studying the relationship between attention sparsity and LM loss. We pre-train an 822M Non-Embedding parameter model in 32K context length. We divide the 32k context length into different number of blocks (ranging from 4 to 512), and let MoBA select a fixed number of topk=3 blocks. In our experiments, a sparsity around 3/32 turns out to be a good choice to balance efficiency (sparsity) and long-context capability.  We will add more details in the next version of our manuscript.\n\n| Number of Blocks | 4    | 8    | 32   | 128  | 512  |\n|------------------|------|------|------|------|------|\n| LM loss          | 2.48 | 2.48 | 2.47 | 2.53 | 2.60 |"}, {"review_id": "2HxpxJkxKN", "reviewer": "Reviewer_cZex", "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.", "questions": "See weaknesses.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths:**\n1. The proposed method is novel and elegant. Its per-head, parameter-free routing is simple yet effective, avoiding complex structures while retaining flexibility.\n\n2. Experiments are comprehensive and well-executed. MoBA performs comparably to full attention while offering large speedups, and ablations support the design choices.\n\n3. Strong scaling to long contexts. Even at a million-token length, it maintains stable loss and efficiency, which is rare for sparse attention methods.\n\n**Weaknesses:**\n- 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.", "was_revised": true, "final_justification": "The response addresses my concern, and I maintain my positive rating.", "rebuttal": "Thank you for your positive feedback and insightful questions! Below we clarify the main concerns.\n\n\n\n**(W1) Regarding the use of mean pooling for block representation**\n\nThank you for your insightful questions! Besides mean pooling, we do ablate another strategy --- in addition to the top block selected by the mean pooling gate, we also select top blocks in terms of max pooling and min pooling respectively. It does lead to slightly better performance, implying that the design of a learnable and more representative block representation method can be a promising future research direction for trainable dynamic sparse attention. We will add more details in the next version of our manuscript."}, {"review_id": "YRBnjOdnZf", "reviewer": "Reviewer_ZZDE", "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.", "questions": "See cons.", "limitations": "na", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Pros:**\n\nThe paper is clearly written with a well-structured and logical flow, making it easy to read.\n\nThe core idea is neat and is shown to be empirically effective.\n\nWhile 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.\n\nCode is released.\n\n**Cons:**\n\nAll 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.\n\nThe 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.\n\nIt would be valuable to include results on the GSM-Infinite benchmark.", "was_revised": true, "final_justification": "Good paper, it can be accepted.", "rebuttal": "Thank you for your positive feedback and insightful questions! Below we clarify the three main concerns.\n\n**(W1) Regarding the restriction that  all attention heads within a group attend to the same selected blocks:**\n\nWe would like to clarify that MoBA **does not** introduce such restriction on block selection. MoBA allows different attention heads to independently select their own top-k KV blocks. Actually, such restriction is introduced by NSA --- quoted from Section 2.3.2 of NSA paper: \"This aggregation ensures consistent block selection across heads within the same group.\"\n\n\n**(W2) Regarding MQAR-style synthetic experiment:**\n\nWe train a MoBA model from scratch on MQAR (Multiquery Associative Recall)-style synthetic data. The synthetic data construction follows Arora, Simran, et al[1], where input sequences consist of a number of key-value pairs followed by queries (for example: A 4 B 3 C 6 F 1 E 2 -> A ? C ? F ? E ? B ? ).\n\nThe model is a 16-layer, 16-head dense Transformer with Multi-Head Attention (MHA). It has 822M non-embedding parameters, a hidden dimension of 2048, and employs the GLU activation in the feed-forward blocks. In our MoBA configuration, we segment the input into chunks of 256 tokens and select the top-4 experts per token; the overall sequence length is 2048.\n\nWe observe that the LM loss for unknown (\"?\") tokens can finally go to zero (less than 1e-7), meaning that MoBA can fit MQAR-style synthetic data. We will add more details in the next version of our manuscript.\n\n[1] Arora, Simran, et al. \"Zoology: Measuring and Improving Recall in Efficient Language Models.\" Proceedings of 12th International Conference on Learning Representations (ICLR). ICLR, 2024.\n\n**(W3) Regarding GSM-Infinite:**\n\nWe evaluated Llama-8B-1M-MoBA and Llama-8B-1M-Full on a subset of the GSM-Infinite task. Because the rebuttal window is short, our protocol deviates slightly from the original paper:\n\n- We sample 200 examples from GSM-infinite, whose ops rangs from 2 to 20 (uniformly) \n- We report exact-match accuracy.\n\n\n| Accuracy (%)                     | MoBA Prefill Full Decode | MoBA Prefill MoBA Decode | Full |\n|----------------------------------|--------------------------|--------------------------|------|\n| gsm_infinite_sample200_hard_128k | 10.5                     | 9.5                      | 12   |\n| gsm_infinite_sample200_medium_128k | 8.5                     | 9.5                      | 9    |\n| **Average**                      | **9.5**                  | **9.5**                  | **10.5** |\n\nBoth MoBA and Full perform poorly yet comparably on this subset, in terms of average accuracy for ops <= 20 in 128K context. The above results suggest that the present base model (Llama-8B) and our long-context training recipe are not yet strong enough to reveal meaningful differences between full and dynamic sparse attention on GSM-Infinite. We will present a more rigorous evaluation in the next version of the manuscript.\n\n\n\nAs a reference, the average accuracy for ops<= 30 of Llama3.1-8B-Instruct reported in GSM-infinite paper[1]  is 21.86%.\n\n[1] Zhou, Yang, et al. \"GSM-Infinite: How Do Your LLMs Behave over Infinitely Increasing Context Length and Reasoning Complexity?.\" arXiv preprint arXiv:2502.05252 (2025)."}]

Author rebuttal:

Meta-review: The paper presents MoBA, a novel approach for scaling large language models (LLMs) to handle long contexts efficiently. MoBA employs a Mixture of Block Attention mechanism, which allows the model to dynamically select the most relevant key blocks for each query token, maintaining sparsity without losing critical information. This approach is designed to adhere to the 'less structure' principle, enabling the model to learn where to focus attention during training. The paper also highlights the method's ability to transition between full and sparse attention, offering a key advantage in efficiency. Experimental results show that MoBA performs comparably to full attention on both general language tasks and long-context benchmarks, and it has already been deployed in production workloads with long-context requirements. The authors address reviewer concerns by providing additional comparisons with other dynamic sparse attention methods, clarifying the configuration of hybrid attention, and discussing the trade-off between efficiency and long-context capability.
Decision: Accept (spotlight)