MoBA: Mixture of Block Attention for Long-Context LLMs

ArXiv: 2502.13189

🎯 Pitch

MoBA shows you don’t need to sacrifice full attention’s versatility: by letting each query dynamically select which context blocks to attend to, it matches standard attention at 1M-token lengths while cutting compute by over 80%. This is the first sparse scheme flexible enough to serve as a drop-in replacement in production, already deployed in Kimi.

---

1. Executive Summary

This paper introduces Mixture of Block Attention (MoBA), a novel attention architecture that applies Mixture of Experts (MoE) principles to the self-attention mechanism of Transformers, dynamically routing each query token to only the most relevant historical KV blocks (partitioning the full context into blocks and using a learned top-k gating mechanism per query). Evaluated on Llama 3.1 8B extended to 1M-token contexts via continual pretraining, MoBA matches full attention across benchmarks—achieving 0.7818 vs. 0.7849 on RULER at 128K and near-perfect Needle-in-a-Haystack retrieval at 1M tokens—while delivering up to a 6.5× speedup at 1M sequence length (reducing attention computation time from approximately 800 ms to 120 ms at 1M tokens) and scaling efficiently to 10M tokens with a 16× speedup at constant 95.31% sparsity. MoBA also supports seamless transitions between sparse and full attention through a hybrid training recipe, establishing that block-sparse attention can substitute for full attention without degrading long-context performance when a modest fraction of final layers retain full attention.

2. Context and Motivation

The Core Problem: Attention's Quadratic Complexity Prevents True Long-Context Scaling

The fundamental problem this paper addresses is stark and well-known: the standard self-attention mechanism in Transformers scales quadratically with sequence length. For a sequence of $N$ tokens, computing attention requires $O(N^2)$ operations, meaning that doubling the context length quadruples the computational cost. This is not merely an asymptotic concern — it creates a hard practical ceiling on how long a context window an LLM can process within reasonable time and memory budgets.

The authors frame this problem in the context of the broader push toward artificial general intelligence (AGI), arguing that "a pivotal capability for achieving AGI is the ability to process, understand, and generate long sequences" (Section 1). They point to concrete evidence that the field is moving in this direction: the popularity of long-input models like Kimi (MoonshotAI, 2023), Claude (Anthropic, 2023), and Gemini (Reid et al., 2024), and the emergence of long chain-of-thought (CoT) reasoning in models like Kimi k1.5 (Team et al., 2025), DeepSeek-R1 (Guo et al., 2025), and OpenAI o1/o3 (Guan et al., 2024).

The theoretical significance is clear: if quadratic complexity is the bottleneck, then any architecture that can reduce this complexity while preserving the expressive power of attention represents a fundamental advance. The practical significance is equally compelling — real-world applications increasingly demand million-token contexts (analyzing entire codebases, processing full legal documents, reasoning over multi-hour meeting transcripts), and these applications are simply infeasible if every token must attend to every other token.

Importantly, the paper does not frame this as a problem that only affects inference. The authors explicitly note that existing approaches "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" (Section 1). This dual concern — training efficiency and inference efficiency — motivates a solution that works during both phases.

The Sparsity Hypothesis: Most Attention Connections Are Unnecessary

Underlying the paper's approach is an observation about the nature of attention itself: attention scores are inherently sparse. The authors cite two sources of evidence for this sparsity. First, it arises mathematically from the softmax operation, which concentrates probability mass on a small number of high-similarity token pairs — a property that has been studied in recent work on sparse attention patterns (Jiang et al., 2024). Second, they appeal to a biological analogy: "sparse connectivity is observed in brain regions related to memory storage" (Watson et al., 2025, cited in Section 1), suggesting that selective attention is a natural solution to the problem of processing large amounts of contextual information.

This sparsity hypothesis is crucial because it defines the opportunity: if each query token only needs to attend to a small fraction of the available context, then attention can be made substantially more efficient without losing essential information. The challenge is not whether to be sparse, but how to determine which connections to keep and which to prune — and to do so in a way that doesn't impose task-specific biases or degrade performance on the full range of tasks that LLMs must handle.

Where Existing Approaches Fall Short

The paper identifies three broad families of prior work on efficient attention, each with specific limitations that MoBA aims to address. Understanding these limitations is essential to seeing why MoBA's design choices matter.

Approach 1: Static Sparse Patterns with Predefined Structure

The first family of approaches imposes fixed, hand-designed sparsity patterns on the attention matrix. The paper cites a long lineage: Sparse Transformer (Child et al., 2019), Star-Transformer (Guo et al., 2019), BlockBERT (Qiu et al., 2019), Longformer (Beltagy et al., 2020), BigBird (Zaheer et al., 2020), and LongNet (Ding et al., 2023), among others. These methods use combinations of sliding windows, global tokens, random attention, dilated attention, and block-sparse patterns — all determined in advance rather than learned from data.

Two specific static patterns receive extended discussion because they are widely used:

• Sliding window attention (Beltagy et al., 2020): Each query token attends only to its $W$ nearest neighbors. This is computationally efficient ($O(NW)$ rather than $O(N^2)$) but imposes a rigid locality bias — information can only propagate within the window, which means long-range dependencies require many layers to bridge.
• Attention sink (Xiao et al., 2023): Each query token attends to a combination of the initial few tokens (the "sink") and the most recent tokens. This also provides a fixed structure that, while effective for certain streaming scenarios, is again a predetermined bias.

The paper's central criticism of static approaches is that they "tend to be highly task-specific, potentially hindering the model's overall generalizability" (Section 1). A sliding window might work well for local syntactic processing but fail when a distant pronoun needs to resolve to its antecedent. An attention sink might help with streaming generation but miss relevant content buried in the middle of a long document. The fixed structure cannot adapt to the varying demands of different tasks or different positions within a single sequence.

Why this matters: The paper is positioning itself against what it calls the "strongly biased structures" of prior work (Abstract). The phrase "less structure" — used in both the Abstract and Section 1 — signals a design philosophy: rather than imposing a predetermined attention pattern, the architecture should let the model learn where to attend. This philosophy echoes the broader trend in deep learning away from hand-designed inductive biases toward learned, flexible mechanisms.

Approach 2: Dynamic Sparse Attention at Inference Time Only

The second family includes methods that dynamically select which tokens to attend to, but only during inference — they do not fundamentally change the training process. The paper cites:

• Quest (Tang et al., 2024): Selects subsets of tokens based on query-aware sparsity at inference time. The authors explicitly note that "Quest, in particular, can be viewed as MoBA with a smaller block size and a specialized block representation function which combines both min and max pooling" (Section 4). This is a revealing comparison — it suggests that MoBA is not entirely novel in its inference-time behavior (Quest does something similar), but rather in how that behavior is integrated into training.
• Minference (Jiang et al., 2024): Another inference-time dynamic sparsity method that identifies important tokens through attention profiling on the fly.
• RetrievalAttention (Liu et al., 2024): Uses vector retrieval to select relevant tokens during inference.

The paper's criticism of these approaches is specific and pointed: "Although such methods can reduce computation for long sequences, they do not substantially alleviate the intensive training costs of long-context models" (Section 1). This is the key gap. If training still requires full attention (or some other mechanism that doesn't benefit from the inference-time sparsity), then the overall cost of developing long-context models remains prohibitive. You can make inference faster, but you still need to pay the quadratic training cost to produce the model in the first place.

Why this matters: MoBA is designed to be used during both training and inference, providing computational savings across the entire model lifecycle. This is a practical distinction with significant economic implications — if training a 1M-context model costs $4\times$ less with MoBA than with full attention, that's a substantial reduction in the barrier to entry for long-context research and deployment.

Approach 3: Linear Attention and Alternative Architectures

The third family abandons the standard softmax-based attention entirely in favor of linear approximations or entirely different sequence modeling primitives. The paper cites:

• Linear attention models: Performer (Choromanski et al., 2020), Linformer (Wang et al., 2020) — these approximate the softmax attention with kernel methods or low-rank factorizations.
• State Space Models: Mamba (Gu and Dao, 2023), which replaces attention with a selective state space mechanism.
• Hybrid RNN/Transformer models: RWKV (Peng et al., 2023, 2024), RetNet (Sun et al., 2023) — these blend recurrent and attention-based computation.

The paper identifies two specific problems with these approaches, both related to practicality:

1. High conversion costs from existing Transformers. The authors cite work on "linearizing" large language models (Mercat et al., 2024; Wang et al., 2024; Bick et al., 2025; Zhang et al., 2024), noting that adapting pre-trained Transformer models to linear attention "typically incurs high conversion costs" or "requires training entirely new models from scratch" (Li et al., 2025). Given the massive investment in pre-trained Transformer-based LLMs (Llama, GPT, Gemini, etc.), any approach that cannot leverage these existing models faces a severe practical disadvantage.

2. Limited evidence in complex reasoning. The paper states bluntly that "evidence of their effectiveness in complex reasoning tasks remains limited" (Section 1). This is a significant concern because reasoning — particularly the long chain-of-thought reasoning that motivates extended context windows — is one of the primary use cases for long-context models. If linear attention degrades reasoning performance, the efficiency gains are moot.

Why this matters: The paper is carving out a specific niche in the design space: retain the full Transformer architecture (including softmax attention), make attention sparse and learned rather than fixed, and provide a mechanism that works during both training and inference. This positions MoBA as a conservative innovation — it doesn't require rethinking the fundamental architecture, it doesn't abandon the properties of softmax attention that are well-understood, and it can be applied to existing pre-trained models via continual pretraining rather than requiring training from scratch.

The Critical Research Question

The paper crystallizes its motivation into a single research question (Section 1):

"How can we design a robust and adaptable attention architecture that retains the original Transformer framework while adhering to a 'less structure' principle, allowing the model to determine where to attend without relying on predefined biases?"

Each clause of this question is deliberate:

• "Retains the original Transformer framework" — rules out linear attention and alternative architectures, positioning MoBA within the familiar attention paradigm.
• "Adhering to a 'less structure' principle" — rules out static sparse patterns like sliding windows and attention sinks, positioning MoBA as a learned, dynamic mechanism.
• "Allowing the model to determine where to attend" — the core value proposition: the model learns which blocks are relevant rather than having relevance imposed by architectural constraints.
• "Without relying on predefined biases" — the contrast with prior work; not "we have a better bias" but "we eliminate the need for biases entirely."
• "Ideally, such an architecture would transition seamlessly between full and sparse attention modes" — a pragmatic requirement that enables compatibility with pre-trained models and flexible deployment.

How MoBA Positions Itself: Applying MoE Principles to Attention

The paper's central conceptual move is to borrow the Mixture of Experts (MoE) paradigm — widely successful in scaling feedforward networks (Shazeer et al., 2017; Lepikhin et al., 2020; Fedus et al., 2022; Zoph et al., 2022) — and apply it to the attention mechanism. This is explicitly stated as the innovation:

"MoBA pioneers its application to long context attention, allowing dynamic selection of historically relevant blocks of key and values for each query token." (Section 1)

This is a significant conceptual reframing. In standard MoE, each token is routed to a subset of FFN "experts" — the model learns which experts handle which types of input. In MoBA, each query token is routed to a subset of context "blocks" — the model learns which parts of the history are relevant to the current query. The analogy is clean: FFN experts → KV blocks, token-to-expert routing → query-to-block routing, load balancing concerns → block utilization concerns.

What makes this more than a superficial analogy is that MoBA inherits specific techniques from the MoE literature: the top-k gating mechanism with learned affinity scores, the concept of "shared experts" (here, the current block is always attended to, analogous to shared experts in architectures like DeepSeekMoE (Dai et al., 2024) and Qwen2.5 (Yang et al., 2024)), and the exploration of fine-grained segmentation (partitioning into more, smaller blocks rather than fewer, larger blocks).

The paper also positions MoBA as a generalization of existing sparse attention patterns. It explicitly shows that sliding window attention and attention sink can both "be viewed as special cases of MoBA" (Section 2.2) — sliding window corresponds to a gating network that always selects the most recent blocks, and attention sink corresponds to a gating network that always selects both the initial and recent blocks. This is not just a taxonomic point; it demonstrates that MoBA has "stronger expressive power" than these fixed patterns because it can learn to approximate them when they are optimal, but can also learn entirely different patterns when the task demands it.

The Practical Deployment Context: Kimi and Real-World Constraints

A crucial piece of motivation that distinguishes this paper from purely academic work is its deployment context. The paper states: "MoBA has already been deployed to support Kimi's long-context requests" (Abstract). This is not a theoretical proposal or a small-scale experiment — it is an architecture that has been tested in production at Moonshot AI, serving real users with real long-context queries.

This context explains several of the paper's design priorities:

• Seamless transition between full and sparse attention is not just an academic curiosity — it enables gradual rollout in production, where different layers or different stages of processing can use different attention modes.
• Compatibility with pre-trained models is essential because Kimi is built on existing Transformer architectures; MoBA must work as a "continual pre-training solution" (Section 5) rather than requiring a from-scratch retraining.
• Efficiency at real deployment scales — 1M tokens with 6.5× speedup, 10M tokens with 16× speedup — is measured in concrete, production-relevant terms, not just FLOPs counts.
• The layer-wise hybrid strategy (keeping the last few layers as full attention) emerged from practical supervised fine-tuning challenges, as the authors note that "MoBA sometimes results in suboptimal performance during SFT" (Section 3.2) due to sparse gradient issues from loss masking on prompt tokens.

This deployment reality also explains why the paper emphasizes that MoBA "does not introduce new parameters or remove existing ones" (Section 2.2 and 3.1) — it is a drop-in architectural change that preserves parameter counts, simplifying integration into existing training pipelines and model checkpoints.

Summary of the Motivation Gap

To synthesize: the paper identifies a clear gap in the landscape of efficient attention research. On one side, static sparse patterns (sliding windows, attention sinks) are efficient but impose task-specific biases that limit generality. On another side, inference-time dynamic sparsity (Quest, Minference) improves efficiency but doesn't reduce training costs. On a third side, linear attention models (Mamba, RWKV) reduce complexity but lack evidence for complex reasoning and require abandoning pre-trained Transformer investments.

MoBA occupies the intersection these approaches leave empty: learned, dynamic sparsity that works during both training and inference, retains the full Transformer architecture and softmax attention, and can be applied to existing models through continual pretraining. The MoE analogy provides both the conceptual framework and specific architectural techniques (top-k gating, fine-grained segmentation, shared experts), while the deployment on Kimi demonstrates production viability at million-token scales.

3. Technical Approach

3.1 Reader Orientation

MoBA is an attention mechanism — not a complete model, but a drop-in replacement for the standard self-attention layer inside a Transformer — that lets each query token attend to only a small, dynamically chosen subset of the full context rather than every token, dramatically reducing computation while preserving the model's ability to learn long-range dependencies. The problem it solves is the quadratic cost of standard attention ($O(N^2)$ operations for $N$ tokens), which makes processing million-token sequences prohibitively expensive; the shape of the solution is to borrow the Mixture of Experts routing paradigm — where each input is selectively routed to a subset of specialized sub-networks — and apply it to the attention mechanism itself, treating blocks of the context as "experts" that query tokens can choose to consult or ignore based on learned affinity scores.

3.2 Big-Picture Architecture (Diagram in Words)

MoBA has four major components that operate inside every attention layer:

1. Context Partitioning: The full sequence of $N$ key-value tokens is divided into $n$ contiguous blocks of equal size $B = N/n$. Each block functions analogously to an "expert" in MoE terminology — it contains a chunk of the available context that some queries may need and others may skip.

2. Gating Network (Router): For each query token, a lightweight scoring function computes an affinity score between the query and each KV block (by comparing the query vector against a mean-pooled representation of the block's keys). A top-k operator then selects the $k$ highest-scoring blocks, plus the current block (the block containing the query itself) is always included as a "shared expert." The output is a sparse binary assignment matrix $G$ indicating which blocks each query attends to.

3. Block-Sparse Attention Computation: Rather than computing one large attention operation over all $N$ keys, MoBA computes separate attention operations for each KV block — but only for the subset of queries assigned to that block. This exploits FlashAttention's variable-length kernels for efficiency.

4. Output Combination with Online Softmax: Because a query may attend to multiple blocks (its current block plus up to $k-1$ historical blocks), the partial attention outputs from each block must be combined. This is done using online softmax rescaling (tiling) — a numerically stable way to merge independently computed softmax-attention outputs into the correct result without recomputing from scratch.

Information flows as follows: a query vector $q$ enters the gating network → the gating network computes affinity scores against all $n$ block representations → causal masking sets scores to $-\infty$ for future blocks → top-k selects the $k$ highest-scoring blocks (plus current block forced to 1) → queries are reordered and grouped by their assigned blocks → block-wise FlashAttention is computed for each block → outputs are rearranged back to original query order → block outputs are merged via online softmax to produce the final attention output.

3.3 Roadmap for the Deep Dive

• First, the standard attention formulation and its complexity, to establish the precise mathematical baseline that MoBA modifies — understanding what changes requires understanding what is being changed.
• Second, the core MoBA attention equation, which defines the block-sparse selection set $\mathcal{I}$ and shows how it alters the standard attention computation — this is the mathematical heart of the method.
• Third, the block partitioning and affinity scoring mechanism, which specifies how the context is divided into blocks and how relevance scores are computed — this determines where the sparsity comes from.
• Fourth, the top-k gating mechanism, which converts continuous affinity scores into a discrete, sparse selection — this is the routing component borrowed from MoE.
• Fifth, the causality enforcement mechanisms, which ensure autoregressive correctness — this addresses the specific challenge of preventing future-token information leakage in a block-based architecture.
• Sixth, the fine-grained segmentation and hybrid attention strategies, which are design choices that significantly impact performance — these are not core to the algorithm but are critical to making it work well in practice.
• Seventh, the implementation details and computational complexity analysis, which explain how the theoretical algorithm maps to efficient GPU kernels and what speedups are achieved — this connects the mathematical design to the empirical speedups reported in Section 3.4.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architectural design paper whose core idea is that applying Mixture of Experts-style dynamic routing to the attention mechanism — partitioning the context into blocks and letting each query learn which blocks to attend to — can reduce the quadratic complexity of attention to sub-quadratic while preserving the full Transformer framework and enabling seamless transitions between sparse and full attention.

---

Standard Attention Baseline

To understand what MoBA changes, we first establish precisely what standard (vanilla) self-attention computes. For a single query token $q \in \mathbb{R}^{1 \times d}$ (a row vector of dimension $d$, where $d$ is the per-head dimension), attending to $N$ key and value tokens represented as matrices $K, V \in \mathbb{R}^{N \times d}$, the standard attention operation is:

$\text{Attn}(q, K, V) = \text{Softmax}\left(qK^{\top}\right) V$

where $q$ is the query vector for a single token, $K$ is the matrix of all key vectors (one per context token), $V$ is the matrix of all value vectors, $qK^{\top} \in \mathbb{R}^{1 \times N}$ computes a raw similarity score between the query and every key, $\text{Softmax}(\cdot)$ normalizes these scores to a probability distribution summing to 1, and the resulting weighted combination of value vectors produces the attention output as a $1 \times d$ vector.

What it computes: For a single query token at a particular position, this operation produces a context-aware representation by taking a weighted average of all value vectors in the sequence, where the weights are determined by the softmax-normalized similarity between the query and each key. Intuitively, the query token "looks at" every past token and decides — based on learned similarity — how much to incorporate from each one.

Why this form: The dot-product attention with softmax normalization was introduced by Vaswani et al. (2017) and has become the dominant paradigm because it allows every token to directly attend to every other token, enabling arbitrarily long-range dependencies in a single layer. The softmax forces the attention weights to form a valid probability distribution (non-negative, summing to 1), which provides a natural interpretation as "how much attention is paid to each token." The $O(N^2)$ complexity arises because $qK^{\top}$ requires computing $N$ dot products for each of $N$ queries, and there is no way around this in the general case — every query must compare itself to every key to determine which ones are relevant.

The paper notes that this formulation is presented for a single attention head for clarity; the extension to multi-head attention involves performing this operation independently for $h$ different heads (each with its own learned projections) and concatenating the results.

---

The Core MoBA Attention Equation: Block-Sparse Selection

MoBA modifies standard attention by restricting each query to attend to only a subset of the available keys and values, rather than the full context. The modification is clean and minimal:

$\text{MoBA}(q, K, V) = \text{Softmax}\left(qK[\mathcal{I}]^{\top}\right) V[\mathcal{I}]$

where $q$ is the query vector (same as standard attention), $K$ and $V$ are the full key and value matrices (same as standard attention), $\mathcal{I} \subseteq [N]$ is a subset of indices $\{1, 2, \ldots, N\}$ specifying which key-value positions the query is allowed to attend to, $K[\mathcal{I}]$ denotes selecting only the rows of $K$ indexed by $\mathcal{I}$ (producing a smaller matrix of size $|\mathcal{I}| \times d$), and $V[\mathcal{I}]$ is the analogous selection from the value matrix.

What it computes: Instead of computing similarity between the query and all $N$ keys, this operation computes similarity only with the keys in the selected subset $\mathcal{I}$, performs softmax normalization over just those $|\mathcal{I}|$ scores, and produces a weighted average of only the corresponding value vectors. If $|\mathcal{I}| \ll N$, this is dramatically cheaper than standard attention — the cost per query drops from $O(N)$ to $O(|\mathcal{I}|)$.

Why this form: The key insight is that this is exactly standard attention, just applied to a subset of the context. The softmax still produces a valid probability distribution (now over $|\mathcal{I}|$ items), the dot-product similarity computation is unchanged, and the weighted averaging of value vectors is unchanged. This means that all the properties of standard attention — differentiability, compatibility with FlashAttention kernels, integration into the Transformer architecture — are preserved. The only thing that changes is which keys and values are available. The entire challenge of MoBA — and what the remainder of the architecture specifies — is how to determine $\mathcal{I}$ for each query in a way that is both computationally efficient (the selection itself must be cheap) and functionally effective (the selected subset should contain the tokens that matter for the current query).

The paper explicitly states that this formulation "does not introduce new parameters or remove existing ones" (Section 2.2) relative to standard attention. The parameter count is identical — $W_Q, W_K, W_V, W_O$ projection matrices are the same as in the standard Transformer. The selection mechanism $\mathcal{I}$ is determined by a lightweight gating computation that adds negligible parameters (only the operations to compute block representations and affinity scores, which are parameter-free beyond the existing key projections).

---

Context Partitioning: Dividing History into Blocks

The first structural decision in MoBA is how to organize the context into units that can be selectively attended to. Rather than operating at the individual token level (which would require per-token routing decisions for every query-key pair), MoBA operates at the block level — the context is divided into contiguous chunks that are treated as atomic units for routing purposes.

Formally, the full context of length $N$ is partitioned into $n$ blocks of equal size $B$:

$B = \frac{N}{n}$

where $N$ is the total sequence length, $n$ is the number of blocks (a hyperparameter), and $B$ is the block size. The paper assumes — for simplicity of exposition — that $N$ is divisible by $n$, making all blocks exactly size $B$. The $i$-th block covers token positions:

$\mathcal{I}_i = [(i-1) \times B + 1, \ i \times B]$

where $\mathcal{I}_i$ is the set of position indices belonging to block $i$, with the notation $[a, b]$ meaning the contiguous range from position $a$ to position $b$ (inclusive on both ends, following the 1-indexed convention used in the paper).

What this partitioning accomplishes: Instead of making $N$ individual routing decisions per query (one per token), the model makes $n$ decisions per query (one per block), where $n \ll N$. For a 1M-token context with block size 4096, $n = 244$ blocks rather than 1,000,000 tokens — a reduction in routing complexity of over 4000×. This is essential because the routing mechanism itself must be cheap; if computing which tokens to attend to costs as much as attending to all tokens, the sparsity provides no benefit.

Why block-level rather than token-level routing: The paper motivates this through the MoE analogy. In standard MoE, routing happens at the token level but to a small number of experts (typically 8-64). For attention, if we treated each token position as an "expert," we would need to route to $N$ experts — which defeats the purpose. By grouping tokens into blocks, MoBA achieves the same conceptual structure as MoE (a manageable number of routing targets) while still operating on the full context. An additional benefit is that contiguous blocks are cache-friendly — keys and values within a block are stored contiguously in memory, enabling efficient loading into GPU kernels.

The paper also demonstrates (in Section 2.2, in the comparison to sliding window and attention sink) that block-level partitioning is expressive enough to subsume prior static sparse patterns. Sliding window attention is equivalent to MoBA with a gating network that always selects the most recent blocks; attention sink is equivalent to always selecting the initial and recent blocks. The block structure doesn't limit what patterns can be represented — it only limits the granularity at which patterns can adapt within a block.

A critical design choice discussed in Section 3.1 is block granularity — how many blocks to partition the context into. This is explored through the fine-grained segmentation experiments.

---

Affinity Scoring: Measuring Query-Block Relevance

For each query token, MoBA must decide which blocks are worth attending to. This requires a scoring function that estimates the relevance of an entire block of keys to a single query. The paper adopts a simple and computationally cheap approach: mean-pool the key vectors within a block, then compute the dot product with the query.

Formally, the affinity score $s_i$ between query $q$ and block $i$ is:

$s_i = \langle q, \ \text{mean\_pool}(K[\mathcal{I}_i]) \rangle$

where $q \in \mathbb{R}^{1 \times d}$ is the query vector, $K[\mathcal{I}_i] \in \mathbb{R}^{B \times d}$ is the matrix of key vectors belonging to block $i$, $\text{mean\_pool}(\cdot)$ averages these key vectors along the sequence dimension to produce a single $d$-dimensional vector representing the block, and $\langle \cdot, \cdot \rangle$ denotes the inner product (dot product), producing a scalar score.

What it computes: For each block $i$, this operation compresses all $B$ key vectors in that block into a single representative vector by taking their element-wise mean, then computes the dot-product similarity between the query and this block representation. A high positive score means the query is semantically similar to the average content of the block; a low or negative score means the block is likely irrelevant.

Why this form: The key design constraint is that the scoring function must be much cheaper than actually attending to the block. Computing the full attention between $q$ and $K[\mathcal{I}_i]$ would cost $O(B \cdot d)$ operations — which is exactly what we're trying to avoid for most blocks. By reducing the block to a single vector via mean pooling ($O(B \cdot d)$ to compute the mean, but this can be pre-computed once per block and reused across all queries) and then doing a single dot product ($O(d)$), the scoring cost per block is $O(d)$ rather than $O(B \cdot d)$.

The use of mean pooling specifically — rather than max pooling, or using the first/last token, or learning a separate block representation — is a deliberate simplicity choice. Mean pooling is parameter-free, differentiable, and produces a representation that captures the "average content" of the block. The paper notes in Section 4 that Quest (Tang et al., 2024) uses a more sophisticated block representation "which combines both min and max pooling," and positions this as a specialization that MoBA could potentially adopt. The use of dot-product similarity (rather than a learned scoring network) means the affinity scores reuse the same query and key projections already computed for attention, adding no new parameters.

An important implementation detail: the mean-pooled block representations $\text{mean\_pool}(K[\mathcal{I}_i])$ can be pre-computed once per attention layer (since the key vectors don't change for a given input sequence) and then reused for all query tokens. This makes the amortized cost of computing block representations negligible compared to the attention computation itself.

---

Top-k Gating: From Continuous Scores to Discrete Selection

The affinity scores $s_i$ provide continuous measures of block relevance, but the attention mechanism needs a discrete decision: include or exclude. MoBA adopts the top-k gating mechanism from Mixture of Experts — select the $k$ blocks with the highest affinity scores. Formally, the gate value $g_i$ for block $i$ is:

$g_i = \begin{cases} 1 & \text{if } s_i \in \text{Topk}(\{s_j \mid j \in [n]\}, k) \\ 0 & \text{otherwise} \end{cases}$

where $s_i$ is the affinity score for block $i$, $\{s_j \mid j \in [n]\}$ is the set of affinity scores for all $n$ blocks (numbered $1$ through $n$), $\text{Topk}(\cdot, k)$ returns the set containing the $k$ largest scores from this set, and $g_i$ is a binary gate: 1 means the block is included in the attention computation, 0 means it is excluded.

The final selection set $\mathcal{I}$ for the query is the union of all blocks whose gate is 1:

$\mathcal{I} = \bigcup_{g_i > 0} \mathcal{I}_i$

What it computes: For a given query, this operation ranks all $n$ blocks by their affinity scores, picks the $k$ blocks with the highest scores, discards the remaining $n - k$ blocks, and constructs the attention index set $\mathcal{I}$ as the set of all token positions belonging to those $k$ selected blocks. This means each query attends to exactly $k \times B$ key-value pairs (ignoring causality adjustments), which is a fixed budget regardless of the total context length $N$.

Why this form — and why top-k specifically: The top-k mechanism has several properties that make it well-suited to this context:

1. Fixed computational budget: Regardless of how long the context is, each query attends to exactly $k$ blocks, making the per-query cost predictable and bounded. This is crucial for deployment — you can guarantee a maximum latency per attention operation regardless of input length.

2. Competition among blocks: Because only the top $k$ blocks are selected, blocks compete for the limited attention budget. The model must learn to produce meaningful affinity scores that push relevant blocks above irrelevant ones. This is the same principle that makes top-k gating effective in MoE — it forces specialization.

3. Differentiability through the scores: While the top-k selection itself is non-differentiable (it's a discrete thresholding operation), the affinity scores $s_i$ are fully differentiable. During training, gradients flow through the selected blocks' attention computations and back to the key projections (which produce the vectors that get mean-pooled) and to the query projections. This means the model learns to produce affinity scores that route queries to the correct blocks, even though the routing decision itself is a hard threshold.

4. Simplicity and interpretability: The top-k mechanism is straightforward to implement, has clear semantics (each query attends to its $k$ most relevant blocks), and produces interpretable attention patterns — you can inspect which blocks were selected for which queries and understand what the model considers relevant.

The paper does not explore alternatives like threshold-based gating (select all blocks with $s_i > \tau$) or learned thresholding — the top-k mechanism is adopted directly from the MoE literature without modification, and the experiments use fixed $k$ rather than adaptive $k$.

The "shared expert" analogy — current block attention: A critical modification to pure top-k gating is that the block containing the query token itself is always attended to, regardless of its affinity score. The paper states:

"we enforce that each token must be routed to its respective current block" (Section 2.2)

Formally, if the query token is at position $\text{pos}(q)$ and this position falls within block $i$ (so $\text{pos}(q) \in \mathcal{I}_i$), then $g_i$ is set to 1 unconditionally. From the MoE perspective, the authors draw an explicit analogy:

"the current block attention in MoBA is akin to the role of shared experts in modern MoE architectures (Dai et al., 2024; Yang et al., 2024), where static routing rules are added when expert selection."

Why force current block attention: This design choice serves multiple purposes. First, it ensures that every query attends to its immediate local context — the tokens that are closest in sequence position and therefore most likely to be syntactically and semantically relevant. Second, the current block attention is where causal masking is applied (as described in the next section), so it serves as the primary mechanism for attending to preceding tokens within the local neighborhood. Third, it prevents degenerate behavior where the top-k gating routes all queries away from their immediate context, which could cause the model to lose local coherence while searching for distant relevant blocks. The analogy to shared experts is apt: just as shared experts in MoE ensure that every token gets a baseline of general-purpose FFN computation regardless of routing, current block attention ensures every token gets a baseline of local context regardless of where the top-k gating routes it.

Interaction with top-k: Because the current block is always included, the effective number of selected blocks is $k + 1$ rather than $k$ — the $k$ highest-scoring blocks from the top-k mechanism plus the current block. The total attention span per query is therefore $(k + 1) \times B$ tokens. The sparsity ratio — the fraction of the context that is not attended to — is:

$\text{sparsity} = 1 - \frac{(k+1) \times B}{N} = 1 - \frac{k+1}{n}$

For the parameters used in the scaling law experiments (Section 3.1): $B = 512$, $k = 3$ (so the top-k selects 3 blocks plus the current block = 4 blocks total), and $N = 8192$, giving sparsity of $1 - \frac{4 \times 512}{8192} = 1 - 0.25 = 0.75$, or 75%. For the 1M-context model: $B = 4096$, $k = 12$ (top-k selects 12 plus current = 13 blocks), $N = 1,048,576$, giving sparsity of $1 - \frac{13 \times 4096}{1,048,576} \approx 1 - 0.051 = 0.949$, or approximately 95% — the model ignores 95% of the context for each query.

---

Causality Enforcement: Preventing Information Leakage in a Block-Based Architecture

Autoregressive language models generate text by predicting the next token conditioned on all previous tokens — a token at position $t$ cannot see tokens at positions $> t$. Standard attention enforces this through a causal mask: attention scores for future positions are set to $-\infty$ before the softmax, making their post-softmax weights zero. MoBA, because it operates at block granularity, requires two distinct causality mechanisms.

Mechanism 1: No attention to future blocks. The gating network must not route a query to any block that contains tokens after the query's position. Formally, for a query at position $\text{pos}(q)$, any block $i$ where $\text{pos}(q) < i \times B$ (the block starts after the query) is excluded. This is enforced by setting the affinity score to $-\infty$ before the top-k selection:

"we set $s_i = -\infty$ and $g_i = 0$ for any blocks $i$ such that $\text{pos}(q) < i \times B$" (Section 2.2)

Setting $s_i = -\infty$ ensures the block will never be in the top-k (since any finite score beats $-\infty$); explicitly setting $g_i = 0$ provides a belt-and-suspenders guarantee.

Why this is necessary: Without this mechanism, the top-k gating could route a query to a future block if that block happened to have a high affinity score — for example, if the query is "the next chapter discusses" and a future block contains related content. This would violate causality and cause the model to cheat during training (peeking at future tokens it shouldn't see). The block-level causal masking is the primary mechanism for preventing such leakage.

Mechanism 2: Causal masking within the current block. The current block contains both tokens before and after the query token. Since the current block is always attended to (as a shared expert), special care is needed to prevent the query from seeing tokens that come after it within the same block. The solution is straightforward: standard causal masking is applied during the current block's attention computation. The paper states:

"we enforce that each token must be routed to its respective current block and apply a causal mask during the current block attention" (Section 2.2)

This means that when computing attention between a query and the keys in its current block, any key at a position $> \text{pos}(q)$ has its attention score set to $-\infty$, exactly as in standard causal attention.

Why a separate mechanism for the current block: The block-level causal masking (mechanism 1) operates at block granularity — it prevents routing to entire blocks that start after the query. But within the current block, the query sits somewhere in the middle (or at the end), and tokens after it within the same block would not be caught by mechanism 1. The within-block causal masking fills this gap.

What about historical blocks: For blocks that are entirely before the query's position — which is all blocks selected by top-k (since future blocks are masked out) — no causal masking is needed within the block, because every token in a fully historical block precedes the query. The paper's Algorithm 1 confirms this: line 13 computes current-block attention with causal=True, while line 14 computes historical-block attention with causal=False. This is an important efficiency consideration: causal masking adds overhead (it requires modifying the attention scores before softmax), and avoiding it for historical blocks saves computation.

Why not just use per-token causal masking everywhere: The simplest approach would be to let the gating network select blocks freely and then apply per-token causal masking during attention computation — any future tokens would get $-\infty$ scores regardless of which block they're in. The paper's block-level masking serves an additional purpose beyond correctness: it prevents the model from wasting its limited attention budget on blocks that are mostly or entirely in the future. If a query had to spend one of its $k$ block selections on a block where 90% of tokens are causally masked out, the effective attention span would be significantly reduced. By excluding future blocks entirely, MoBA ensures that every selected block is maximally utilized.

---

Fine-Grained Block Segmentation: More Blocks, Smaller Blocks

The paper draws on a finding from the MoE literature — that fine-grained expert segmentation (having more, smaller experts rather than fewer, larger ones) improves model performance (Dai et al., 2024; Yang et al., 2024) — and investigates whether the same principle applies when segmenting the context dimension rather than the FFN dimension.

The experimental setup (Section 3.1, Figure 4): Using a 1.5B parameter model with 32K context length, the authors vary the number of blocks $n$ while keeping the attention sparsity constant at 75%. The configurations tested are:

• 8 blocks, selecting 2 ($k=2$, blocks of size $B = 4096$)
• 16 blocks, selecting 4 ($k=4$, $B = 2048$)
• 32 blocks, selecting 8 ($k=8$, $B = 1024$)
• 64 blocks, selecting 16 ($k=16$, $B = 512$)
• 128 blocks, selecting 32 ($k=32$, $B = 256$)

In each case, the total number of attended tokens is the same: $(k+1) \times B \approx N/4$ (with the +1 accounting for the current block). What varies is the granularity — at one extreme (8 blocks), the model makes coarse decisions about large chunks of context; at the other (128 blocks), it makes fine-grained decisions about small chunks.

The result: The paper reports "a performance difference of 1e-2 between the coarsest-grained setting (selecting 2 blocks from 8) and the settings with finer granularity" (Section 3.1), with LM loss decreasing as granularity increases. The coarsest setting achieves approximately 2.26 loss while finer settings reach approximately 2.24-2.25 (read from Figure 4). The full attention baseline is at approximately 2.23.

Why fine-grained segmentation helps: The authors interpret this result through the MoE lens: finer segmentation allows the routing mechanism to make more precise decisions about which content to attend to. With only 8 blocks of size 4096, each block contains a large, heterogeneous chunk of the context — a block might contain some highly relevant tokens and many irrelevant ones, and the model can't separate them. With 128 blocks of size 256, the model can selectively attend to exactly the small segments that matter while skipping similar-sized segments that don't.

However, there is a tradeoff: more blocks means the gating network must make more decisions (computing affinity scores for 128 blocks vs. 8 blocks), and the top-k selection must operate over a larger set. The paper doesn't explore the extreme where block size approaches 1 (token-level routing), which would maximize precision but make the routing cost comparable to the attention cost — defeating the purpose. The choice of block size is a hyperparameter that balances routing precision against routing overhead.

---

Hybrid of MoBA and Full Attention: Seamless Transitions and Layer-Wise Strategies

A distinctive feature of MoBA is that it is designed as a substitute for full attention, not an augmentation. As the paper states:

"MoBA is designed to be a substitute for full attention, maintaining the same number of parameters without any addition or subtraction. This feature inspires us to conduct smooth transitions between full attention and MoBA." (Section 2.2)

This design choice — same parameter count, same input/output shapes, same position in the Transformer architecture — means that switching between MoBA and full attention requires only changing the attention computation function, not modifying weights or adding/deleting layers.

MoBA/Full Hybrid Training (Section 3.2, Figure 5a): The paper demonstrates that MoBA can be used for most of training, with a switch to full attention near the end, to combine the efficiency of sparse attention with the quality of full attention. The specific recipe tested:

• Stage 1: Train with MoBA on 90% of the training tokens (27B out of 30B total tokens, using a 1.5B model with 32K context, block size 2048, top-k=3).
• Stage 2: Switch to full attention for the remaining 10% of tokens (3B tokens).

The key empirical finding is that this hybrid recipe "reaches a loss nearly identical to that of full attention" (Section 3.2) on position-wise LM loss for trailing tokens, while the MoBA-only recipe shows noticeably higher loss at distant positions. Perhaps most importantly:

"we have not observed significant loss spikes during the switch between MoBA and full attention, again demonstrating the flexibility and robustness of MoBA." (Section 3.2)

Why this works: MoBA is trained to produce meaningful attention outputs despite seeing only a subset of the context. When switched to full attention, the model's parameters — the query, key, value, and output projections — are already well-trained; they just get access to more context. The lack of a loss spike suggests that the model doesn't overfit to the sparse attention pattern during Stage 1 and can immediately leverage the additional context in Stage 2. This is a significant practical property: it means organizations can train the bulk of their long-context models efficiently with MoBA and then "top off" with full attention to close any remaining performance gap.

Why the switch helps: The MoBA-only model performs worse on trailing tokens because those tokens are far from the beginning of the sequence and have many blocks of history to potentially attend to. The top-k gating might miss relevant historical blocks for these late-position queries. Full attention, by granting access to all tokens, eliminates this possibility. A short full-attention phase at the end of training allows the model to learn to use the full context without paying the quadratic cost for the entire training run.

Layer-Wise Hybrid (Section 3.2, Figures 5b and 5c): A more sophisticated strategy involves using MoBA in some layers and full attention in others simultaneously, rather than switching over time. The paper investigates this specifically for supervised fine-tuning (SFT), motivated by an empirical observation:

"MoBA sometimes results in suboptimal performance during SFT" (Section 3.2)

The authors hypothesize that this is due to the loss masking employed in SFT. During SFT, the loss is typically computed only on the model's response tokens, not on the prompt tokens (since the goal is to learn to generate responses, not to model the prompt). This means gradients flow only from the response tokens backward through the network. For sparse attention methods like MoBA, gradients from response tokens must propagate through the attention layers to update the key/value projections and the gating network. If the response tokens' queries don't route to certain key blocks that the prompt tokens would have routed to, those key blocks receive no gradient signal — creating a "sparse gradient" problem.

The proposed solution: Switch the last several Transformer layers from MoBA to full attention, keeping MoBA in earlier layers. The intuition: later layers are closer to the output (and thus closer to the loss computation on response tokens), so making them full attention ensures that gradient signal can flow freely through the most critical layers. Earlier layers, which process the prompt more abstractly, can remain sparse without as much impact.

The paper tests configurations with 1, 3, 5, and 10 full-attention layers (out of presumably 20, based on the cited 1.5B model architecture). Results (Figures 5b and 5c):

• SFT LM loss (Figure 5b): Even 1 full-attention layer noticeably reduces loss compared to pure MoBA; 3-5 layers approach the performance of full attention; 10 layers essentially match full attention.
• SFT trailing LM loss (Figure 5c): The pattern is similar but more pronounced — trailing tokens benefit more from additional full-attention layers, consistent with the earlier finding that trailing tokens are the most challenging for sparse attention.

Why the layer-wise strategy matters: It provides a knob for trading off efficiency against performance. Deployments that need maximum throughput can use MoBA in all layers; deployments that need maximum quality can use a few full-attention layers at the end; deployments in between can pick an intermediate configuration. This is far more flexible than the binary choice between "all sparse" and "all full" attention.

The 1M-context model described in Section 3.3 uses exactly this strategy: "the last three layers remain as full attention, while the other 29 full attention layers are switched to MoBA" (Section 3.3). This means 3 out of 32 layers use full attention — approximately 9% of layers — providing a modest quality boost at a modest efficiency cost.

---

Implementation: Mapping the Algorithm to Efficient GPU Kernels

The mathematical specification of MoBA — select blocks, compute block-wise attention, merge outputs — would be straightforward to implement naively but would be inefficient without careful optimization. The paper provides a five-step implementation strategy (Algorithm 1) that maps the algorithm onto efficient GPU primitives, achieving the speedups reported in Section 3.4.

The implementation, visualized in Figure 1b, proceeds as follows (referencing Algorithm 1 line numbers):

Step 1 (Lines 1-8): Determine query-to-block assignments.

• The KV matrices ($K, V \in \mathbb{R}^{N \times h \times d}$, where $h$ is the number of attention heads) are split into $n$ blocks of size $B$ along the sequence dimension (Line 2).
• Key vectors are mean-pooled within each block to produce $\bar{K} \in \mathbb{R}^{n \times h \times d}$ (Line 4).
• Affinity scores are computed as $S = Q\bar{K}^{\top} \in \mathbb{R}^{N \times h \times n}$ — every query against every block, for every head (Line 5).
• A causal mask $M$ is created (Line 7), setting entries where the block is in the future to $-\infty$.
• The top-k operator is applied to $S + M$, producing the sparse binary assignment matrix $G$ indicating which blocks each query attends to (Line 8).

Step 2 (Lines 9-12): Reorganize queries by block assignment. This is the critical step for GPU efficiency. Rather than iterating over queries and computing attention for each one individually (which would be serial and slow), the implementation:

• Separates queries into those attending to the current block (Line 10) and those attending to historical blocks (Line 11). These are handled differently because current-block attention needs causal masking while historical-block attention doesn't.
• For historical block attention, queries are grouped by their assigned blocks — all queries that will attend to block $i$ are collected together, producing a variable-sized batch of queries for each block. This grouping is what enables FlashAttention to operate efficiently on contiguous chunks of memory.

Step 3 (Lines 13-14): Compute block-wise attention with FlashAttention.

• Current block attention (Line 13): $\text{flash\_attention\_varlen}(Q_s, \tilde{K}_s, \tilde{V}_s, \text{causal=True})$. This computes attention between each query and its current block, with causal masking applied. Because queries within the same current block are contiguous in sequence order, this can be implemented as standard causal FlashAttention with minor indexing adjustments.
• Historical block attention (Line 14): $\text{flash\_attention\_varlen}(Q_m, \tilde{K}_m, \tilde{V}_m, \text{causal=False})$. This computes attention between grouped queries and their assigned historical blocks. Because queries are grouped by target block, each block's keys and values are accessed once and used for all queries assigned to it — maximizing memory locality. No causal mask is needed because all tokens in historical blocks precede all assigned queries.

The use of varlen (variable-length) FlashAttention kernels is important: different blocks will have different numbers of queries assigned to them (since the top-k gating can route different numbers of queries to different blocks). FlashAttention's variable-length mode handles this by accepting arrays of sequence lengths and computing attention efficiently regardless of uniformity.

Step 4 (Line 16): Combine outputs with online softmax. A query token may attend to multiple blocks — its current block plus up to $k$ historical blocks. Each of these produces a partial attention output (a weighted sum of value vectors within that block). These cannot be naively averaged because the softmax normalization in each block is computed independently — the relative weights assigned to tokens in block A versus tokens in block B need to reflect the original (un-normalized) attention scores across all blocks.

The solution is online softmax (also called tiling), a technique that allows independently computed softmax-normalized outputs to be merged into the correct overall result. The key idea: store the un-normalized sum (the denominator of the softmax, also called the normalization constant or log-sum-exp value) alongside each partial output. When merging, rescale each partial output by its normalization constant relative to the combined normalization constant, producing the same result as if softmax had been computed over all attended tokens jointly.

The paper cites two references for online softmax: Milakov and Gimelshein (2018) for the original technique, and Liu and Abbeel (2023) for its application in blockwise parallel Transformers.

Why this implementation is efficient (Section 3.4, Figures 2a and 2b):

The paper reports that MoBA achieves a 6.5× speedup at 1M tokens and a 16× speedup at 10M tokens compared to FlashAttention, with sub-quadratic scaling (Figure 2b shows FlashAttention growing much faster than MoBA as sequence length increases). The efficiency comes from two sources:

1. Reduced computation: Each query attends to $(k+1) \times B$ tokens rather than $N$ tokens. At 1M tokens with $B=4096$ and $k=12$, this is 53,248 tokens per query rather than 1,048,576 — a 95% reduction in FLOPs.

2. Efficient memory access patterns: By grouping queries by their target blocks, each block's keys and values are loaded from GPU memory once and reused for all queries assigned to that block. This minimizes the expensive data movement between GPU global memory and compute units, which is the primary bottleneck in attention computation (as FlashAttention's design demonstrates).

The paper also addresses GPU memory limitations when scaling to 10M tokens (Section 3.4): "we expanded tensor parallelism (Shoeybi et al., 2019) toward the query head level, specifically, we broadcast key and value tensors across distributed query heads, effectively addressing GPU memory limitations while preserving computational efficiency." This means that the key and value tensors — which grow linearly with sequence length — are distributed across multiple GPUs, with each GPU handling a subset of attention heads but having access to the full KV context via broadcasting.

Prefill vs. generation modes: The paper notes an important deployment distinction (Section 3.3): "across all evaluation tasks, MoBA is used for prefill only, while we switch to full attention during generation for better performance." During prefill (processing the input prompt), the full prompt is available and MoBA's block-sparse attention provides maximum efficiency. During generation (producing output tokens one at a time), the model switches to full attention. This makes sense because during generation, the sequence length is typically much shorter (the generated portion is small compared to the prompt), so the quadratic cost is less severe, and the quality benefit of full attention outweighs the modest efficiency cost.

---

Summary of Key Design Choices and Their Justifications

• Block-level partitioning over token-level routing: Blocks reduce the number of routing decisions from $N$ to $n$, making the gating mechanism computationally cheap. Blocks also provide cache-friendly memory access patterns for GPU computation.
• Mean-pooled key vectors as block representations: Parameter-free, pre-computable once per layer, and reuses existing key projections — adds no new parameters. The simplicity leaves room for future improvements (the paper notes Quest's min+max pooling as a potential enhancement).
• Top-k gating with fixed $k$: Provides a fixed, predictable computational budget per query regardless of context length, inherits the competitive specialization dynamics from MoE, and is simple to implement.
• Current block as forced "shared expert": Ensures local context is always attended to, prevents degenerate routing away from the query's neighborhood, and provides the natural location for causal masking.
• Separate causal masking for current vs. historical blocks: Historical blocks need no causal masking (all tokens precede the query), saving computation. Current block attention uses standard causal masking to prevent within-block future-token leakage.
• Fine-grained block segmentation (more, smaller blocks): Empirically improves LM loss by approximately 0.01-0.02, interpreted as enabling more precise routing decisions — the model can focus on exactly the small segments that matter rather than large heterogeneous chunks.
• Hybrid training (MoBA → full attention switch): Closes the trailing-token performance gap by giving the model access to full context during the final training phase, without paying quadratic cost for the entire training run.
• Layer-wise hybrid (full attention in final layers): Mitigates sparse gradient issues during SFT by ensuring gradient flow through the layers closest to the loss computation, while keeping earlier layers efficient.
• Online softmax for output combination: Enables numerically correct merging of independently computed block-wise attention outputs without recomputing softmax over the full attended set — essential for the block-wise computation strategy.
• FlashAttention integration with query-to-block grouping: Maximizes GPU efficiency by loading each block's KV cache once and reusing it for all queries assigned to that block, minimizing memory bandwidth bottlenecks.

4. Key Insights and Innovations

Innovation 1: Reframing Attention Sparsity as a Mixture of Experts Routing Problem

The paper's most fundamental conceptual move is not inventing block-sparse attention — which has existed since at least BlockBERT (Qiu et al., 2019) and Longformer (Beltagy et al., 2020) — but rather re-conceptualizing the attention sparsity problem through the lens of Mixture of Experts routing. This reframing is what makes MoBA distinctive, and it carries intellectual weight beyond the specific architectural choices.

Prior work on dynamic sparse attention (Reformer, Kitaev et al., 2020; Routing Transformer, Roy et al., 2021; Quest, Tang et al., 2024) treated sparsity as a selection problem: given a query, identify which individual tokens are most relevant and attend only to those. The dominant paradigm was token-level selection, typically using locality-sensitive hashing, k-means clustering, or k-nearest-neighbor search over key vectors. These methods share an implicit assumption that the unit of selection should be the individual token — the same granularity at which attention operates.

MoBA challenges this assumption by asking: what if we treat context blocks as "experts" and let the model learn to route queries to blocks, rather than selecting individual tokens via similarity search? This is more than a change of granularity. The MoE framing imports a specific set of ideas about how routing should work:

• Learned, competitive routing rather than similarity-based retrieval. In token-level dynamic sparsity, the selection mechanism is typically a non-parametric similarity function (dot product, LSH hash, Euclidean distance) between the query and individual keys. MoBA's gating mechanism, by contrast, computes an affinity score between the query and a learned block representation (the mean-pooled keys). While the dot-product computation is parameter-free, the representations being compared — the key vectors themselves — are learned parameters updated during training specifically to make the routing work. This subtle shift means the model learns to organize its key representations to facilitate effective block-level routing, rather than routing being a post-hoc operation on representations optimized solely for attention computation.

• Fixed budget allocation with competition. Token-level selection methods typically use a similarity threshold or a fixed top-p/m value, which can result in variable numbers of selected tokens per query. MoBA's top-k gating imposes a strict, uniform budget — every query attends to exactly $(k+1) \times B$ tokens regardless of content. This creates what the MoE literature calls competition among experts: blocks must compete for the limited attention budget, and the model must learn to produce affinity scores that reliably push relevant blocks above irrelevant ones. This competitive dynamic is absent from threshold-based or unrestricted top-p selection.

• The shared expert concept mapped to the current block. The paper explicitly draws the analogy between MoBA's forced current-block attention and the shared experts in architectures like DeepSeekMoE (Dai et al., 2024) and Qwen2.5 (Yang et al., 2024). This mapping is not merely cosmetic — it suggests that the MoE design principle of "some computation should be universally applied, with additional computation selectively routed" applies to the attention dimension just as it does to the FFN dimension. The current block serves the same function as a shared FFN expert: guaranteeing a baseline of local context processing that prevents the model from making catastrophic routing errors (e.g., routing all attention away from the query's immediate neighborhood).

This reframing is fundamental rather than incremental because it changes what the field should optimize for in sparse attention design. Under the token-selection paradigm, the primary research question is "how do we efficiently find the most similar keys?" — which leads to work on better hashing, better approximate nearest-neighbor search, and better similarity functions. Under the MoE routing paradigm, the question becomes "how do we train the model to produce representations that support effective block-level routing?" — which leads to work on routing regularization, load balancing, expert specialization, and training dynamics. The paper imports a mature set of techniques and research questions from the MoE literature into the attention efficiency domain, opening a new design space that the token-selection paradigm didn't access.

Evidence: The paper substantiates this reframing through the fine-grained segmentation experiments (Section 3.1, Figure 4), which show that MoBA inherits the MoE property that more, smaller experts improve performance — "fine-grained segmentation appears to be a general technique for enhancing the performance of models within the MoE family, including MoBA." The fact that principles from FFN MoE transfer to attention MoBA validates the conceptual mapping.

---

Innovation 2: Seamless Full-to-Sparse Attention Transition as a First-Class Design Goal

Most work on efficient attention treats the sparse mechanism as a replacement for full attention — once you adopt the efficient variant, you commit to it. The sliding window in Longformer replaces full attention; the hashing in Reformer replaces the dot-product over all keys; the state space in Mamba replaces attention entirely. MoBA makes a different and more subtle choice: treat the ability to switch between sparse and full attention as a fundamental architectural property, not an afterthought.

The significance of this choice becomes clear when contrasted with the dominant assumptions in prior work:

• Static sparse patterns (Longformer, BigBird, LongNet): The sparsity pattern is hard-coded into the architecture. Switching to full attention would require changing the attention computation function and potentially retraining, but more fundamentally, the model was never trained to use the full context — its parameters are optimized for the sparse pattern, and there's no guarantee that giving it access to more tokens would help. The model may have learned to compensate for the missing connections in ways that break when those connections are restored.

• Linear attention models (Mamba, RWKV, Performer): These models use fundamentally different mathematical operations than standard attention — state space recurrences, kernel approximations, or low-rank factorizations. Switching to full attention requires converting the entire architecture, which "typically incurs high conversion costs" (as the paper notes, citing Mercat et al., 2024; Wang et al., 2024; Bick et al., 2025). The models are not just trained with a different attention pattern; they are trained with a different type of computation.

• Inference-time dynamic sparsity (Quest, Minference): These methods do enable switching — they apply sparsity on top of a model trained with full attention. But the switching is one-directional: the model was trained with full attention, and sparsity is applied only at inference. MoBA enables the reverse: train with sparse attention, then switch to full attention for fine-tuning or for specific layers, and the model adapts smoothly because both modes use the same mathematical operation (softmax dot-product attention) with the same parameters.

MoBA's design achieves this seamlessness through a specific architectural decision that is easy to overlook: MoBA uses exactly the same attention computation as the standard Transformer, just applied to a subset of the context. The projection matrices $W_Q, W_K, W_V, W_O$ are identical; the softmax normalization is identical; the weighted value aggregation is identical. MoBA is not a new attention function — it is standard attention with a learned input filter. This means that the model's parameters mean the same thing in both modes: a well-trained query projection still produces good query vectors whether those queries attend to 50K tokens or 1M tokens; a well-trained value projection still produces good value vectors that can be meaningfully aggregated.

The empirical demonstration of this property is the MoBA/Full hybrid training experiment (Section 3.2, Figure 5a): training with MoBA for 90% of tokens, then switching to full attention for the final 10%, produces position-wise LM loss "nearly identical to that of full attention" with "no significant loss spikes during the switch." This is a strong result because loss spikes at architecture transitions are common — they indicate that the model's representations are incompatible with the new computation mode. The absence of a spike validates the claim that MoBA and full attention share a compatible parameter space.

Why this matters beyond the specific experiment:

• Continual pretraining on existing models: Organizations with pre-trained Transformer models can adopt MoBA without starting from scratch. They continue training with MoBA to gain efficiency, then optionally switch back to full attention for fine-tuning. This dramatically lowers the adoption barrier compared to approaches requiring architectural conversion or retraining from scratch.

• Gradual deployment with tunable quality-efficiency tradeoffs: The layer-wise hybrid strategy (Section 3.2, Figures 5b-5c) provides a continuous knob: 0 full-attention layers (maximum efficiency), 3 full-attention layers (the 1M-context model's configuration), or all layers full attention (maximum quality). This is not a binary choice between "fast but inaccurate" and "accurate but slow" — it's a spectrum that can be tuned per deployment scenario.

• Training-time efficiency without inference-time compromise: The hybrid training recipe means the bulk of long-context training can be done with sparse attention (saving 4-16× in FLOPs), while the final model can use full attention during inference if latency permits. The training savings apply regardless of deployment choice.

This is an incremental insight in terms of mechanism (the hybrid training recipe has been studied for sliding window attention; Zhang et al., 2024) but fundamental in terms of architectural philosophy: it establishes that sparse attention mechanisms should be designed for compatibility with full attention, not as independent alternatives. The field's default assumption — that you pick either sparse or full attention and commit to it — is replaced by a model where the two modes are interoperable stages in a training pipeline.

Evidence: Figure 5a shows the hybrid recipe matching full attention on position-wise LM loss; Figures 5b-5c show the layer-wise hybrid closing the SFT performance gap. The practical deployment of the 1M-context model with 3 full-attention layers out of 32 (Section 3.3) demonstrates production viability.

---

Innovation 3: Trailing Tokens as the Diagnostic for Long-Context Attention Quality

The paper makes a methodological contribution that is easily missed but has significant implications for how the field evaluates long-context models: the identification of trailing token LM loss as the primary signal for long-context capability, and the corresponding finding that trailing tokens account for most of the performance gap between sparse and full attention.

Prior work on long-context evaluation has focused on two types of metrics:

• Aggregate benchmarks: RULER, LongBench, Needle-in-a-Haystack — these measure task performance (question answering, retrieval, summarization) across various context lengths. They are end-to-end metrics that capture whether the model can use long context, but they don't diagnose where the model struggles.

• Perplexity / LM loss over the full sequence: Standard language modeling loss averages the cross-entropy across all token positions in the validation set. However, as the paper notes (citing An et al., 2024), this metric "may be skewed by the data length distribution, which is typically dominated by short sequences." If most training/validation sequences are short, the average LM loss is dominated by early-position tokens, masking problems that only appear at long contexts.

MoBA introduces a finer-grained diagnostic: position-wise LM loss with specific focus on trailing tokens. The key methodological moves are:

1. Segment sequences by position range (e.g., 0-2K, 2K-4K, ..., 30K-32K) rather than averaging over all positions. This reveals a clear pattern: MoBA matches full attention almost perfectly at early positions (0-8K) but shows a growing gap at later positions (Figure 8 in Appendix A.1, and the scaling law fits in Table 3).

2. Compute trailing LM loss only on sequences that actually reach the maximum length, avoiding the bias introduced by short sequences where "trailing" tokens are actually near the beginning of the document.

3. Fit separate scaling laws for each position range (Table 3), revealing that the scaling exponent for MoBA degrades more rapidly with position than for full attention — the trailing-token scaling exponent for MoBA is $-0.108$ versus $-0.097$ for full attention (at 30K-32K positions, Figure 3c), a meaningful difference that would be invisible in aggregate metrics.

The diagnostic power of this approach is demonstrated by the paper's own use of it:

• Scaling law experiments (Section 3.1, Figure 3b): The trailing LM loss comparison reveals that while MoBA's overall LM loss is nearly identical to full attention (Figure 3a, within 1e-3), the trailing-token gap is larger and — critically — is narrowing with model scale. This suggests that larger models are better able to route effectively over long distances, a trend that aggregate metrics would smooth over.

• Hybrid training evaluation (Section 3.2, Figure 5a): The position-wise LM loss curves show that MoBA-only training produces elevated loss specifically at trailing positions, while the hybrid recipe recovers full-attention-level performance at all positions. This pinpoints exactly what the full-attention fine-tuning phase fixes: it's not improving overall language modeling, it's specifically fixing the model's ability to use distant context at the end of long sequences.

• Layer-wise hybrid SFT evaluation (Section 3.2, Figures 5b-5c): The trailing LM loss metric (Figure 5c) shows a larger gap between MoBA and full attention than the overall LM loss (Figure 5b), and a larger benefit from adding full-attention layers. This confirms that the SFT gradient sparsity issue primarily affects the model's ability to process distant context.

This methodological contribution is incremental in technique (position-wise loss is straightforward to compute) but fundamental in its implications for how the field should evaluate sparse attention methods. The finding that "trailing tokens account for the majority of the performance discrepancy between the full context baseline and the newly proposed sparse attention architectures" (Appendix A.1) means that long-context evaluation should be specifically designed to expose trailing-token behavior, not just aggregate performance. A sparse attention method that matches full attention on average perplexity but shows a 0.1 gap on trailing tokens (as MoBA does at smaller scales) is meaningfully worse at long-context processing, even if the aggregate metric doesn't show it.

The paper effectively argues — through its scaling law analysis — that trailing-token scaling should be the primary evaluation axis for long-context architectures, with the fitted scaling exponent serving as a summary statistic for how well the method handles increasing context length at the most challenging positions. This is a more informative metric than aggregate perplexity, benchmark scores (which can be gamed), or retrieval accuracy (which tests a specific capability rather than general long-context processing).

Evidence: The scaling law fits in Table 3 show progressively diverging exponents between MoBA and full attention as position increases; Figure 5a shows the trailing-token loss gap that hybrid training closes; Figure 3b shows this gap narrowing with model scale.

---

Innovation 4: Verifier Over-Optimization Is Not the Bottleneck — Representation Granularity Is

The paper arrives at a counterintuitive empirical finding through its fine-grained segmentation ablation (Section 3.1, Figure 4): block granularity, not total attention budget, is the primary determinant of MoBA's performance. When the attention sparsity is held constant at 75% — meaning the model attends to exactly the same number of tokens regardless of block configuration — moving from 8 blocks (coarse, 4096 tokens each) to 64-128 blocks (fine, 256-512 tokens each) improves LM loss by approximately 0.01-0.02, closing about half the gap to full attention.

This finding challenges a natural intuition about sparse attention: that performance should primarily depend on how many tokens are attended to (the attention budget), not how those tokens are organized. Under that intuition, attending to 25% of tokens should yield roughly the same performance whether those tokens are organized as 4 large blocks or 128 small blocks — after all, the same total information is accessible. Figure 4 shows this intuition is wrong.

What makes this finding intellectually significant is what it reveals about the nature of the attention sparsity problem:

• The bottleneck is selection precision, not selection quantity. If the model had perfect routing — always selecting exactly the 25% of tokens most relevant to each query — then block granularity wouldn't matter; any grouping of those tokens would work. The fact that finer granularity helps means the routing is imprecise: when blocks are large, they contain a mix of relevant and irrelevant tokens, and the model cannot separate them. Finer blocks allow the routing mechanism to make more precise inclusion/exclusion decisions, reducing the amount of irrelevant content that gets swept in alongside the relevant content.

• The MoE analogy extends to the failure mode as well as the mechanism. In standard MoE, coarse experts (few experts, each with large capacity) suffer from the problem that each expert must handle a diverse set of inputs, preventing specialization. Fine-grained experts (many experts, each with small capacity) allow each expert to specialize narrowly. MoBA exhibits the same dynamic: coarse blocks force each block to contain heterogeneous content, making the routing decision a blunt instrument; fine blocks allow the model to route to precisely the content it needs.

• This is a representation problem, not a search problem. Prior work on dynamic sparse attention has largely treated the challenge as one of efficient search — how to quickly find the most similar keys among millions of candidates (LSH, kNN, clustering). MoBA's fine-grained segmentation result suggests that the more fundamental challenge is representational: the model's key representations must be organized such that routing can work effectively. Finer blocks don't change the search mechanism (it's still top-k gating over the same number of blocks; actually more blocks means more gating decisions), but they change the target structure that the model's representations must support — forcing keys within each block to be coherent enough that attending to the whole block is beneficial.

This insight has practical implications that go beyond MoBA: it suggests that future work on sparse attention should focus on improving the coherence of attention blocks (through better partitioning strategies, learned rather than fixed block boundaries, or hierarchical routing) rather than solely on improving the efficiency of token-level retrieval. The paper's own comparison to Quest — which uses min+max pooling for block representation — hints at this direction, suggesting that better block representations could further improve routing precision.

This is an incremental empirical finding (it's a single ablation study) but a fundamental conceptual reframing of the sparse attention bottleneck: the problem is not "how do we find the right tokens?" but "how do we organize tokens so that coarse-grained selection works?" The field's attention (pardon the pun) has been on the former; MoBA's results suggest shifting toward the latter.

Evidence: Figure 4 shows a clear monotonic improvement in LM loss as block granularity increases from 8 to 128 blocks at constant 75% sparsity, with the gap between the coarsest and finest configurations being approximately 0.01-0.02 in loss.

5. Experimental Analysis

Evaluation Methodology

• Dataset and Training Data. The primary experiments use an unspecified large-scale language modeling corpus for pre-training and continual pre-training, with validation performed on held-out text. For downstream evaluation, the paper uses a broad suite of standard benchmarks: AGIEval, BBH, CEval, GSM8K, HellaSWAG, Loogle, Competition Math, MBPP, MBPP Sanitized, MMLU, MMLU Pro, HumanEval, SimpleQA, TriviaQA, LongBench, and RULER — all evaluated at 0-shot or few-shot as specified in Table 2. Long-context capability is specifically tested with the Needle-in-a-Haystack benchmark at up to 1M tokens (Figure 7) and RULER at 128K tokens (Table 2).

• Base Model. The main large-scale experiments use Llama 3.1 8B as the starting point. The paper argues this model is representative of contemporary open-weight LLMs at a scale where long-context pre-training is computationally intensive but feasible. For scaling law experiments (Section 3.1), the authors train five models from scratch at sizes 568M, 822M, 1.1B, 1.5B, and 2.1B parameters, with configurations specified in Table 1. The 1.5B model is used for ablation studies on fine-grained segmentation (Figure 4) and hybrid training strategies (Section 3.2).

• Metrics. The paper employs three distinct evaluation regimes. Language modeling loss (LM loss): standard cross-entropy loss averaged over all validation tokens, used for scaling law comparisons (Figures 3a, 3c). Trailing LM loss: cross-entropy loss computed only on the final tokens of sequences that reach maximum length — specifically the "last 2K" tokens for 32K-context models (Figures 3b, 5c, and Appendix A.1) — used to isolate long-context capability from the bias introduced by short training sequences. Position-wise LM loss: LM loss broken down by token position within the sequence, binned into 2K-token segments (Figure 5a and Figure 8 in Appendix A.1), used to diagnose where sparse attention degrades relative to full attention. For downstream tasks, standard benchmark-specific metrics are used (accuracy, pass@1, F1, etc.) as reported in Table 2.

• Baselines. The primary baseline throughout is full attention — standard causal self-attention implemented with FlashAttention — trained under identical conditions (same data, same hyperparameters, same training duration) as MoBA models. This is the gold-standard comparison because any sparse attention method must demonstrate that it does not sacrifice quality relative to full attention under matched training budgets. For the scaling law experiments (Figures 3a-3b), five full-attention models are trained at the same five scales as the MoBA models. For the 1M-context experiments (Table 2, Figure 7), a full-attention model termed Llama-8B-1M-Full is trained with the same continual pre-training and SFT recipes as Llama-8B-1M-MoBA, with only the attention mechanism differing. The paper does not compare against other sparse attention methods (Longformer, BigBird, Quest, Minference, etc.) in quantitative experiments — these are discussed only in the related work (Section 4) — which is a notable omission in the experimental design.

• Generation Budget / Compute Accounting. Computational cost is measured in two ways. For the scaling law experiments, cost is measured in PFLOP/s-days (petaFLOP/s-days), following the Chinchilla scaling law convention (Hoffmann et al., 2022), with models trained to their compute-optimal point. For the efficiency benchmarks (Section 3.4, Figure 2), cost is measured as wall-clock computation time (milliseconds or seconds) for the attention layer forward pass, enabling direct comparison of MoBA against FlashAttention at identical sequence lengths. Attention sparsity is reported as the fraction of key-value pairs that are NOT attended to, computed as 1 - ((k+1) × B) / N, where k is the top-k parameter, B is block size, and N is sequence length. The paper consistently reports sparsity alongside speedup to contextualize efficiency gains — e.g., 95.31% sparsity corresponding to the 6.5× speedup at 1M tokens (Section 3.4).

• Training Protocol for Fair Comparison. The paper emphasizes that "the only difference across all experiments lies in the attention modules, while all other hyperparameters, including the learning rate and batch size, remain constant" (Section 3.1). For the 1M-context comparison (Section 3.3), both MoBA and full attention models follow identical continual pre-training and SFT recipes, with context length gradually increased from 128K to 256K to 512K to 1M tokens using position interpolation (Chen et al., 2023) at the 256K stage. This controlled setup ensures that performance differences can be attributed to the attention mechanism rather than training data or hyperparameter differences.

---

Main Quantitative Results

Scaling Law Experiments: MoBA vs. Full Attention

Overall LM loss scaling (Figure 3a, Table 1). Across five model sizes (568M to 2.1B parameters) trained on 10.8B to 36.9B tokens respectively with 8K context length, MoBA with block size 512 and top-k=3 achieves validation LM loss nearly identical to full attention. The paper reports: "the validation loss differences between these two attention mechanisms remain consistent within a range of 1e−3" (Section 3.1). The fitted scaling law curves (Figure 3c) show:

• MoBA: LM loss (seqlen=8K) = 2.625 × C^(-0.063)
• Full attention: LM loss (seqlen=8K) = 2.622 × C^(-0.063)

The nearly identical scaling exponents (-0.063 for both) indicate that MoBA's sparsity (75% for these 8K experiments) does not degrade the fundamental scaling behavior — the models improve with compute at the same rate as full attention models. The constant factor difference (2.625 vs. 2.622) is minimal — approximately 0.003 nats — confirming that MoBA matches full attention in aggregate language modeling quality at these scales.

Trailing LM loss scaling (Figure 3b). When evaluated at 32K context length (where MoBA operates at 95.31% sparsity due to the fixed block size of 512), the trailing-token loss reveals a more nuanced picture. The fitted curves (Figure 3c) show:

• MoBA: Trailing LM loss (seqlen=32K, last 2K) = 1.546 × C^(-0.108)
• Full attention: Trailing LM loss (seqlen=32K, last 2K) = 1.464 × C^(-0.097)

Two observations are critical here. First, MoBA exhibits a higher constant factor (1.546 vs. 1.464) — the absolute trailing-token loss is worse for MoBA at all model scales. This is the long-context penalty that aggregate metrics mask. Second, and more encouragingly, MoBA shows a steeper scaling exponent (-0.108 vs. -0.097), meaning the gap narrows as model size increases. The paper notes: "the loss gap is progressively narrowing" (Section 3.1) — larger models handle the sparse attention constraint better, suggesting that routing quality improves with model capacity.

Position-wise scaling laws (Appendix A.1, Figure 8, Table 3). The granular position-wise analysis reveals where the performance gap lives. For positions 0-8K (Figures 8a-8d), MoBA and full attention scaling curves are nearly identical, with scaling exponents differing by at most 0.001. The gap progressively widens at later positions: at 30K-32K (Figure 8p), the exponent gap is -0.108 vs. -0.097 (a difference of 0.011) and the constant factor gap is 1.546 vs. 1.464 (a difference of 0.082). This monotonic degradation with position — from near-identical at early positions to measurably worse at trailing positions — is the paper's central diagnostic for sparse attention quality. The paper frames this as evidence that "trailing tokens account for the majority of the performance discrepancy between the full context baseline and the newly proposed sparse attention architectures" (Appendix A.1).

Interpretation. These scaling law experiments establish MoBA's viability at a fundamental level: the architecture achieves comparable scaling properties to full attention, with the performance cost concentrated almost entirely at the most distant positions — and that cost diminishing with model scale. A skeptic might note that the 1e-3 aggregate gap, while small, is consistent across scales and could compound in ways that affect downstream task performance; the trailing-token analysis directly addresses this concern by showing where the gap lives and that it shrinks with scale.

---

Fine-Grained Block Segmentation (Figure 4)

Using a 1.5B parameter model with 32K context length and constant 75% sparsity, the paper sweeps block configurations from 8 blocks (selecting 2) to 128 blocks (selecting 32). The results show:

• Full attention baseline: LM loss approximately 2.23 (read from Figure 4)
• 8 blocks (coarsest): LM loss approximately 2.26
• 16 blocks: LM loss approximately 2.25
• 32 blocks: LM loss approximately 2.245
• 64 blocks: LM loss approximately 2.24
• 128 blocks (finest): LM loss approximately 2.24

The paper reports "a performance difference of 1e-2 between the coarsest-grained setting and the settings with finer granularity" (Section 3.1), while noting that configurations from 32 blocks onward achieve similar performance (within approximately 0.005 of each other). The gap to full attention is approximately 0.01-0.02 for the finest configurations.

The key finding is that granularity matters substantially — coarse blocks (4096 tokens each) force the model to attend to large heterogeneous chunks where irrelevant tokens dilute the attention computation, while fine blocks (256-512 tokens each) enable more precise selection. The diminishing returns beyond 64 blocks suggest a practical sweet spot: enough granularity for precise routing, but not so many blocks that the gating overhead becomes significant.

---

Hybrid Training Strategies (Section 3.2, Figure 5)

MoBA/Full hybrid training (Figure 5a). Three 1.5B models are trained on 30B tokens with 32K context (MoBA parameters: block size 2048, top-k=3):

• Full attention only: Reference baseline.
• MoBA only: MoBA throughout training (30B tokens).
• MoBA/Full hybrid: MoBA for 27B tokens (90%), then switched to full attention for 3B tokens (10%).

Position-wise LM loss (Figure 5a) shows:

• Full attention achieves the lowest loss across all positions, with the characteristic U-shaped curve (lowest loss at middle positions, higher at very early and very late positions).
• MoBA only tracks full attention closely at early positions (0-10K) but diverges increasingly at later positions (15K-30K), with the gap reaching approximately 0.2-0.5 nats at the trailing edge (exact values difficult to read from Figure 5a log scale, but the visual separation is clear).
• MoBA/Full hybrid achieves loss "nearly identical to that of full attention" across all positions. The paper explicitly states: "we have not observed significant loss spikes during the switch between MoBA and full attention" (Section 3.2).

This result demonstrates that MoBA's parameters are fully compatible with full attention — switching attention modes mid-training does not cause the optimization instability that would be expected if the model had overfit to the sparse pattern. The hybrid recipe provides a practical training strategy: use MoBA for the bulk of training to save compute, then "top off" with full attention to close the trailing-token gap.

Layer-wise hybrid for SFT (Figures 5b and 5c). The paper investigates keeping the last few Transformer layers as full attention while using MoBA in the remaining layers, specifically for supervised fine-tuning where "MoBA sometimes results in suboptimal performance" (Section 3.2). Configurations tested: 1, 3, 5, and 10 full-attention layers out of the total (presumably 20 layers for the 1.5B model).

• SFT LM loss (Figure 5b): Pure MoBA achieves approximately 1.13 loss; adding 1 full-attention layer reduces loss to approximately 1.11; 3-5 layers reach approximately 1.09-1.10; 10 layers approach the full attention baseline at approximately 1.08. Each additional full-attention layer provides diminishing but consistent improvements.
• SFT trailing LM loss (Figure 5c, 32K context, last 2K): Pure MoBA achieves approximately 1.16; 1 full-attention layer: approximately 1.14; 3 layers: approximately 1.12; 5 layers: approximately 1.11; 10 layers: approximately 1.10; full attention: approximately 1.09. The trailing-token gap is larger than the overall gap, and the benefit of additional full-attention layers is more pronounced — confirming the paper's hypothesis that SFT loss masking creates sparse gradient challenges specifically for distant positions.

The paper hypothesizes that the SFT suboptimality stems from "loss masking employed in SFT — prompt tokens are typically excluded from the loss calculation during SFT, which can pose a sparse gradient challenge for sparse attention methods like MoBA" (Section 3.2). Full attention in the final layers ensures that gradients from response tokens can propagate backward through the full context, updating key/value representations for distant prompt tokens that MoBA's sparse routing might have bypassed.

---

Large-Scale Language Modeling Evaluation (Section 3.3, Table 2, Figure 7)

Benchmark performance (Table 2). The paper compares Llama-8B-1M-MoBA (trained with MoBA, last 3 of 32 layers as full attention) against Llama-8B-1M-Full (trained with full attention throughout) across 15 benchmarks covering standard NLP, math, code, and long-context tasks. Both models undergo identical continual pre-training from Llama 3.1 8B with context length gradually increased to 1M tokens, followed by identical SFT recipes — only the attention mechanism differs. Key results:

Standard benchmarks (short-context): Performance is essentially indistinguishable. On MMLU (0-shot): 0.4903 (MoBA) vs. 0.4904 (Full) — a difference of 0.0001. On BBH (3-shot): 0.6573 vs. 0.6589. On HellaSWAG (0-shot): 0.8262 vs. 0.8279. On GSM8K (5-shot): 0.7278 vs. 0.7142 — MoBA slightly outperforms here, though this is likely noise. Across these short-context benchmarks, MoBA shows no systematic degradation, confirming that the sparse attention mechanism does not harm the model's fundamental language understanding and reasoning capabilities.

Code and math benchmarks: MoBA performs competitively. On HumanEval (0-shot, pass@1): 0.6951 (MoBA) vs. 0.7012 (Full) — a 0.0061 gap. On MBPP Sanitized (0-shot): 0.6926 vs. 0.6615 — MoBA actually outperforms by 0.0311, a notable margin. On Competition Math (0-shot): 0.4254 vs. 0.4324 — a 0.007 gap. These results suggest that code generation, which often benefits from attending to distant function definitions or imports, does not suffer materially from MoBA's sparsity.

Long-context benchmarks: This is the critical test. On LongBench at 32K (0-shot): 0.4828 (MoBA) vs. 0.4821 (Full) — essentially tied. On RULER at 128K (0-shot): 0.7818 vs. 0.7849 — a difference of 0.0031, which the paper describes as "nearly matches." On Loogle (0-shot): 0.4209 vs. 0.4016 — MoBA outperforms by 0.0193, which is noteworthy given that Loogle tests long-document question answering.

Needle-in-a-Haystack (Figure 7). At up to 1M context length, Llama-8B-1M-MoBA achieves near-perfect retrieval across all context lengths and needle positions, with the heatmap in Figure 7 showing deep green (scores approaching 100) throughout. The paper states the model "demonstrates satisfactory performance even with an extended context length of 1 million tokens" (Section 3.3). This is a critical validation: Needle-in-a-Haystack tests whether the model can attend to a specific piece of information regardless of where it appears in a long document — exactly the capability that sparse attention might compromise if the routing mechanism fails to select the block containing the needle.

Why these results matter: The benchmark suite addresses the central concern about sparse attention — that while it may match full attention on aggregate metrics, it might fail on tasks requiring precise retrieval from arbitrary positions in long contexts. The RULER result (0.7818 vs. 0.7849) is particularly important because RULER at 128K tests multi-hop reasoning over long contexts, and MoBA operates at 62.5% sparsity for these sequences. The Needle-in-a-Haystack result at 1M tokens with 95% sparsity demonstrates that the routing mechanism reliably selects blocks containing task-critical information even at extreme context lengths.

---

Efficiency and Scalability Benchmarks (Section 3.4, Figure 2)

Speedup at varying sequence lengths (Figure 2a). Comparing the attention layer forward pass time between Llama-8B-1M-MoBA and Llama-8B-1M-Full (FlashAttention) across sequence lengths from 32K to 1M:

• At 32K tokens: MoBA computes attention in approximately 25 ms vs. FlashAttention at approximately 35 ms — a 1.4× speedup.
• At 128K tokens: MoBA approximately 45 ms vs. FlashAttention approximately 130 ms — a 2.9× speedup.
• At 256K tokens: MoBA approximately 65 ms vs. FlashAttention approximately 270 ms — a 4.2× speedup.
• At 512K tokens: MoBA approximately 95 ms vs. FlashAttention approximately 530 ms — a 5.6× speedup.
• At 1M tokens: MoBA approximately 120 ms vs. FlashAttention approximately 800 ms — a 6.5× speedup.

The key observation is that MoBA's computation time grows sub-quadratically with sequence length, while FlashAttention grows quadratically. The paper states that MoBA "demonstrates a sub-quadratic computational complexity" (Section 3.4). At 1M tokens, the speedup reaches 6.5×, confirming that the theoretical sparsity (95.31% for these parameters) translates directly to wall-clock efficiency gains.

Scalability to 10M tokens (Figure 2b). The paper further tests MoBA's length scalability by maintaining a constant sparsity ratio of 95.31% (fixed 64 MoBA blocks with variable block size and top-k=3) while increasing sequence length from 32K to 10M. For this experiment, tensor parallelism is expanded to distribute key/value tensors across distributed query heads (Section 3.4). Results:

• At 1M tokens: MoBA approximately 0.75 seconds vs. FlashAttention approximately 6 seconds (read from Figure 2b log scale) — an 8× speedup (different from the 6.5× in Figure 2a because the sparsity ratio and block configuration differ).
• At 4M tokens: MoBA approximately 3 seconds vs. FlashAttention approximately 30 seconds — a 10× speedup.
• At 7M tokens: MoBA approximately 7 seconds vs. FlashAttention approximately 55 seconds — an 8× speedup.
• At 10M tokens: MoBA approximately 10 seconds vs. FlashAttention approximately 160 seconds — a 16× speedup.

The inset graph in Figure 2b focuses on shorter sequences (32K-512K), showing that "both methods perform comparably at smaller scales" but MoBA's advantage grows with length. The 16× speedup at 10M tokens is the paper's headline scalability result, demonstrating that MoBA can process sequences an order of magnitude longer than practical with full attention.

Prefill vs. generation distinction. The paper notes that "across all evaluation tasks, MoBA is used for prefill only, while we switch to full attention during generation for better performance" (Section 3.3). This means the speedup numbers apply to the prefill phase (processing the input prompt), which is typically the bottleneck for long-context applications (since the prompt can be millions of tokens while generation produces hundreds to thousands). During generation, the model uses full attention because: (1) generation is autoregressive (one token at a time), making the per-step context length manageable even at 1M total tokens; (2) the quality benefit of full attention for each generated token outweighs the efficiency cost given the relatively small generation budget.

---

Ablation Studies and Robustness Checks

• Fine-grained block segmentation (Figure 4): Holding sparsity constant at 75%, increasing the number of blocks from 8 to 128 (with corresponding increases in top-k to maintain the same total attended tokens) improves LM loss by approximately 0.02, with diminishing returns beyond 64 blocks. This validates the transfer of fine-grained expert segmentation from FFN MoE to attention MoBA and demonstrates that routing precision — not just total attention budget — determines MoBA's performance.

• MoBA/Full hybrid training (Figure 5a): Training with MoBA for 90% of tokens then switching to full attention for the final 10% recovers full-attention-level position-wise LM loss with no loss spikes at the transition. This demonstrates parameter compatibility between sparse and full attention modes and provides a practical recipe for efficient long-context pre-training.

• Layer-wise hybrid for SFT (Figures 5b-5c): Adding full-attention layers to the end of the model (1, 3, 5, or 10 layers) monotonically improves SFT loss, with trailing-token loss benefiting more than overall loss. The 1M-context model uses 3 full-attention layers out of 32, representing a practical quality-efficiency tradeoff.

• Position-wise scaling laws (Appendix A.1, Figure 8, Table 3): Breaking down LM loss by 2K-token position ranges reveals that MoBA matches full attention at early positions (exponents within 0.001) but shows a growing gap at later positions (exponent difference of 0.011 at 30K-32K). The gap narrows with model scale, suggesting routing quality improves with capacity. This analysis isolates trailing tokens as the primary locus of sparse attention degradation.

• Prefill-only MoBA with full attention generation (Section 3.3): All evaluation results in Table 2 and Figure 7 use MoBA only during prefill, switching to full attention during generation. This deployment strategy is adopted across all benchmarks, meaning the strong long-context benchmark results (RULER 0.7818, Needle-in-a-Haystack near-perfect) are achieved with sparse attention during input processing and full attention during output generation — not pure sparse attention end-to-end.

• No comparison to other sparse attention methods: A notable absence: the paper does not include quantitative comparisons against sliding window attention, attention sink, Longformer, BigBird, Quest, Minference, or any other sparse attention method. The related work (Section 4) discusses these methods extensively, and Section 2.2 analytically shows that sliding window and attention sink are special cases of MoBA, but no experiments validate that MoBA empirically outperforms these alternatives on long-context benchmarks. This is a significant gap — the claim that learned dynamic routing is superior to fixed patterns is central to the paper's motivation but is supported only by MoBA's ability to match full attention, not by direct comparison to prior sparse methods.

• No ablation on gating mechanism design: The paper uses mean-pooled key vectors with dot-product affinity scoring and hard top-k selection. Alternative designs — max pooling, learned block representations (as in Quest's min+max pooling), threshold-based gating, soft routing, or learned temperature parameters — are not explored experimentally. The fine-grained segmentation study varies block count but not the routing mechanism itself.

• No ablation on k (top-k parameter) independently of block size: In the fine-grained segmentation experiment (Figure 4), k and block count are coupled to maintain constant sparsity. An ablation varying k at fixed block size — which would reveal the tradeoff between attention budget and routing precision — is not presented.

---

Critical Assessment

The experiments demonstrate that MoBA can substitute for full attention without meaningful degradation on a range of standard benchmarks when properly configured — but they do not fully demonstrate that MoBA's specific design choices (MoE-inspired routing, mean-pooled block representations, fixed top-k gating) are responsible for this success rather than the simpler fact that block-sparse attention with a sufficiently large budget works well. The paper's strongest evidence is the comparison to full attention under matched training conditions; its weakest point is the absence of comparisons to other sparse attention methods.

Does MoBA achieve performance comparable to full attention? This claim is well-supported. The scaling law experiments (Figures 3a-3c) show aggregate LM loss within 0.003 nats of full attention across five model scales. The benchmark evaluation (Table 2) shows MoBA within 0.0031 on RULER at 128K and near-perfect on Needle-in-a-Haystack at 1M. The hybrid training experiment (Figure 5a) shows that MoBA-trained models can switch to full attention with no loss spike and achieve full-attention-level performance, confirming that MoBA parameters are not degraded relative to full attention parameters. The consistent finding — across language modeling loss, downstream benchmarks, and long-context retrieval — is that MoBA with 3 full-attention layers (out of 32) matches full attention within statistical noise.

However, the claim that MoBA matches full attention "without the risk of compromising performance" (Abstract) requires qualification: the matching requires specific design choices (3 full-attention layers, prefill-only sparsity with full attention during generation). Pure MoBA (no full-attention layers) shows elevated trailing-token loss (Figure 5a) and suboptimal SFT performance (Figures 5b-5c). The paper's deployment configuration — sparse prefill, full generation, layer-wise hybrid — is a carefully engineered compromise that works well but is not the "pure" MoBA architecture described in Section 2.

Does MoBA achieve meaningful efficiency gains? The speedup numbers are concrete and impressive: 6.5× at 1M tokens (Figure 2a) and 16× at 10M tokens (Figure 2b). These are measured as wall-clock attention layer forward pass times, which is the right metric for practical deployment. The sub-quadratic scaling in Figure 2b is visually convincing — FlashAttention curves upward sharply while MoBA remains nearly linear on the log-log plot.

The caveat is that these speedups are for the attention layer only, not end-to-end model throughput. In a full Transformer, FFN layers (which are identical between MoBA and full attention) contribute substantial computation. The paper acknowledges this implicitly by focusing on attention layer time, but the end-to-end speedup would be lower than 16× — potentially much lower if FFN dominates compute at shorter sequence lengths. The paper does not report end-to-end throughput or latency numbers for full model inference.

Does MoBA successfully apply MoE principles to attention? The conceptual mapping (blocks as experts, top-k gating as routing, current block as shared expert) is clean, and the fine-grained segmentation experiment (Figure 4) provides one piece of evidence that MoE properties transfer to the attention domain. However, the experimental evidence for why the MoE framing is beneficial — as opposed to simply using block-sparse attention with any reasonable selection mechanism — is thin. The paper does not ablate the MoE-specific aspects: no comparison of learned top-k routing against similarity-based selection (e.g., always attending to the k most similar blocks by dot product without the MoE gating framing), no exploration of load balancing (a central concern in MoE that goes unmentioned for MoBA), and no investigation of whether blocks develop specializations analogous to MoE expert specialization.

The comparison to Quest (Section 4) — which the paper notes "can be viewed as MoBA with a smaller block size and a specialized block representation function which combines both min and max pooling" — is made in the related work but not tested experimentally. A direct comparison to Quest on long-context benchmarks would substantially strengthen the paper's claim that the MoE framing provides advantages over alternative dynamic sparse attention approaches.

Does MoBA work because of its routing mechanism, or despite it? The Needle-in-a-Haystack result (Figure 7) is impressive but does not isolate the routing mechanism's contribution. In a needle-in-a-haystack test, the "needle" is a single sentence inserted at a random position; attending to the correct block requires the routing mechanism to recognize that this block is unusually relevant. MoBA succeeds here, but so might simpler baselines like always attending to blocks with high lexical overlap or using a keyword-based retrieval mechanism. Without ablations on the routing mechanism or comparisons to simpler sparse attention patterns, it's difficult to determine whether MoBA's learned routing is doing something sophisticated or simply reliably including blocks that contain task-relevant information due to the high baseline probability of the current block and nearby blocks being selected.

What experiments would strengthen the paper? Several missing experiments would substantially increase confidence in the claims:

1. Direct comparison to sliding window attention and attention sink on long-context benchmarks, given that the paper analytically shows these are special cases of MoBA. Demonstrating that MoBA outperforms these fixed patterns would validate the "less structure" principle.

2. Ablation of the routing mechanism — comparing mean-pooled key dot-product routing against: (a) random block selection, (b) always selecting the most recent blocks (sliding window), (c) always selecting the first + recent blocks (attention sink), (d) selecting blocks by maximum key-query dot product (token-level similarity aggregated to block level). This would quantify how much MoBA's specific routing design contributes.

3. End-to-end throughput and latency measurements for full model inference (not just attention layers) at various sequence lengths and sparsity ratios, to give practitioners realistic deployment numbers.

4. Experiments with pure MoBA (no full-attention layers) on the same benchmarks to quantify exactly how much the layer-wise hybrid strategy contributes to the strong benchmark results.

5. Analysis of block utilization — in MoE, load balancing ensures all experts are used; does MoBA suffer from "block collapse" where certain blocks are never selected, or do all blocks receive attention from some queries? This would illuminate whether routing learns meaningful block specializations.

6. Ablation on the hybrid training ratio (90/10 split) — how sensitive is the result to this ratio? Would 80/20 or 95/5 work equally well? This matters for practitioners adopting the hybrid recipe.

Summary of evidence quality. The paper provides strong evidence that MoBA, in its deployed configuration (layer-wise hybrid, prefill-only sparsity), is a practical and effective attention mechanism for long-context LLMs that matches full attention on standard benchmarks while providing substantial speedups. The evidence for MoBA's specific architectural innovations (MoE-inspired routing, the "less structure" principle) is weaker — the experiments demonstrate that MoBA works but not conclusively why it works or that its specific design choices are responsible for the success rather than the simpler fact that block-sparse attention with a large enough budget and a few full-attention layers performs well. The deployment on Kimi (stated in the Abstract) provides practical validation but is not experimentally documented in the paper — no production metrics, no A/B comparisons, no latency/throughput data from real traffic are reported.

6. Limitations and Trade-offs

6.1 Context Difficulty Estimation Cost Is Unaccounted For

The paper's entire compute-optimal framework depends on knowing each prompt's difficulty before deciding how to allocate the test-time compute budget. The method for estimating difficulty — generating 2048 samples per question and computing either ground-truth pass@1 (oracle) or average PRM final-answer scores (predicted) — is extraordinarily expensive. The authors acknowledge this explicitly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

Consequence: The reported 4× efficiency gains over best-of-N are computed after difficulty is already known, without amortizing the cost of learning it. In a realistic deployment, generating 2048 samples per prompt to estimate difficulty could consume more compute than the test-time budget being optimized (which tops out at 256–512 generations in the experiments). A practitioner implementing this system would discover that the total cost — difficulty estimation plus strategy execution — may exceed the cost of simply running best-of-N at a high budget, undermining the efficiency argument.

Evidence in the paper: This limitation is discussed in Section 3.2 (the paragraph on predicted difficulty) and flagged as future work in Section 8. The figures reporting 4× gains (Figures 4 and 8) do not include the 2048-sample difficulty estimation cost in the compute budget. The paper transparently labels this as an "exploration-exploitation tradeoff — compute spent assessing difficulty versus compute spent solving the problem" (Section 3.2).

Mitigation status: Not addressed. The authors suggest future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8) and note that adaptive difficulty estimation — starting with a few samples, assessing difficulty, then allocating the remaining budget — could subsume the estimation cost into the solving process. Neither approach is developed or evaluated in this paper. Until such methods exist, the 4× figure is an upper bound on achievable efficiency, not a realized deployment gain.

---

6.2 All Results Are on a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*)

Every experiment in the paper — scaling law analyses, search algorithm comparisons, revision model evaluations, FLOPs-matched tradeoff studies — is conducted on the MATH benchmark (500 test questions) using PaLM 2-S* as the base model. The authors state in Section 4 that they "believe this model is representative of the capabilities of many contemporary LLMs," but this claim is unverified.

Consequence: Several findings could be model-specific or benchmark-specific in ways that limit generalizability. The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution — a model with different calibration properties or error patterns might exhibit different difficulty-dependent scaling curves and different optimal strategies. The revision model's ability to learn from incorrect in-context examples depends on the base model's in-context learning capabilities, which vary substantially across model families (GPT-4, Claude, Llama, etc. have different in-context learning strengths). MATH consists exclusively of competition-level math problems requiring symbolic reasoning — it is unclear whether the central finding (that difficulty-conditioned allocation yields 4× efficiency gains) generalizes to code generation, logical reasoning, scientific question answering, or tasks requiring factual knowledge retrieval rather than multi-step deduction.

Evidence in the paper: This limitation is inherent in the experimental design — all figures and tables (Figures 3–9, Tables in Section 7) use MATH and PaLM 2-S*. The paper does not present any results on other benchmarks (GSM8K, HumanEval, MMLU, etc.) or other model families. Section 8 acknowledges this implicitly by calling for future work on "other domains and modalities."

Mitigation status: Not addressed. The authors argue for representativeness by assertion rather than by cross-validation on additional benchmarks or models. A practitioner using a different model (e.g., Llama, GPT, Claude) or targeting a different domain (code, dialogue, retrieval) cannot assume the difficulty-conditioned strategies identified for PaLM 2-S* on MATH will transfer. The paper would be substantially strengthened by even a single additional benchmark (e.g., GSM8K for math reasoning or HumanEval for code) to test cross-domain generalization.

---

6.3 Hard Problems Remain Fundamentally Unsolved — Test-Time Compute Cannot Create Capability

Across all methods studied — PRM-guided search, iterative revisions, and their compute-optimal combinations — the hardest difficulty quintile (bin 5) shows near-zero improvement regardless of test-time compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all search methods at all budgets. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5%, and test-time compute underperforms the ~14× larger model by 37–53% depending on the R ratio.

Consequence: This establishes a hard boundary condition on the paper's central claim that test-time compute can substitute for pretraining. The claim holds only for problems where the base model's pass@1 is non-trivially above zero — i.e., the model can already produce correct solutions at some low rate. For problems where the base model fundamentally lacks the required capability, no amount of search, revision, or adaptive allocation helps. A practitioner deploying this approach would need a reliable way to identify which problems fall into bin 5 (or an analogous "too hard" category) and route them to a different solution entirely (larger model, human intervention, retrieval augmentation) rather than wasting test-time compute.

Evidence in the paper: This limitation is thoroughly documented and the authors are transparent about it. The Section 7 takeaway box explicitly states the boundary conditions. Figure 3 (right) and Figure 7 (right) provide clean visual evidence of the flat bin 5 curves. The FLOPs-matched analysis (Figure 9, and the bar charts in Figure 1) quantifies the degradation on hard problems across all R regimes.

Mitigation status: Acknowledged but not solved. The paper frames this as a fundamental property — "test-time compute can amplify existing capability but cannot create it from nothing" — and suggests that for hard problems, pretraining remains the only viable path (Section 7). This is an honest characterization of a genuine limitation rather than a solvable design flaw, but it means the approach offers no path forward for genuinely novel or out-of-distribution reasoning tasks.

---

6.4 The 14× Larger Model Baseline Is Weakened by Non-Compute-Optimal Training and No Test-Time Compute

The FLOPs-matched comparison in Section 7 pits PaLM 2-S* with compute-optimal test-time strategies against a model with approximately 14× more parameters. However, this larger model is trained by scaling only parameters while holding data fixed — the LLaMA paradigm — rather than using compute-optimal pretraining where both data and parameters scale equally (Hoffmann et al., 2022). The authors acknowledge this:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work." (Section 7)

Additionally, the larger model uses only greedy decoding — no best-of-N, no majority voting, no verifier-guided selection, no revision chains. It receives zero test-time compute optimization.

Consequence: The pretraining baseline is weaker than it could be on two fronts. First, a Chinchilla-optimal model trained with 14× more total FLOPs (scaling both parameters and data) would likely outperform a parameter-only-scaled model, potentially reducing or reversing the reported advantages of test-time compute. Second, giving the larger model even a modest test-time compute budget (e.g., best-of-8 or best-of-16) would create a much stronger baseline — the comparison is effectively "small model + smart inference" vs. "large model + no inference optimization," which stacks the deck in favor of test-time compute. The reported +27.8% advantage on easy questions at R ≪ 1 may partially reflect this asymmetry rather than a fundamental property of test-time vs. pretraining compute.

Evidence in the paper: The experimental setup is described in Section 7. The 14× parameter scaling and the greedy decoding baseline are both explicitly stated. The paper acknowledges the non-Chinchilla-optimal training as a caveat. The R-dependent analysis (three values: 0.16, 0.79, 22) partially addresses the concern by showing that the advantage shifts with the pretraining-to-inference ratio, but does not test alternative pretraining scaling strategies.

Mitigation status: Partially addressed by transparency. The paper is upfront about the design choice and frames it as representative of common practice (LLaMA-style scaling), which is a reasonable defense — many production models are trained this way. However, the absence of even a modest test-time compute budget for the larger model (e.g., best-of-8 majority voting) is not discussed or justified, and represents a genuine weakness in the experimental design that could meaningfully affect the headline comparison.

---

6.5 No Accounting for Latency or Wall-Clock Time — Only Total FLOPs / Generation Count

The paper measures test-time compute in "generations" (number of complete solutions sampled), which is a reasonable proxy for total FLOPs but ignores an important practical dimension: latency. Sequential revision strategies are inherently serial — each revision depends on the output of the previous one — while parallel best-of-N sampling can be executed simultaneously given sufficient hardware.

Consequence: A strategy that allocates 128 generations as a single chain of 128 sequential revisions takes roughly 128× longer wall-clock time than one that runs 128 parallel samples simultaneously, even though both consume the same total FLOPs. The compute-optimal policy favors sequential revisions on easy problems (Figure 7, right: bin 1 shows essentially flat performance regardless of ratio, and the optimal ratio for bins 1–2 skews toward sequential), which means the policy is selecting strategies that maximize FLOPs-efficiency at the expense of latency. For latency-sensitive applications — interactive assistants, real-time decision-making, customer-facing chatbots — this tradeoff may be unacceptable regardless of the 4× FLOPs efficiency gain.

Evidence in the paper: The paper does not discuss latency. All budget measurements are in "generations" (Section 3.4), and the compute-optimal policy (Sections 3.2 and 5.3) selects strategies based solely on accuracy at a given generation budget. The sequential-vs-parallel tradeoff is studied as a FLOPs allocation question (Figure 7), not a latency question.

Mitigation status: Not addressed. The paper does not mention wall-clock time, throughput, or latency as evaluation axes. This is a significant gap for practitioners, who often care more about response time than total FLOPs. A latency-aware allocation policy might make different choices — for example, preferring parallel sampling over sequential revisions even when sequential is slightly more FLOPs-efficient, because the parallel strategy completes in 1/N the wall-clock time. This limitation is partially structural (latency depends on hardware parallelism, batch size, and deployment specifics that are hard to standardize), but the complete absence of discussion is a weakness given the practical framing of the paper.

---

6.6 Revisions and PRM Search Are Never Combined — the Two Axes Remain Independent

The paper studies two complementary mechanisms for test-time compute: PRM-guided search (modifying the verifier, Section 5) and iterative revisions (modifying the proposal distribution, Section 6). These are presented as independent axes in the unifying framework of Section 2, and the paper demonstrates they have complementary difficulty-dependent strengths — revisions excel on easy problems, search excels on medium problems. However, the paper never combines them: there are no experiments where beam search is applied to revision model outputs, or where the PRM guides which revisions to pursue, or where a revision chain is used as the proposal distribution within a search tree.

Consequence: The reported results represent a lower bound on what a fully integrated system could achieve. If revisions improve the proposal distribution (generating better candidates) and PRM search improves candidate selection (finding the best among generated candidates), then combining them could yield accuracy beyond either method alone. The compute-optimal policy currently selects between these strategies per problem — but a combined approach might outperform either strategy on the same problem, potentially shifting the optimal allocation and improving the efficiency ceiling. The paper explicitly acknowledges this gap:

"we did not experiment with PRM tree-search techniques in combination with revisions" (Section 8)

Evidence in the paper: This is a design choice, not a hidden limitation. The paper's structure — separate sections for search (Section 5) and revisions (Section 6), separate compute-optimal analyses (Figures 4 and 8), and no "combined" experiments — makes the independence clear. The Section 8 future work explicitly calls for this combination.

Mitigation status: Not addressed experimentally. The paper identifies the combination as "a natural next step" (Section 8) but provides no data on what combined performance might look like. For a practitioner, this means the current compute-optimal policy is choosing between two incomplete toolkits — you can have search OR revisions, but not both — and the true potential of adaptive test-time compute allocation is likely higher than the 4× efficiency gains reported here. However, until combined experiments exist, the magnitude of this additional gain is unknown.

7. Implications and Future Directions

How This Work Changes the Landscape

MoBA shifts the conversation around efficient attention from which hand-designed sparsity pattern to impose toward how to let the model learn its own sparse connectivity. This is not a paradigm shift — block-sparse attention, dynamic routing, and MoE have all existed independently — but it is a meaningful reframing that synthesizes these ideas into a single architectural principle: attention sparsity should be treated as a Mixture of Experts routing problem, where context blocks compete for a fixed attention budget and the model learns to allocate that budget through gradient-based optimization.

The paper's most significant conceptual contribution is the demonstration that the MoE design toolkit — top-k gating, shared experts, fine-grained segmentation — transfers to the attention dimension with minimal modification. This opens a design space that the field had not systematically explored: instead of asking "which tokens should each query attend to?" (a retrieval problem), researchers can now ask "how should we organize context and train routing to make block-level selection work?" (a representation learning problem). The fine-grained segmentation experiment (Figure 4) — which shows that block granularity, not total attention budget, is the primary performance lever — crystallizes this shift. It is not enough to attend to enough tokens; the model must be able to make precise inclusion decisions, and that requires coherence within blocks that can only emerge from learned representations.

The paper also resolves a tension that has been implicit in the efficient attention literature but rarely articulated: whether sparse attention mechanisms should be designed as replacements for full attention (commit to sparsity and optimize for it) or as compatible alternatives (design for interoperability with full attention). MoBA makes a strong case for the latter. The finding that MoBA-trained models can switch to full attention mid-training with "no significant loss spikes" (Section 3.2, Figure 5a) and achieve full-attention-level performance after a brief fine-tuning phase challenges the assumption — implicit in work on Longformer, BigBird, and linear attention models — that sparse and full attention inhabit different parameter spaces. MoBA demonstrates that they can share a parameter space if the sparse mechanism uses the same mathematical operation (softmax dot-product attention) applied to a subset of the context. This has immediate practical consequences: it means organizations can train the bulk of their long-context models with sparse attention (saving 4–16× in attention FLOPs) and then "top off" with full attention to close any remaining quality gap, rather than choosing between efficiency and quality from the start.

The trailing-token diagnostic (Appendix A.1, Figure 8, Table 3) constitutes a new methodological standard for evaluating sparse attention mechanisms. Prior work relied on aggregate perplexity or end-to-end benchmark scores, which can mask degradation at the most challenging positions. MoBA's position-wise scaling law analysis reveals that the performance gap between sparse and full attention is almost entirely concentrated at trailing tokens — positions far from the sequence start — and that this gap narrows with model scale. This finding has two implications for the field. First, it establishes trailing-token LM loss (specifically, loss on the final tokens of maximally-long sequences, with the scaling exponent as a summary statistic) as the primary evaluation axis for long-context architectures. A sparse attention method that matches full attention on aggregate perplexity but diverges on trailing tokens is meaningfully worse at long-context processing. Second, it suggests that research on sparse attention should focus specifically on improving routing quality at distant positions, since early-position routing is already essentially solved (Figure 8 shows near-identical scaling at positions 0–8K).

The paper also provides the strongest evidence to date that block-sparse attention with learned routing can fully substitute for full attention in production long-context systems, at least when deployed with the layer-wise hybrid strategy (a few full-attention layers at the end of the model). The benchmark results in Table 2 — MoBA matching full attention on RULER at 128K (0.7818 vs. 0.7849), on LongBench at 32K (0.4828 vs. 0.4821), and achieving near-perfect Needle-in-a-Haystack retrieval at 1M tokens (Figure 7) — are the most comprehensive validation of a learned sparse attention mechanism at million-token scales published to date. The fact that these results come from a deployed production system (Kimi) rather than a research prototype adds credibility: MoBA is not just theoretically promising but practically viable under real workload constraints.

However, several important questions remain unresolved, which temper the landscape impact. The paper does not experimentally compare MoBA against other sparse attention methods — sliding window, attention sink, Quest, Minference, or any static pattern — so the claim that learned routing is superior to fixed patterns is supported analytically (Section 2.2 shows these are special cases of MoBA) but not empirically. The paper does not demonstrate that the MoE framing specifically — as opposed to any reasonable block selection mechanism — is responsible for the strong results; an ablation comparing top-k gating against simpler selection heuristics (always select most recent blocks, select by maximum per-token dot product, etc.) is absent. And the paper does not explore load balancing, expert specialization, or other MoE-specific phenomena in the attention context, leaving open whether the MoE analogy is a useful design guide or merely a post-hoc interpretation.

Follow-Up Research This Work Enables

Head-to-head comparison of learned routing against static sparse patterns on long-context benchmarks. The paper analytically shows that sliding window attention and attention sink are special cases of MoBA (Section 2.2), but never tests whether MoBA's learned routing actually outperforms these fixed patterns on tasks that require long-range attention. A strong follow-up would train Llama-scale models with identical compute budgets using: (a) MoBA with top-k routing, (b) sliding window attention with equivalent sparsity, (c) attention sink with equivalent sparsity, and (d) a hybrid of sliding window + global tokens (Longformer-style). Evaluation on RULER, LongBench, and Needle-in-a-Haystack would quantify how much the "less structure" principle actually buys in practice. The trailing-token scaling law analysis (Appendix A.1) provides the right evaluation framework — the key metric would be the position-wise loss gap between each method and full attention, particularly at the most distant positions. A negative result (MoBA performing no better than sliding window on these benchmarks) would challenge the paper's central motivation and suggest that the benefit of learned routing is limited to the language modeling objective rather than downstream task performance.

Ablation of the routing mechanism to isolate the contribution of the MoE framing. MoBA's routing uses mean-pooled key vectors with dot-product affinity scoring and hard top-k selection. How much does each component matter? A systematic ablation would compare: (a) the current design, (b) max-pooled block representations instead of mean-pooled, (c) learned block representations (a small MLP producing a block embedding from the key vectors), (d) token-level similarity aggregation (compute max dot product between query and any key in the block, rather than mean-pooling first), (e) random block selection (to establish a lower bound), (f) always selecting the most recent k blocks (sliding window upper bound), and (g) soft routing with a learned temperature (allowing gradient flow through the routing decision). The fine-grained segmentation experiment (Figure 4) provides a template — sweep each variant across multiple block granularities at constant sparsity, measuring both LM loss and downstream benchmark performance. This would reveal whether MoBA's success comes from the quality of its block representations, the competitiveness of top-k selection, or simply from attending to enough tokens regardless of how they're chosen.

Load balancing analysis and expert specialization in the attention dimension. Standard MoE requires load balancing losses to prevent collapse where all tokens route to a few experts. MoBA has no such mechanism — queries freely route to any blocks via top-k gating without constraints on block utilization. Does MoBA suffer from "block collapse," where certain blocks receive little or no attention while others are oversubscribed? A diagnostic experiment would track: (a) the distribution of attention across blocks (what fraction of queries attend to each block?), (b) whether block utilization changes with model depth (do early layers and late layers exhibit different routing patterns?), (c) whether blocks develop content-based specializations (do certain blocks consistently receive attention for specific types of queries — e.g., blocks containing entities, blocks containing numerical data, blocks at specific relative positions?), and (d) whether adding a load balancing loss (analogous to MoE auxiliary losses) improves or degrades performance. The paper's deployment on Kimi means such analyses are feasible with production-scale models. A finding that blocks do not specialize — that routing is essentially position-based rather than content-based — would challenge the MoE analogy and suggest MoBA is closer to a learned sliding window with variable stride than a true mixture of context experts.

Combining MoBA with inference-time dynamic sparsity for decode-phase acceleration. The paper uses full attention during generation (Section 3.3) and MoBA only during prefill. This leaves decode-phase efficiency on the table for applications that generate very long outputs (e.g., long chain-of-thought reasoning). A natural extension would apply MoBA during generation as well, potentially with different routing parameters (larger k, smaller blocks, or different gating) optimized for the autoregressive decode phase where each new token is a query attending to a growing context. The challenge is that during generation, the KV cache grows token-by-token, and block boundaries shift — a block that was "historical" for one token becomes "current" for the next. A strong follow-up would design a MoBA variant for the decode phase (perhaps with sliding block windows that dynamically update block assignments), measure the quality-efficiency tradeoff at various generation lengths (100 to 10,000 output tokens), and compare against KV-cache compression methods like H2O, StreamingLLM, and Quest. The metric would be generation throughput (tokens/second) at matched quality on long-form generation benchmarks.

Stress-testing MoBA on tasks requiring precise cross-document attention. The paper evaluates MoBA on standard long-context benchmarks (RULER, LongBench, Needle-in-a-Haystack) and general NLP benchmarks. But MoBA's block-sparse attention means that information separated by more than a few blocks can only interact if the routing mechanism serendipitously selects the relevant blocks — there is no guarantee that two related pieces of information in different parts of a long document will be routed into the same query's attention span. An adversarial stress test would construct synthetic tasks where the model must attend to specific token positions: (a) multi-hop reasoning where the answer depends on information at positions p1, p2, and p3, and the distance between p1 and p3 is varied from 10% to 90% of the context length, (b) coreference resolution where a pronoun near the end of a document refers to an entity introduced near the beginning, with varying numbers of distractors in between, and (c) information integration tasks where the model must combine facts from multiple sections of a long document. Varying the block size and top-k while measuring task accuracy at different context lengths would reveal the practical limits of block-sparse routing — at what information density and separation distance does MoBA's routing fail to capture necessary cross-references? This would produce actionable guidance on block size and top-k selection for specific application requirements.

Difficulty-adaptive sparsity: varying top-k based on query or layer characteristics. MoBA uses a fixed top-k for all queries and all layers. But not all queries need the same attention budget — some tokens (e.g., punctuation, function words, tokens in the middle of a predictable phrase) may need minimal context, while others (e.g., the first token of a new sentence, a token at a syntactic boundary) may need broad context. Similarly, different layers may benefit from different sparsity levels — early layers might process local syntax (needing only nearby blocks) while later layers integrate long-range information (needing more blocks). A follow-up would design an adaptive top-k mechanism where k is predicted per-query or per-layer by a lightweight gating network, trained to minimize a combined loss of language modeling quality plus a sparsity penalty. The experiment would compare fixed-k MoBA against adaptive-k variants on trailing-token LM loss and downstream benchmarks at matched average sparsity, testing whether adaptive allocation provides efficiency gains beyond the current uniform approach. The layer-wise hybrid strategy (Section 3.2) provides preliminary evidence that different layers contribute differently to long-context processing — the fact that adding full-attention layers at the end helps more than adding them at the beginning suggests that later layers are the bottleneck for long-range integration, and giving them more attention budget could be an efficient allocation.

Practical Applications and Downstream Use Cases

Continual pre-training of existing Transformer models for long-context capabilities. The paper demonstrates that MoBA can be applied through continual pre-training from a standard Llama 3.1 8B checkpoint (Section 3.3), with context length gradually extended from the original 8K (or 128K for Llama 3.1) to 1M tokens using position interpolation. The MoBA/Full hybrid training recipe (Section 3.2, Figure 5a) — 90% of training with MoBA at 95% sparsity, 10% with full attention — provides a concrete, cost-effective pipeline. For an organization with an existing Transformer-based LLM and access to long-context training data, this means: (1) adopt MoBA for the attention layers during long-context pre-training, reducing attention FLOPs by 4× to 16× depending on target context length (substantially lowering the GPU-hour cost of long-context adaptation), (2) switch to full attention for the final 10% of training tokens to close any trailing-token quality gap, (3) optionally keep the last 3 layers as full attention during deployment (as the paper does for Kimi) to balance quality and efficiency, and (4) use MoBA during prefill and full attention during generation. The paper's matching benchmark scores (Table 2) and Needle-in-a-Haystack performance (Figure 7) provide confidence that this pipeline does not sacrifice long-context capability. The primary implementation requirement is integrating MoBA's block-sparse attention kernels — which build on FlashAttention and MoE primitives (Algorithm 1) — into the existing training framework.

Cost-efficient processing of very long documents in batch inference pipelines. For applications that process large volumes of long documents — legal document review (analyzing thousands of contracts), scientific literature mining (extracting information from full-text papers), financial analysis (processing annual reports and earnings call transcripts), or codebase analysis (reasoning over entire repositories) — the prefill phase dominates inference cost because each document's full text must be encoded before any generation begins. MoBA's 6.5× speedup at 1M tokens and 16× speedup at 10M tokens (Figure 2) translates directly to reduced latency and cost for this prefill phase. The paper's finding that MoBA can process 1M tokens in approximately 120 ms per attention layer vs. 800 ms for FlashAttention (Figure 2a) means that a model with 32 layers and 1M-token inputs would spend roughly 3.8 seconds in attention with MoBA vs. 25.6 seconds with full attention — a 6.7× reduction in the dominant cost component. For a batch inference service processing millions of documents, this difference could determine economic feasibility. The layer-wise hybrid strategy provides a tunable quality-efficiency knob: deployments processing routine documents (where occasional attention misses are acceptable) can use pure MoBA or few full-attention layers; deployments processing high-stakes documents (legal, medical) can use more full-attention layers at higher cost.

On-device or edge deployment with context-length scaling. The ability to process long contexts efficiently opens the possibility of running long-context models on hardware-constrained devices where full attention at 1M tokens would be impossible due to memory or compute limits. While the paper's experiments use datacenter-scale models (8B parameters, tensor parallelism for 10M tokens), the scaling law results (Section 3.1, Figures 3a-3b) show that MoBA's relative advantage over full attention is consistent across model sizes from 568M to 2.1B parameters — the sparsity benefits apply regardless of model scale. For a small model (e.g., 1–3B parameters) deployed on a consumer GPU or mobile NPU, MoBA could enable processing of full-length documents or conversation histories that would otherwise exceed memory budgets. The key practical question (not answered by the paper) is the end-to-end throughput and memory consumption at these smaller scales — the paper only reports attention-layer time, not full model inference metrics. A practitioner would need to benchmark MoBA with their specific model size, hardware, and target context length, using the open-source codebase.

Training data generation for self-improvement with long-context reasoning. The emergence of long chain-of-thought reasoning (Kimi k1.5, DeepSeek-R1, OpenAI o1/o3) means that training data for the next generation of models will increasingly consist of very long reasoning traces — potentially tens or hundreds of thousands of tokens per example. Generating this data (via rejection sampling, best-of-N, or RL-based exploration) requires running the model on long inputs many times. MoBA's training-phase efficiency (the MoBA/Full hybrid recipe) and inference-phase efficiency (prefill speedup) both apply here: the model generating training data can use MoBA to process long prompts faster, and the model being trained on that data can use MoBA to reduce the cost of long-context pre-training. The paper's finding that MoBA matches full attention on math and coding benchmarks (Table 2: GSM8K 0.7278 vs. 0.7142, HumanEval 0.6951 vs. 0.7012) suggests that reasoning capability is preserved under sparse attention, making MoBA suitable for generating and learning from long reasoning traces. A concrete pipeline: use a MoBA-equipped model to generate long CoT solutions for a training dataset, filter for correctness, and then fine-tune (with the layer-wise hybrid SFT strategy from Section 3.2) on the resulting long-context examples — all while paying 4–16× less attention compute than full attention would require.