MoBA: Learned Block Sparsity for Long-Context Attention

Core Snapshot

TL;DR MoBA is an attention architecture that transfers the Mixture of Experts idea from feed-forward layers to the context dimension: each query token is routed to the top-k key-value blocks that are most relevant under a simple mean-pooled affinity score. The result is a drop-in sparse attention module that keeps the original Transformer parameter count, preserves causality through forced current-block attention, and can be switched back to full attention when needed. In the paper's scaling-law and Llama-based experiments, MoBA matches full attention closely on loss and long-context benchmarks while delivering attention-layer speedups up to 6.5x at one million tokens and 16x at ten million tokens.

Background positioning This is not a purely theoretical paper, but an industrial method report from Moonshot AI with deployment evidence: MoBA is described as supporting Kimi's long-context requests. Its novelty is architectural and empirical rather than mathematical proof. It sits between static sparse attention and linear attention: it does not abandon softmax attention, nor does it impose a hard window, but it introduces a learned block-level routing mechanism whose sparsity pattern can approximate several existing static architectures.

Problem and Motivation

Long-context Transformers are constrained by the quadratic attention cost in the sequence length $N$. The paper identifies a concrete failure mode that is more precise than "attention is too expensive": when we try to fix the cost by imposing a structure, that structure may encode a wrong inductive bias. Sliding-window attention assumes relevant tokens are nearby; attention-sink methods assume relevant tokens are at the beginning plus a recent window. These assumptions are often useful, but they are task-dependent. A long-document question may need a distant paragraph, a repository may need a header far from the query, and chain-of-thought generation may need to attend to an intermediate reasoning state. A predefined sparse pattern cannot adapt to all of these cases at training time.

A second failure mode appears in dynamic sparse attention systems such as Quest, MInference, or retrieval-based inference accelerators. Those methods often select tokens or blocks during inference, which reduces prefill or decode cost, but they do not fundamentally solve the training cost of a long-context model. If the model is trained with full attention, the training step still pays the quadratic cost. The paper's question is therefore not only "how to make attention faster at inference," but "how to make a trainable long-context attention module that keeps the Transformer backbone and learns where to attend."

Linear attention models address a different problem. They replace softmax attention with state-space or RNN-like recurrences, giving asymptotic linear cost. But they are not a simple replacement for a pretrained Transformer. Converting a softmax attention model into Mamba, RetNet, RWKV, or related architectures is nontrivial, and the paper argues that evidence for complex reasoning remains limited. MoBA therefore adopts a conservative stance: keep exact softmax attention over a small set of historical blocks, and let the model discover which blocks matter.

Core Chapter: From Full Attention to Block-Level Mixture Routing

Starting Point: Exact Softmax Attention and Its Cost

The baseline is standard causal attention, written for a single head as:

Here $\mathbf{q}$ is one query vector of dimension $d$, and $\mathbf{K},\mathbf{V}$ are the key and value matrices for the $N$ previously visible tokens. This expression defines the exact softmax attention that every modern Transformer decoder relies on: the query computes similarity scores against every key, those scores are normalized, and the result is a convex combination of all value vectors. The paper's method preserves this exact operation locally, but not globally: it applies it only to a subset $I$ of keys and values. The baseline matters because it gives the compatibility criterion for the rest of the paper. MoBA is not a new approximate attention family in the same sense as linear attention; it is a selector for the softmax attention already present in Transformers. If a model trained with MoBA is later switched back to full attention, the parameter shapes and the underlying attention function remain the same. The limitation is equally clear: because this formula scores every key, its cost scales with $N^2$ per head, and at long context lengths the attention layer becomes a dominant prefill cost.

The Change: Treat KV Blocks as Experts and Route Queries with Top-k Gating

MoBA rewrites attention as:

The selected set $I$ is no longer the entire sequence. It is the union of the key-value blocks chosen by a gating rule. The block partition is:

where $B$ is block size and $n$ is the number of blocks; the paper assumes $N$ is divisible by $n$ for simplicity. The selected indices are:

where $g_i$ indicates whether block $i$ is active. This formula is the architectural pivot: the attention pattern is no longer fixed by a window or sink rule, but by the binary vector $\{g_i{\}}_{i=1}^n$. The model can, in principle, choose any subset of blocks for any query. The top-k gate chooses exactly the blocks with the highest affinity scores.

The affinity between a query and a block is defined by the inner product between the query and the mean-pooled key vector of that block:

Here $mean\_pool(\mathbf{K}[I_i])$ is the average of the keys inside block $I_i$, producing one representative vector per block. The role of this formula is to create a cheap block-level relevance score before expensive attention is computed. It is not the attention itself, because it collapses a whole block into one vector; it is a router. The boundary of this design is important: if a block contains many irrelevant keys and a few highly relevant keys, mean pooling can dilute the signal. The paper does not report an ablation replacing mean pooling with max pooling, attention pooling, or token-level scoring, so we cannot conclude that mean pooling is optimal; it is a practical choice that is computationally light and compatible with FlashAttention-style block organization.

The hard routing rule is:

where $k$ is the number of selected blocks and $\mathbb{I}$ denotes the indicator function. This equation applies the Mixture of Experts "top-k" convention to attention blocks rather than feed-forward experts. In standard MoE, a token is dispatched to a small set of FFN experts; in MoBA, a query is dispatched to a small set of KV blocks. If $k=n$, the selected set becomes all blocks, and the formula recovers full attention; if $k$ is small and the router repeatedly selects nearby blocks, it can approximate a window. The hard top-k gate is efficient but not smooth: the gradient does not flow through the discrete selection in the same way it flows through a softmax gate. The paper does not present a formal gradient analysis for this hard routing step, so the explanation of training stability is largely empirical.

Causality is preserved by a second routing rule. The paper states that no query may attend to future blocks, and every query must also attend to the block that contains itself:

For future blocks, the score is treated as masked, effectively preventing selection. This rule is more subtle than it first appears. The current block contains both the query token and tokens that come after it, so attending to the whole block would leak future information if the block spans the current position. MoBA avoids that by always selecting the current block but applying a causal mask inside the block. In MoE terms, the current block behaves like a shared expert: unlike the optional historical blocks, it is always active. This also gives the model a local attention bias without hard-coding a window, because the current block is necessarily included regardless of the top-k scores.

The sparsity induced by the block selection can be summarized by the attention-pair fraction that is removed:

This formula gives a useful reading of the paper's efficiency claims. In the 8K scaling-law setting, $B=512$ and $k=3$, so the model attends to roughly $1536$ of $8192$ positions and the paper reports sparsity of about $81.25\%$. In the 32K trailing-loss setting with the same $B$ and $k$, the reported sparsity is about $95.31\%$. The boundary condition is also visible here: if $B$ is fixed while $N$ grows, $\rho$ grows and the model may become too sparse unless $k$ or $B$ is adjusted. The paper's long scaling to ten million tokens keeps sparsity roughly constant by increasing block size, which is why the implementation section discusses varying-length FlashAttention and online softmax combination rather than assuming one fixed block geometry works at every length.

Why the Change Is Plausible: Expressivity, Compatibility, and Hybrid Repair

The strongest conceptual argument in the paper is that existing static sparse patterns are special cases of MoBA. Sliding-window attention can be viewed as a gate that always selects recent blocks. Attention-sink attention can be viewed as a gate that always selects the initial blocks and recent blocks. Because MoBA's gate can select arbitrary blocks, it is strictly more expressive than those fixed patterns in the space of block-level sparse attention. This is not a theorem in the paper, but an architectural observation: by changing the routing rule, one can recover known static patterns. This matters for the "less structure" principle because it says MoBA does not impose a new bias stronger than the ones it generalizes; it learns block selection from query-key affinity while keeping full attention as a limiting case.

Compatibility with pretrained models comes from parameter preservation. The paper explicitly frames MoBA as a substitute for full attention that does not add or remove parameters. This is why the hybrid recipe is credible: a layer can run as MoBA during long-context prefill or early training and run as full attention during a later stabilization stage. In the hybrid training experiment, three 1.5B models are trained on 30B tokens with context length 32K, using block size 2048 and top-k 3. The MoBA-only recipe shows higher trailing position-wise loss in Figure 5a, while the MoBA/full hybrid reaches a loss nearly identical to full attention. This suggests that MoBA can be used for efficient long-context training, but a short full-attention phase can recover the exactness needed for trailing-token modeling. The paper does not provide a numerical loss gap for the hybrid configuration in the supplied text, so the claim should be read as the visual and qualitative result reported in Figure 5a.

Fine-grained segmentation is another design choice with empirical support. The paper compares block granularity while holding sparsity fixed at 75 percent: for a 32K context, the model divides the sequence into 8, 16, 32, 64, or 128 blocks and selects 2, 4, 8, 16, or 32 blocks accordingly. Figure 4 reports that the coarsest setting, selecting 2 out of 8 blocks, is about one hundredth in validation LM loss worse than finer-grained variants. This is the attention-side analogue of a known MoE lesson: smaller experts can give better specialization, and coarser experts can over-penalize or under-specialize. The paper does not explain why a particular granularity is best for a particular task, so the safe conclusion is that block size and top-k should be tuned jointly, not treated as neutral implementation constants.

The layer-wise hybrid design responds to a specific failure mechanism during supervised fine-tuning. The paper observes that MoBA can be suboptimal in SFT, shown in Figure 5b, and hypothesizes that the cause is loss masking: SFT often masks prompt tokens from the loss, so gradients begin only on response tokens. A sparse attention pattern may then receive weak or uneven backpropagation through the long context. Switching the last several Transformer layers to full attention while keeping earlier layers as MoBA reduces SFT loss in Figure 5b and Figure 5c. This is a plausible gradient-flow repair, but the paper labels it as a speculation rather than a proof. The evidence is behavioral: the hybrid pattern improves loss curves; the supplied text does not report per-layer gradient norm measurements that would directly confirm the masking mechanism.

Experiments and Evidence

Scaling Laws: Is MoBA a Drop-in Replacement?

The scaling experiments are designed to isolate attention as the only architectural variable. Table 1 lists five model configurations, all trained with block size 512 and top-k 3 in the initial 8K sequence setting.

The important control is that every model, both MoBA and full attention, uses the same hyperparameters except for the attention module itself. At 8K context, the paper reports that the validation loss curves in Figure 3a are nearly indistinguishable, with differences within $10^{-3}$. That number is not just a small loss gap; it is the first quantitative anchor for the claim that MoBA does not obviously damage the model's ability to learn from the training distribution, even when it attends to only about $18.75\%$ of the token pairs in the reported sparsity calculation.

Table 2 in the paper, reproduced from the supplied text, fits power-law forms to average LM loss and trailing-token loss. The notation $C$ is the scaling compute variable.

The average 8K loss is almost identical between MoBA and full attention, but the trailing loss row is more informative about long-context behavior. MoBA's trailing loss has a slightly worse intercept, $1.546$ versus $1.464$, but a steeper exponent, $-0.108$ versus $-0.097$. The paper interprets this narrowing gap in Figure 3b as evidence of long-context scalability: the initial penalty may disappear as compute increases. This is a stronger claim than the visual reading suggests, because it relies on extrapolating fitted power laws. A reader should treat trailing-token scaling as suggestive evidence, not as a proof that MoBA will always catch up at arbitrary context lengths.

Hybrid Training and Layer-wise Recovery

The hybrid training recipe shows why MoBA alone may not be enough. In the 30B-token, 32K-context experiment, the paper compares three recipes: full attention throughout, MoBA throughout, and a two-stage recipe using MoBA for 90 percent of tokens and full attention for the remaining 10 percent. Figure 5a shows that MoBA-only produces higher position-wise losses for trailing tokens, while the hybrid recipe nearly matches full attention. The paper also notes that it did not observe significant loss spikes when switching between MoBA and full attention. That absence of spikes is an engineering signal: if switching attention patterns created a distribution shift in the model's attention behavior, training would likely destabilize. The result supports the paper's design choice of preserving parameter shapes and the softmax operation itself.

For SFT, the paper proposes switching the last several layers from MoBA to full attention. This is motivated by the observation in Figure 5b that MoBA can be worse during SFT, and the authors speculate that prompt loss masking creates sparse gradients. Figure 5c shows that layer-wise hybridization reduces SFT loss. The supplied text does not give the exact numerical improvement from the layer-wise experiment, so the evidence should be interpreted as a qualitative curve comparison. Still, the design is useful because it identifies a concrete stage, SFT, where pure MoBA may need rescue by exact attention in the layers where final predictions are made.

Real-World LLM and Benchmark Evidence

The strongest external validation comes from the Llama-8B experiment. The paper starts from Llama 3.1 8B Base and continually pretrains from 128K to 256K, 512K, and 1M tokens, using position interpolation at the 256K stage. After reaching 1M context, MoBA is activated for 100B tokens with block size 4096 and top-k 12. The paper reports a sparsity of about $95.31\%$ in this setting and uses the layer-wise hybrid structure: the last three Transformer layers remain full attention, while the other layers use MoBA. During downstream evaluation, MoBA is used for prefill only, and full attention is used during generation. That caveat is crucial: the benchmark numbers test long-context understanding with sparse prefill, not necessarily sparse decode.

Table 2 of the paper compares Llama-8B-1M-MoBA against Llama-8B-1M-Full on many benchmarks.

These numbers show near parity rather than universal improvement. On RULER at 128K, MoBA scores 0.7818 against full attention's 0.7849, a difference of only 0.0031. On CEval, GSM8K, Loogle, and MBPP Sanitized, MoBA is ahead by 0.0108, 0.0136, 0.0193, and 0.0311 respectively. On MMLU, BBH, HellaSWAG, HumanEval, and RULER, the differences are small but favor full attention. The pattern is consistent with the paper's claim: the sparse architecture does not appear to break general capabilities in the evaluated regime, but it is not a magic performance upgrade. The most meaningful result is that the drop is small enough to justify using MoBA when efficiency is the constraint.

Figure 7 evaluates needle-in-the-haystack up to one million tokens. The paper reports satisfactory performance, but the supplied text does not give the full heatmap score distribution, so we can only conclude that retrieval failure at extreme length is not reported as a catastrophic problem in their test. The benchmark table and the needle experiment together suggest that MoBA can extend effective context in a pretrained Transformer, but the evaluation still has an important boundary: MoBA is used only during prefill, while generation falls back to full attention. That choice improves task quality but means the method is not presented as a complete KV-cache pruning solution for long generation.

Efficiency and Implementation

The efficiency section separates the attention layer from the rest of the network. This is a clean experimental design because FFN parameters and operations are identical between MoBA and full attention in their comparison, so the measured difference isolates the attention module. In Figure 2a, MoBA is more efficient than full attention across context lengths and reaches a speedup ratio of up to 6.5x when prefilling one million tokens. In the ten-million-token scaling test, the paper reports a 16x reduction in attention computation time relative to standard Flash Attention. To maintain constant sparsity at longer lengths, it keeps top-k and the number of MoBA blocks fixed while increasing block size proportionally, and expands tensor parallelism to the query-head level to fit keys and values across distributed devices.

The implementation story is as important as the algorithmic one. MoBA cannot be read as a simple masked-attention trick. It arranges queries according to their selected KV blocks, computes block attention with FlashAttention for varying lengths, and combines outputs with online softmax because a query may have both a current-block causal contribution and multiple historical-block contributions. Algorithm 1 formalizes this: split KV into blocks, compute mean-pooled gating scores, apply causal mask and top-k, separate current-block attention from MoBA attention, run two FlashAttention variable-length calls, then combine with online softmax. The efficiency is real but conditional on a substantial engineering stack. Without these kernels, the block routing overhead and memory rearrangement could erase much of the theoretical FLOP reduction.

Evidence Quality and Limitations

The evidence is strongest in the scaling-law and Llama-8B sections. Multiple model sizes are trained, controls keep hyperparameters fixed, and downstream tasks include long-context retrieval, reasoning, code, knowledge, and Chinese evaluation suites. The reported speedups are concrete and attributed to attention-layer measurement rather than end-to-end wall-clock claims. The fine-grained segmentation ablation is also useful because it changes a design axis while holding sparsity fixed.

The evidence is weaker in three places. First, the paper's SFT layer-wise hybrid is motivated by a gradient speculation and evaluated by loss curves, not by direct gradient diagnostics. Second, the long-context efficiency claims include prefill-focused evaluation, while generation falls back to full attention; therefore MoBA's value for memory-bound long decoding is not the central claim. Third, the top-k router uses a mean-pooled block representation, but there is no reported comparison with more expressive scorers such as max pooling, learned block summaries, or hierarchical block scoring. The paper leaves open whether the mean-pool shortcut is the bottleneck at the hardest long-document reasoning cases.

MoBA also exposes a design tension between flexibility and training stability. Hard top-k selection gives interpretability and compatibility, but it may be brittle when multiple blocks have nearly equal relevance. If a task requires attending to several small, distant passages, the fixed top-k rule may select too many redundant blocks or too few informative ones depending on block size. The current implementation mitigates this with hybridization and fine-grained blocks, but those are empirical patches rather than a closed-form solution.

Deep Insight and Conclusion

MoBA's main contribution is not merely "faster attention." It is the reframing of attention sparsity as an expert routing problem along the context dimension. Once KV blocks are experts, existing sparse attention patterns become routing rules, and the model can learn the routing from query-key affinity. That gives a rare combination: parameter-compatible substitution, causal preservation, switchable full attention, and a practical subquadratic prefill path. The paper shows that this is enough to train and deploy a long-context model with benchmark quality close to full attention.

The lasting practical lesson is that efficient attention should be designed as a continuum between exact and sparse, not as a binary architectural fork. Full attention is the upper bound; sliding window and sink attention are fixed corners of the MoBA routing space. The useful middle ground is hybrid MoBA: long-context prefill can use block routing, trailing-token SFT can switch the last layers to full attention, and generation can fall back to exact attention if quality demands it. For future work, the natural next step is to learn or regularize the block affinity more robustly, for example by replacing hard mean-pool routing with a differentiable block selector or by studying whether MoBA can be extended to full generation while controlling KV memory. That would test whether the method's biggest promise, efficient reasoning at long context, can survive once the safety net of full-attention decoding is removed.