InfLLM-V2: Dense-Sparse Switchable Attention for Seamless Short-to-Long Adaptation

ArXiv: 2509.24663

🎯 Pitch

Existing trainable sparse attention methods like NSA break the standard pretrain-on-short, finetune-on-long workflow by introducing new parameters that misalign with the pretrained dense model, causing training instability. InfLLM-V2 solves this by reusing dense attention weights within a parameter-free sparse architecture, achieving a 4× speedup over dense attention while retaining over 98% of long-context performance.

---

1. Executive Summary

This paper introduces dense-sparse switchable attention framework (InfLLM-V2), a trainable sparse attention mechanism that seamlessly adapts models pretrained on short sequences for long-context processing without introducing additional parameters. InfLLM-V2 reuses existing dense attention parameters through parameter-free architecture modifications — shared key-value projections and a unified sparse attention that merges selected attention and sliding attention — while maintaining the ability to dynamically switch back to dense attention for short-sequence efficiency. Evaluated on an 8B-parameter GQA model with long-context understanding benchmarks (RULER, LongBench) and chain-of-thought reasoning tasks (MATH-500, AIME, LiveCodeBench), InfLLM-V2 achieves 4× speedup over dense FlashAttention while retaining 98.1% and 99.7% of the original performance respectively, establishing that trainable sparsity can match full-attention performance in the pretrain-on-short, finetune-on-long paradigm only when the sparse architecture preserves parameter alignment with the pretrained dense model, as NSA's multi-parameter design disrupts training stability and degrades to 59.92% on RULER versus InfLLM-V2's 82.62%.

2. Context and Motivation

The Core Problem: Standard Attention Does Not Scale to Long Sequences

The fundamental challenge this paper tackles is a structural inefficiency in the Transformer architecture that becomes crippling as sequences grow long. The self-attention mechanism — the heart of every modern large language model — computes attention scores between every pair of tokens in a sequence. For a sequence of length $n$, this requires $O(n^2)$ computation and $O(n^2)$ memory. When $n$ reaches tens or hundreds of thousands of tokens, both the compute budget and the GPU memory needed to store the full attention matrix become prohibitive.

This is not a theoretical concern. The paper identifies several real-world scenarios where long-sequence processing is the decisive capability bottleneck (Section 1):

• Long-input scenarios: deep research tasks that require synthesizing information across entire books or codebases (Zheng et al., 2025; Xu & Peng, 2025), chatbots that maintain coherent personality and factuality over multi-turn conversations spanning weeks, and software engineering agents that must reason across entire repository histories to fix bugs (Jimenez et al., 2023; Yang et al., 2025).
• Long-output scenarios: chain-of-thought reasoning models that produce thousands of tokens of intermediate deliberation (OpenAI et al., 2024; DeepSeek et al., 2025) and autonomous agents that generate and execute multi-step plans (Wang et al., 2024).

In all these settings, the model's ability to attend across the full context directly determines its performance. Yet the $O(n^2)$ scaling of standard attention means that doubling the context length quadruples the attention cost. For a model deployed at scale — serving millions of queries daily — this quadratic growth makes long-context processing economically unviable without architectural intervention.

Why This Is Urgent Now: The Mismatch Between Capability and Efficiency

The urgency of this problem has intensified with two converging trends. First, frontier models are being trained and deployed with ever-larger context windows — 32K, 128K, even 1M tokens — because downstream applications demand them. Second, the pretrain-on-short, finetune-on-long workflow has become the dominant paradigm for building long-context LLMs. The reasoning is practical: pretraining on long sequences from scratch is astronomically expensive (cost scaling linearly with sequence length), so practitioners pretrain on manageable short sequences (often 4K tokens) and then finetune on longer ones to extend the effective context window.

This paradigm creates a latent architectural tension that the paper brings into focus. When finetuning a dense-attention model on longer sequences, the attention mechanism itself remains unchanged — it still computes full $O(n^2)$ attention. The model can learn to attend over longer contexts, but it does so at ever-increasing computational cost. Efforts to replace dense attention with sparse attention at this finetuning stage have run into a fundamental problem: the sparse architecture must not disrupt what the model already learned during short-sequence pretraining. As the paper demonstrates through its NSA baseline experiments (Section 4, Table 1, Figure 5), architectures that introduce new parameters or restructure the attention computation create a training discontinuity — the loss spikes when switching architectures, erasing pretrained knowledge and requiring the model to re-learn basic capabilities from scratch during long-context finetuning.

Prior Approaches: Training-Free Sparsity and Its Fundamental Limit

The first wave of solutions, which the paper categorizes as training-free sparse attention (Section 2.1), attempts to accelerate inference without modifying the model's weights. These methods exploit the empirical observation that attention patterns in trained Transformers are naturally sparse — most tokens attend strongly to only a small subset of the context.

Predefined sparse patterns rely on hand-crafted heuristics. Sliding window attention (Beltagy et al., 2020) restricts each token to attending only to nearby neighbors, exploiting the local structure of language. Other approaches designate special "global" tokens — initial tokens, segment separators — that all tokens must attend to (Xiao et al., 2024b; Chen et al., 2024; Child et al., 2019), based on the observation that these tokens serve as attention sinks that absorb excess attention probability. While computationally predictable, these hand-designed patterns are inherently rigid: they cannot adapt to the semantic content of the query or the context.

Dynamic sparse patterns address this rigidity by computing relevance scores between the query and candidate context blocks, then selecting only the most relevant blocks for full attention computation. InfLLM (Xiao et al., 2024a), the direct predecessor of this paper's method, partitions the context into contiguous blocks and uses attention score approximations to select which blocks each query token should attend to. MInference (Jiang et al., 2024) builds on similar principles with optimized block-sparse kernels. Other methods (Tang et al., 2024; Zhang et al., 2025b; Lai et al., 2025) refine the block selection criteria or the granularity of sparsity patterns.

The paper also notes that research on attention sparsity has spawned related work on KV cache eviction and compression (Section 2.1) — methods like H2O (Zhang et al., 2023) and SnapKV (Li et al., 2024) that discard or compress key-value pairs with low attention probabilities to reduce memory consumption. While these techniques are complementary to sparse attention, they operate at the memory-management level rather than fundamentally altering the attention computation.

The fundamental trade-off of training-free methods: Because training-free methods do not modify model weights, they inherit the attention patterns learned during dense pretraining. The model was trained to attend broadly across the full context; constraining it to attend only to a subset risks discarding tokens that the model would have used. To avoid catastrophic performance degradation, training-free methods must be conservative — they retain a relatively large fraction of tokens, which limits the achievable sparsity and therefore the speedup. As the paper states (Section 2.1):

"Training-free methods, while focusing on improving the inference efficiency of dense attention models, are often constrained by insufficient sparsity levels in order to avoid severe performance degradation and finally suffer from limited acceleration benefits."

This trade-off is visible in the paper's experimental results (Table 1): training-free InfLLM applied to a full-attention model achieves only 27.94% on RULER at 32K, while MInference — a more sophisticated dynamic method — achieves 73.22% but still substantially trails the full-attention baseline's 84.26%. The gap between training-free sparsity and full attention represents the price of not training the model to operate under the sparsity constraint.

Trainable Sparse Attention: The Promise and the Perils

The second wave of approaches, trainable sparse attention (Section 2.2), aims to close this gap by incorporating sparsity into the training process itself. If the model learns to produce good attention patterns under a sparsity constraint, it should, in principle, achieve strong performance at much higher sparsity levels than training-free methods permit.

The paper identifies three recent trainable sparse attention methods, each with distinct limitations:

SeerAttention (Gao et al., 2024) uses self-distillation to train a router that selects relevant context blocks for each query block. It treats query tokens in blocks rather than individually, which means it can only accelerate the prefilling phase (where all queries are processed in parallel) but not the autoregressive decoding phase (where tokens are generated one at a time). For long-output tasks — chain-of-thought reasoning, agent trajectories — decoding efficiency is equally critical.

MoBA (Lu et al., 2025) applies a mixture-of-blocks approach during the short-to-long adaptation phase, training routers between query blocks and KV blocks. It shares the same prefilling-only limitation as SeerAttention.

NSA (Yuan et al., 2025) is the most architecturally ambitious prior work and the paper's primary point of comparison. NSA designs three distinct attention components — Compressed Attention, Selected Attention, and Sliding Attention — each with its own set of key-value projection parameters, and combines their outputs through a learned gating mechanism. By operating at token-level granularity rather than block-level, NSA can accelerate both prefilling and decoding. Its CUDA kernel implementation achieves practical wall-clock speedups.

Where NSA falls short — the architectural mismatch problem: The paper's central critique of NSA is not that it fails to produce speedups, but that its architecture is fundamentally incompatible with the pretrain-on-short, finetune-on-long workflow that dominates real-world LLM development. The technical argument, laid out in Section 3.1 and Figure 1, has several layers:

1. Parameter proliferation: NSA introduces three separate sets of key-value projection matrices ($W^{cmp}_K, W^{cmp}_V$ for Compressed Attention, $W^{slc}_K, W^{slc}_V$ for Selected Attention, $W^{win}_K, W^{win}_V$ for Sliding Attention) plus an MLP for compressing KV tensors and a gating module with learned parameters. When adapting a dense pretrained model, these parameters must be initialized — typically by replicating the original dense KV projections — and then trained from a state that has never seen gradient signals for this three-way decomposition.

2. Multi-output attention: Standard dense attention produces a single attention output per head. NSA produces three distinct outputs ($O_{cmp}, O_{slc}, O_{win}$) that are linearly combined through learned gate values. This means the model must learn, during finetuning, how to coordinate three parallel attention mechanisms that were never present during pretraining. The gating mechanism itself introduces additional parameters that start from random initialization.

3. Training instability: The consequence of these two factors is visible in Figure 5. When the pretrained dense model is converted to NSA and finetuned on long sequences, the training loss spikes dramatically at the transition point — jumping from approximately 1.2 (where full attention and InfLLM-V2 continue smoothly) to above 1.4. This spike indicates that the architectural change has partially destroyed the representations learned during pretraining. The model must expend training compute to recover from this disruption before it can make progress on the long-context task.

4. Short-sequence overhead: Even when processing short sequences where sparse attention is unnecessary, NSA must still execute all three attention modules and the gating mechanism because the architecture provides no fallback to dense attention. This imposes a constant computational tax on every forward pass regardless of sequence length, making NSA inefficient for the mixed-length workloads that characterize production deployments.

The quantitative evidence for NSA's inadequacy in the pretrain-on-short, finetune-on-long setting is stark (Table 1): NSA achieves only 59.92% on RULER compared to full attention's 84.26%, despite being trained with the same long-context data. Its LongPPL of 4.24 versus full attention's 2.06 confirms that it has not adequately learned long-range dependencies. On general short-sequence tasks (Table 4), NSA degrades to 60.63% average across seven benchmarks, compared to the pretrained model's 67.73% and full-attention finetuning's 67.41%. These results indicate that NSA's architectural disruption is not merely a training inefficiency — it causes permanent capability regression.

The Block Selection Bottleneck: A Hidden Efficiency Problem

Beyond the architectural mismatch, the paper identifies a second, more subtle problem that affects any block-sparse attention method: the block selection step itself introduces significant computational overhead (Section 3.4). To decide which blocks each query token should attend to, the method must compute approximate attention scores between the query and some compressed representation of the context. This computation writes intermediate attention score matrices to GPU high-bandwidth memory (HBM) — and the volume of this I/O scales as $h_q n^2 / s_{C1}$, where $h_q$ is the number of query heads, $n$ is sequence length, and $s_{C1}$ is the compression stride. Since $s_{C1} \ll n$, this approaches the same memory bandwidth bottleneck that sparse attention was designed to avoid.

The paper's profiling reveals that this block selection overhead can dominate the total attention time, eating into the speedup gains from sparsifying the main attention computation. Without addressing this, even a perfectly aligned sparse architecture would fail to deliver practical speed improvements at long sequence lengths.

How InfLLM-V2 Positions Itself

The paper situates InfLLM-V2 as a direct response to both the architectural mismatch problem (which causes training instability and performance degradation) and the block selection bottleneck (which limits practical speedup). Its design philosophy, articulated in Section 3.2, is guided by three principles:

Parameter-free adaptation: InfLLM-V2 introduces zero additional parameters. It reuses the existing dense attention's key-value projections ($W_K, W_V$) for sparse attention, eliminating the need for separate projection matrices for different attention modes. This means that when switching from dense pretraining to sparse finetuning, no parameters need to be randomly initialized — every weight starts from its pretrained, converged value. The architectural transition is therefore a change in computation pattern only (which blocks are attended to), not a change in the model's representational capacity.

Unified attention output: Instead of producing three separate attention outputs that must be gated together (as in NSA), InfLLM-V2 produces a single attention output through a unified Sparse Attention module. The Selected Attention and Sliding Attention patterns are fused by expanding the local block window in Selected Attention to fully cover the sliding window region (Figure 3). Compressed Attention is used only for block selection — its attention scores guide which blocks to select — but its output is discarded; it does not contribute to the final attention result. This single-output design mirrors dense attention, ensuring that the gradient flow and representational structure learned during pretraining remain intact.

Dense-sparse switchability: Because the parameters are shared and the computation is structurally aligned, the model can freely switch between dense attention (for short sequences) and sparse attention (for long sequences) at inference time. This is enabled by a simple sequence-length check: if $n$ is below a threshold, run the standard dense FlashAttention kernel; if above, run the sparse attention kernel. There is no learned gating, no additional components to execute — just a routing decision that costs nothing.

The paper's contribution is thus not a fundamentally new sparse attention pattern (the block-sparse structure builds on InfLLM and NSA) but rather an architectural philosophy: that trainable sparse attention for the pretrain-on-short, finetune-on-long paradigm must maintain strict parameter identity with the dense pretrained model. Any deviation — additional KV projections, extra attention outputs, learned gating mechanisms — creates a distributional shift that the finetuning process cannot fully recover from within a practical compute budget.

This position is supported empirically through the side-by-side comparison with NSA in the experimental section (Tables 1-4, Figure 5), which demonstrates that InfLLM-V2 achieves near-parity with full attention across all benchmarks (RULER: 82.62% vs. 84.26%; LongBench: 42.54 vs. 42.30; long reasoning: 42.66% vs. 42.79%) while NSA substantially underperforms. The training loss curves (Figure 5) provide the mechanistic explanation: InfLLM-V2 continues smoothly from the pretrained loss, while NSA experiences a disruptive spike at the architectural transition.

The Efficiency Sweet Spot

The paper also addresses a pragmatic concern that prior trainable sparse attention work has underemphasized: the method must be fast not just for long sequences but for all sequences. In production deployments, models process a mixture of short and long inputs — a chatbot might handle a one-sentence query followed by a 100-page document analysis. If the sparse attention mechanism imposes overhead on short sequences (because it must execute compression, block selection, and gating regardless of whether sparsity is beneficial), the average-case speedup may be much smaller than the long-sequence speedup touted in benchmarks.

InfLLM-V2's switchable design solves this by falling back to dense FlashAttention for short sequences, where the $O(n^2)$ cost is manageable and the constant-factor overhead of block selection would outweigh any theoretical savings. The efficiency results (Figure 6) demonstrate that this design delivers speedups at all sequence lengths — the overhead of block selection has been minimized through the fused kernel implementation, and the method transitions gracefully from dense to sparse as sequence length grows.

3. Technical Approach

3.1 Reader Orientation

InfLLM-V2 is a trainable sparse attention mechanism that can be swapped into any existing Transformer-based language model at finetuning time, requiring zero new parameters, and that can dynamically switch between dense attention (for short sequences) and sparse attention (for long sequences) at inference time. The paper is primarily an architectural design and empirical analysis paper whose core idea is that trainable sparse attention for the pretrain-on-short, finetune-on-long workflow must maintain strict parameter identity with the pretrained dense model — any architectural deviation creates a training discontinuity that practical finetuning budgets cannot fully recover from.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components, all operating within a standard Transformer decoder with grouped-query attention (GQA):

1. Shared Key-Value Projections — A single set of projection matrices $W_K$ and $W_V$ that serve both dense and sparse attention paths. These are initialized from the pretrained dense model and fine-tuned on long sequences without duplication or expansion.

2. Block Compression Module — A parameter-free pooling pipeline that compresses the full key sequence into a coarse-grained representation used solely for selecting which blocks to attend to. It uses a three-stage cascade (mean pooling → GQA head-group summation → max pooling) to preserve fine-grained information in the selection scores.

3. Block Selection — The compressed key representation is used to compute approximate attention scores, from which the top-k most relevant blocks are selected for each query token, along with fixed initial and local blocks. This produces the sparse attention mask.

4. Unified Sparse Attention — A single attention module that computes exact attention only over the selected blocks (initial + local + top-k), producing a single attention output per head. It fuses the Selected Attention and Sliding Attention patterns from prior work by expanding the local block window.

5. Dense-Sparse Switch — A sequence-length-gated router: if $n$ is below a threshold, run standard dense FlashAttention using the shared KV projections; if above, run the sparse attention pipeline. This requires no learned parameters and no additional computation for the unchosen path.

Information flows as follows: a hidden state tensor $X \in \mathbb{R}^{n \times d}$ enters an attention layer → $X$ is projected to $Q, K, V$ using the shared projections → if the sequence is short, the standard dense FlashAttention kernel computes full attention and returns the output → if the sequence is long, the block compression module pools $K$ into a coarse representation, block selection uses this representation plus $Q$ to determine which blocks are visible, and the sparse attention kernel computes attention only over the selected blocks → the output proceeds to the output projection $W_O$ as in standard attention.

3.3 Roadmap for the Deep Dive

• First, the formal GQA background and NSA architecture (Section 3.1), which sets up the precise problem InfLLM-V2 solves — without understanding what NSA does and why it creates an architectural mismatch, the design choices in InfLLM-V2 are opaque.
• Second, the shared KV projection and aligned computation (Section 3.2), which are the two design principles that eliminate the architectural mismatch.
• Third, the block compression and selection mechanism (Section 3.3), which determines which blocks receive full attention and how the model preserves granularity during coarse-grained selection.
• Fourth, the efficient CUDA kernel implementation (Section 3.4), which addresses the block-selection bottleneck through kernel fusion and an LSE approximation technique, and which is essential for the practical speedups the paper reports.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily an architectural design and systems paper whose core idea is that sparse attention can match dense attention performance in the pretrain-on-short, finetune-on-long paradigm only when the sparse architecture preserves parameter identity with the dense pretrained model — and that the block selection step itself must be carefully optimized to deliver on the theoretical speedups promised by sparsity.

---

Grouped-Query Attention (GQA) Background

Grouped-query attention (Ainslie et al., 2023) is a variant of multi-head attention that reduces the memory footprint of the key-value cache during autoregressive decoding. In standard multi-head attention, every query head has its own dedicated key and value head, requiring $h_q$ sets of KV caches. GQA instead uses $h_{kv} < h_q$ KV heads, with multiple query heads sharing each KV head.

Formally, given an input sequence of hidden states $X \in \mathbb{R}^{n \times d}$ where $n$ is the sequence length and $d$ is the model dimension, GQA computes:

$Q = XW_Q, \quad K = XW_K, \quad V = XW_V$

where $W_Q \in \mathbb{R}^{d \times (h_q d_h)}$, $W_K, W_V \in \mathbb{R}^{d \times (h_{kv} d_h)}$, and $d_h$ is the per-head dimension.

These tensors are reshaped into $h_q$ query heads $\{Q_i\}_{i=1}^{h_q}$ and $h_{kv}$ KV heads $\{K_j, V_j\}_{j=1}^{h_{kv}}$, each with shape $n \times d_h$. The query heads are partitioned into groups of size $G = h_q / h_{kv}$. For the $i$-th query head, the attention is computed against its corresponding KV head with index $j = \lfloor (i - 1) / G \rfloor + 1$:

$S_i = \text{Softmax}\left(\frac{Q_i K_j^\top}{\sqrt{d_h}}\right), \quad O_i = S_i V_j$

The final output concatenates all attention outputs and projects through $W_O \in \mathbb{R}^{(h_q d_h) \times d}$:

$\text{Attention}(X) = \text{Concat}(O_1, \ldots, O_{h_q}) W_O$

What this computes: For each query token, a scaled dot-product attention score against all key tokens in the same KV head group, normalized by softmax to produce a probability distribution, then used to compute a weighted sum of value vectors. The computation cost is $O(n^2 \cdot d_h \cdot h_q)$.

Why GQA matters for InfLLM-V2: The block-sparse attention pattern in InfLLM-V2 requires that all query heads within a group share the same sparse mask (the same set of selected KV blocks). GQA naturally provides this grouping structure — the paper explicitly sets $G = 16$ (Section 3.4), following NSA's configuration, because this group size is well-suited for block-sparse attention kernel implementations where a warp of threads processes one group of query heads sharing the same block selection pattern.

---

The NSA Architecture and Why It Creates a Mismatch

NSA (Yuan et al., 2025) is InfLLM-V2's primary architectural foil — understanding its design is essential for understanding InfLLM-V2's design choices. NSA extends GQA with three parallel attention mechanisms, each operating on a different view of the context, and combines their outputs through a learned gating mechanism.

The three attention modules:

1. Compressed Attention: Compresses the full key sequence into a shorter representation using a learned MLP-based compression, then computes attention between the uncompressed queries and the compressed keys. This provides a coarse-grained, global view of the entire context at reduced computational cost. Its outputs are denoted $O_{cmp}$.

2. Selected Attention: Leverages the attention scores from Compressed Attention to identify the most important blocks of the original (uncompressed) key sequence. For each query token, it computes exact attention only over these selected blocks, providing fine-grained attention to the most relevant parts of the context. Its outputs are denoted $O_{slc}$.

3. Sliding Attention: Restricts each query token to attend only to a local window of $w$ preceding tokens, capturing local syntactic and semantic patterns that coarse-grained global attention might miss. Its outputs are denoted $O_{win}$.

The gating mechanism: Each module's output is multiplied by a learned scalar gate value and summed:

$\text{Output} = g_{cmp} O_{cmp} + g_{slc} O_{slc} + g_{win} O_{win}$

The gate values $g_{cmp}, g_{slc}, g_{win}$ are computed from the input features $X$ via a learned MLP followed by a sigmoid activation, producing three scalars per token that sum to 1 (due to the sigmoid's independent per-dimension output, though in practice they are independently activated and not constrained to sum to 1 — the paper does not specify exact normalization, only that they are "derived from the input features X via an MLP and a sigmoid activation").

The parameter cost: NSA introduces three separate sets of KV projection matrices — $W_K^{cmp}, W_V^{cmp}$ for Compressed Attention, $W_K^{slc}, W_V^{slc}$ for Selected Attention, and $W_K^{win}, W_V^{win}$ for Sliding Attention — plus the compression MLP and the gating MLP. When converting a pretrained dense GQA model to NSA, these six projection matrices (three for K, three for V) must be initialized — typically by replicating the original $W_K$ and $W_V$ — and then trained. The compression MLP and gating MLP start from random initialization.

Why this creates a training discontinuity: During dense pretraining, the model learned to produce a single attention output per head from a single KV projection. The attention gradients flowed through one computation path, shaping the KV projections to serve a single purpose. In NSA, these same parameters are suddenly asked to serve three distinct purposes simultaneously — compressed attention requires KV representations that are robust to aggressive pooling, selected attention requires KV representations that produce informative scores for block importance ranking, and sliding attention requires KV representations that capture fine-grained local patterns. The model has no prior experience coordinating these three pathways because they did not exist during pretraining.

The quantitative consequence is visible in Figure 5: when the pretrained model is converted to NSA and long-context finetuning begins, the training loss jumps from approximately 1.2 to above 1.4, while InfLLM-V2's loss continues smoothly from the pretrained value. This spike represents the model's confusion as its carefully learned KV representations are repurposed for a fundamentally different computation. Even after 600 steps of long-context finetuning (using 5B tokens), NSA's training loss remains elevated compared to full attention and InfLLM-V2, and its downstream performance on RULER (59.92% vs. 84.26% for full attention) and LongPPL (4.24 vs. 2.06) confirms that the damage is not fully recoverable within a practical finetuning budget.

---

Shared Key-Value Projection: Eliminating Parameter Proliferation

InfLLM-V2's first architectural principle is that sparse and dense attention must use exactly the same KV projection parameters. There is no separate projection for compressed attention, no separate projection for local attention, no separate projection for selected attention. The model uses a single shared set $W_K, W_V$, initialized with the pretrained dense attention parameters, and used for all attention computation — both the coarse-grained block selection scoring and the fine-grained sparse attention over selected blocks.

Formally, regardless of whether the model is in dense mode or sparse mode, the same computation produces the keys and values:

$K = XW_K, \quad V = XW_V$

What this means operationally: When the model switches from pretraining on short sequences (dense attention) to finetuning on long sequences (sparse attention), the only thing that changes is which subsets of the key-value pairs each query attends to. The KV representations themselves are computed identically. The gradients that flow during sparse finetuning update the same $W_K$ and $W_V$ that were optimized during dense pretraining, and because the computation structure is similar (scaled dot-product attention, just over a subset of keys), the gradient signals are compatible with the pretrained parameter state.

Why this eliminates the training discontinuity: The model is not being asked to learn new representational capacities. It already knows how to produce good keys and values for attention scoring and value aggregation — it learned that during dense pretraining. The only new thing it learns during sparse finetuning is which keys are worth attending to when it cannot attend to all of them. This is a strictly easier learning problem than the NSA approach, because:

• The model does not need to learn how to coordinate three separate attention outputs.
• The model does not need to learn gating weights from scratch.
• The model does not need to adapt its KV representations to serve multiple distinct purposes simultaneously.
• The block selection mechanism (described below) is parameter-free, so the model learns to produce KV representations that are informative for block selection through the same attention gradients it would receive in dense mode — there is no separate training signal or auxiliary loss.

The evidence for this elimination is in Figure 5: InfLLM-V2's training loss at the transition point is essentially continuous with the pretrained loss, with no visible spike. This means the model can immediately begin learning the long-context task without first recovering from architectural disruption.

Implementation detail: For the block compression used in block selection (Section 3.3), InfLLM-V2 uses the same $K$ — produced by the shared $W_K$ — as input to the compression pipeline. There is no separate compressed key projection. This means the model must learn to produce key representations that are informative both for exact attention scoring (when used in the sparse attention module) and for coarse-grained importance estimation (when pooled and used for block selection). The paper's results suggest this dual-use is not problematic because pooling is a linear operation in expectation (mean pooling), and the softmax attention scoring is approximately linear in the relevant regime.

NSA initialization comparison: When adapting NSA from a dense pretrained model, the three sets of KV parameters are initialized "by replicating the original KV parameters in dense attention" (Section 4.1). This means all three KV projections start with identical weights. During training, they diverge as each pathway receives different gradient signals. InfLLM-V2 avoids this divergence entirely by having only one set of parameters — there is nothing to replicate and nothing to diverge.

---

Aligned Computation: Fusing Selected and Sliding Attention into a Unified Sparse Module

InfLLM-V2's second architectural principle is that sparse attention must produce a single attention output per head, mirroring the structure of dense attention. This directly addresses NSA's three-output-plus-gating design.

The union operation: In NSA, Selected Attention and Sliding Attention produce two separate sets of attention outputs that the gating mechanism combines. InfLLM-V2 observes that these two attention patterns operate over overlapping sets of tokens. Selected Attention always includes a fixed set of local blocks $I_{local}(i)$ around the query token's block $b_i = \lfloor (i-1)/B \rfloor + 1$ (where $B$ is the block size). Sliding Attention allows the $i$-th token to attend to a window $\{i - w + 1, \ldots, i\}$ of width $w$. Since the local blocks in Selected Attention and the window in Sliding Attention cover adjacent or overlapping regions of the sequence, InfLLM-V2 merges them into a single attention mask.

The merge is accomplished by expanding the number of local blocks in the unified Sparse Attention to strictly cover the region of the Sliding Attention:

$N_{local} \geq \lceil w / B \rceil + 1$

where $N_{local}$ is the number of local blocks on each side of the query's block, $w$ is the sliding window width (in tokens), and $B$ is the block size.

What this means concretely: If the block size $B = 64$ and the sliding window $w = 128$, then $\lceil 128 / 64 \rceil + 1 = 3$ local blocks on each side are sufficient. Figure 3 visualizes this: NSA's Selected Attention with init=1, topk=2, local=1 produces one block of local context on each side; NSA's Sliding Attention with w=6 (tokens, shown as small squares) covers a narrower range; InfLLM-V2's Sparse Attention with init=1, topk=2, local=3 produces three local blocks on each side, which encompasses the sliding window plus additional context. The union mask ensures that every token that would have been attended to by either Selected Attention or Sliding Attention under NSA is attended to by the unified Sparse Attention in InfLLM-V2.

The complete set of attended blocks for a query token with index $i$ is:

$I(i) = I_{init} \cup I_{local}(i) \cup I_{topk}(i)$

where:

• $I_{init} = \{1, 2, \ldots, N_{init}\}$ — a fixed set of initial blocks (typically $N_{init} = 1$, the first block), always attended to by all tokens
• $I_{local}(i) = \{b_i - N_{local} + 1, \ldots, b_i - 1, b_i\}$ — the blocks immediately surrounding and including the query token's own block
• $I_{topk}(i)$ — the top-$k$ blocks selected from the remaining blocks (those not in $I_{init}$ or $I_{local}(i)$) based on the compressed attention scores $S_{cmp}$

If $T_j = \{jB + 1, \ldots, (j+1)B\}$ denotes the set of token indices in the $j$-th block, then the query token attends to the union $\bigcup_{j \in I(i)} T_j$.

Why eliminate Compressed Attention's output: NSA's Compressed Attention module serves two purposes: it produces attention scores used for block selection in Selected Attention, and it produces its own attention output $O_{cmp}$ that contributes to the final gated output. InfLLM-V2 retains only the first purpose — the compressed attention scores $S_{cmp}$ are used solely for ranking blocks and selecting the top-$k$ — and discards the output. This is justified because:

1. The compressed attention operates on pooled representations and therefore provides lower-fidelity attention than the exact attention over selected blocks. For tokens that end up in the selected set, the exact attention is strictly better. For tokens that do not make the cut, the compressed attention would have been the only way to influence them — but by definition, the model judged these tokens to be low-importance, so the loss of this low-fidelity signal is minimal.
2. Eliminating the output removes a degree of freedom that the model would otherwise need to coordinate during training. With only one attention output, the model's training objective is unambiguous: produce good representations for the tokens you chose to attend to. With multiple outputs, the model must learn how to divide representational labor across modules — a coordination problem that the dense pretrained model never faced.

The single-output design's alignment with pretraining: During dense pretraining, the model's output projection $W_O$ was trained to expect a particular statistical distribution of attention outputs. If sparse attention suddenly produces three outputs gated together, the distribution of the combined output may differ from what $W_O$ expects, even if the KV representations are identical. This is a second-order effect — even with shared KV projections, NSA's multi-output gating could cause a subtle distributional shift at the attention output that propagates through the residual stream and disrupts the pretrained representations in subsequent layers. InfLLM-V2's single-output design avoids this entirely: the sparse attention output is structurally identical to a dense attention output (just computed over a subset of keys), so the downstream layers see the same type of signal they were trained on.

The paper's specific configuration: For the experiments, InfLLM-V2 uses $|I| = 96$ total selected blocks (including $|I_{init}| = 1$, $|I_{topk}| = 63$, and $|I_{local}| = 32$) for both training and inference. With block size $B = 64$, the total number of visible tokens per query is $|I| \cdot B = 96 \times 64 = 6,144$ — approximately 6K tokens. For a 32K sequence, this represents a sparsity of approximately $6K / 32K \approx 19\%$, meaning each token attends to only about 19% of the context. The sparsity ratio improves (becomes more aggressive) as sequence length increases.

---

Block Compression: Three-Stage Coarse-to-Fine-Grained Pooling

The block selection mechanism needs a way to estimate which blocks are important for each query token without computing the full $O(n^2)$ attention matrix. InfLLM-V2 uses a three-stage compression pipeline that transforms the full key sequence into a coarse-grained importance score per block.

Why three stages are necessary: A single-stage compression with block size $B$ would simply average-pool each block of $B$ consecutive key vectors into a single representation, then compute attention scores between the query and these pooled block representations. This would lose fine-grained information — a block might contain one highly relevant token and many irrelevant ones, and mean-pooling would dilute the signal. NSA acknowledged this problem and used an MLP for compression (learning a data-dependent pooling), but this introduced additional parameters and required separate gradient signals. InfLLM-V2 solves the granularity problem without parameters by cascading three pooling operations at different resolutions.

Stage 1: Fine-grained mean pooling into coarse key representation

The full key sequence $K \in \mathbb{R}^{n \times d_h}$ (for one KV head) is pooled into a shorter sequence $K_{C1}$ using mean pooling with block size $l_{C1}$ and stride $s_{C1}$:

$K_{C1}^i = \text{Mean}(K_{i \cdot s_{C1} : i \cdot s_{C1} + l_{C1}})$

where $K_{C1}^i \in \mathbb{R}^{d_h}$ is the $i$-th compressed key vector, computed by averaging over a window of $l_{C1}$ consecutive original key vectors starting at position $i \cdot s_{C1}$.

The first-stage attention scores are then computed between the uncompressed queries $Q$ and the compressed keys:

$S_{C1} = \text{Softmax}(Q (K_{C1})^\top)$

where $S_{C1} \in \mathbb{R}^{h_q \times n \times (n/s_{C1})}$ contains one attention score per query head per query token per compressed key position.

What this computes: A coarse attention map where each compressed key position represents $l_{C1}$ original tokens with stride $s_{C1}$. Because $s_{C1} \ll n$, this matrix is substantially smaller than the full $n \times n$ attention matrix, but it still preserves enough resolution to distinguish relevant blocks from irrelevant ones.

Stage 2: Head-group summation for shared block selection

In GQA with group size $G = h_q / h_{kv}$, all query heads within a group attend to the same KV head. InfLLM-V2 forces all heads within a group to share the same block selection pattern, which is achieved by summing the first-stage attention scores across the head group dimension:

$S_{shared} = \sum_{h=1}^{G} S_{C1}^{(h)}$

where $S_{C1}^{(h)}$ is the attention score tensor for the $h$-th query head within the group, and $S_{shared} \in \mathbb{R}^{(h_q/G) \times n \times (n/s_{C1})}$ contains one aggregated importance score per group per query token per compressed key position.

What this computes: An importance score that reflects the consensus of all query heads in the group about which compressed key positions are relevant. This summarization is essential for efficient kernel implementation — it means the block selection pattern is computed once per group rather than once per head, reducing the block selection overhead by a factor of $G$, and it means the sparse attention kernel can process an entire group of query heads with a single mask, improving memory coalescing and reducing control divergence.

Stage 3: Max pooling into block-level scores

The aggregated scores $S_{shared}$ still operate at the compressed key granularity (stride $s_{C1}$). To produce per-block scores for the block size $B$, max pooling is applied with a sliding window:

$S_{cmp}^i = \text{Max}(S_{shared}^{i \cdot s : i \cdot s + l})$

where $s$ is the stride and $l$ is the window length for this final pooling stage. The result $S_{cmp}$ contains one score per block (aligned to the attention block size $B$).

Why max pooling instead of mean pooling: Max pooling preserves the most salient features — if any token within the block's coverage region has high relevance to the query, max pooling ensures that high relevance is reflected in the block score. Mean pooling would dilute such signals, potentially causing the model to miss blocks that contain a single highly relevant token surrounded by irrelevant ones. This is the same intuition that motivates max pooling in computer vision for preserving edge and texture features.

The specific configuration used in experiments: The paper sets $l_{C1} = B/2$, $s_{C1} = B/4$, $l = 5$, and $s = 4$, with $B = 64$. This means $l_{C1} = 32$ and $s_{C1} = 16$. The three stages together achieve the same compression ratio as a single-stage compression with block size $B$, but the intermediate fine-grained stage (with stride 16) preserves information that would be lost by immediately pooling to stride 64. The max pooling stage with $l = 5$ sub-blocks of stride $s = 4$ means each block-level score is computed as the max over 5 overlapping windows, each covering 4 compressed-key positions. With $s_{C1} = 16$, each compressed-key position represents 32 original tokens, so the coverage pattern is dense and overlapping.

Why the compression module is parameter-free: In NSA, the compression is performed by a learned MLP that takes the original keys as input and produces compressed representations. This MLP requires separate training and introduces a discrepancy between the representations used for block selection (learned compressed) and the representations used for the actual attention (original uncompressed). InfLLM-V2 uses simple pooling — mean then max — which is parameter-free and operates directly on the same key representations that will be used in the sparse attention. This means the block selection scores are a direct function of the attention-relevant properties of the keys, with no learned intermediary that could drift or introduce bias.

A subtle design choice: The paper notes that because the Compressed Attention output is eliminated (only its scores are used for block selection), there is no gradient signal that flows through the output of the compressed attention. This is actually intentional — the compressed keys $K_{C1}$ are produced by mean pooling, and the pooling operation itself has no parameters, so there is nothing to train. The gradients that update $W_K$ (and thereby change the key representations) come entirely from the sparse attention module, where exact attention is computed over the selected blocks. This is another manifestation of the paper's design philosophy: the model does not need a separate training signal for block selection because the same KV representations that work well for exact attention will, when pooled, produce informative block selection scores. If a block is selected and turns out to be important (high attention score in the sparse attention module), the gradients will reinforce the key representations that led to it being selected; if a block is selected and turns out to be unimportant, the gradients will diminish those representations.

---

Efficient Implementation: Fused Head Group Summation with LSE Approximation

The paper identifies a critical efficiency bottleneck: computing and storing the first-stage attention scores $S_{C1}$ requires writing $h_q \cdot n^2 / s_{C1}$ scalar values to GPU high-bandwidth memory (HBM). For an 8B model with $h_q = 32$, $n = 32K$, and $s_{C1} = 16$, this is $32 \times (32768)^2 / 16 \approx 2.15$ billion values — roughly 8.6 GB for float32. Reading and writing this much data to HBM dominates the execution time, negating the speedup from sparsifying the main attention computation.

The key insight: Only the reduced scores $S_{shared}$ (after head-group summation) are needed for block selection. If the summation over the head group dimension can be fused into the attention computation kernel — so the intermediate per-head scores never leave the fast on-chip SRAM — the HBM write volume drops by a factor of $G$. With $G = 16$, this reduces the write volume to approximately 0.54 GB, a 16× improvement.

The challenge: online softmax and group summation do not commute

The standard FlashAttention algorithm (Dao, 2024) uses online softmax to compute attention scores without materializing the full unnormalized score matrix. The online softmax maintains running statistics — the row-wise maximum $m$ and the row-wise sum of exponentials $\ell$ — and incrementally updates them as blocks of the key matrix are processed. The final softmax-normalized attention score for position $(i, j)$ is:

$\text{Softmax}(S)_{ij} = \frac{\exp(S_{ij} - m_i)}{\sum_{j'} \exp(S_{ij'} - m_i)} = \frac{\exp(S_{ij} - m_i)}{\ell_i}$

where $m_i = \max_j S_{ij}$ and $\ell_i = \sum_j \exp(S_{ij} - m_i)$. Critically, to compute this, the algorithm needs to know $m_i$ and $\ell_i$ — statistics that depend on all positions $j$ in the row — before it can output the final normalized score for any particular position $j$.

InfLLM-V2 needs to sum the normalized scores across the head group before writing them to HBM. But the summation over the group dimension and the online softmax normalization along the sequence dimension are not commutative:

$\sum_{h=1}^{G} \text{Softmax}(S_{C1}^{(h)}) \neq \text{Softmax}\left(\sum_{h=1}^{G} S_{C1}^{(h)}\right)$

Summing before softmax would produce a different result than summing after softmax. The correct computation requires normalizing each head independently (with its own $m_i^{(h)}$ and $\ell_i^{(h)}$) and then summing the normalized scores. But this means the algorithm cannot fuse the summation into a single-pass online softmax — it would need to compute the normalization statistics first, then recompute to apply them.

Solution: Two-pass approach with log-sum-exp

The paper's solution is a two-pass computation that trades computation for reduced I/O:

1. First pass (coarse-grained): Compute only the log-sum-exp (lse) statistics needed for softmax normalization, using a coarser approximation (described below). The lse for head $h$ at query position $i$ is:

$\text{lse}_i^{(h)} = m_i^{(h)} + \log(\ell_i^{(h)})$

where $m_i^{(h)} = \max_j S_{C1, ij}^{(h)}$ and $\ell_i^{(h)} = \sum_j \exp(S_{C1, ij}^{(h)} - m_i^{(h)})$. These statistics are compact (one scalar per head per query position) and can be stored in SRAM.

2. Second pass (fine-grained): Recompute the attention scores $S_{C1}$, normalize them using the stored lse statistics, sum across the head group, and write only the reduced scores $S_{shared}$ to HBM. The normalized score for head $h$ at positions $(i, j)$ is:

$\text{Softmax}(S_{C1})_{ij}^{(h)} = \exp(S_{C1, ij}^{(h)} - \text{lse}_i^{(h)})$

The summed score across the group is:

$S_{shared, ij} = \sum_{h=1}^{G} \exp(S_{C1, ij}^{(h)} - \text{lse}_i^{(h)})$

What this computes operationally: In the first pass, the kernel processes the coarse key blocks sequentially, accumulates the running maximum and log-sum-exp for each head independently, and stores only the final lse values to SRAM. In the second pass, it reprocesses the fine-grained key blocks, computes the exact attention scores, normalizes them using the first-pass lse, sums across heads within the group, and writes the reduced scores to HBM. The HBM write volume is reduced by a factor of $G = 16$ — only $S_{shared}$ is written, not the per-head scores. The cost is that the key blocks must be loaded from HBM twice, doubling the computational workload of the block selection kernel.

LSE Approximation: reducing the two-pass overhead

To reduce the 2× computational overhead of the two-pass approach, the paper proposes approximating the lse computation using coarser-grained keys. Instead of computing the lse on the fine-grained compressed keys $K_{C1}$ (with block size $l_{C1} = 32$, stride $s_{C1} = 16$), the first pass computes the lse on an even coarser representation $K_{C2}$:

$K_{C2}^i = \text{Mean}(K_{i \cdot s_{C2} : i \cdot s_{C2} + l_{C2}})$

$S_{C2} = \text{Softmax}(Q (K_{C2})^\top)$

with $s_{C2} = 4 s_{C1} = 64$ and $l_{C2} = 4 l_{C1} = 128$.

What this means: The first pass operates on keys that are pooled at 4× coarser resolution. The number of key positions is reduced by a factor of 4, so the attention score computation in the first pass is 4× cheaper. The lse computed from this coarser attention map is used as an approximation for the lse that would have been computed from the fine-grained attention map. The second pass then computes the exact fine-grained attention scores and normalizes using the approximate lse.

The total overhead is reduced from 2× (exact two-pass) to $1 + 1/4 = 1.25\times$ (one pass at 1× cost for the coarse lse, one pass at 1× cost for the fine-grained attention, but the coarse pass is at 4× lower resolution, so it's $1 + 0.25 = 1.25$). The trade-off is that the lse approximation may cause slight inaccuracies in the block selection scores — some blocks might be ranked slightly differently than under exact normalization.

Why this approximation is acceptable: The block selection scores are used only for ranking blocks, not for computing the final attention output. Small normalization errors that shift scores uniformly (e.g., all scores slightly overestimated) do not change the ranking. Only errors that are differentially distributed — causing some blocks to be incorrectly ranked above others — could affect block selection. The paper's experimental results (Table 1, bottom row) confirm this: InfLLM-V2 with LSE Approximation achieves 82.62% on RULER, essentially identical to the 82.09% without the approximation, and in some subtasks it's actually slightly better. This indicates that the approximation noise is below the threshold where it would meaningfully affect block selection quality.

The efficient block selection kernel (Algorithm 1): The paper provides pseudocode (Algorithm 1) for the fused computation of $S_{shared}$. The key steps, assuming $h_{kv} = 1$ for simplicity:

1. Divide the query $Q$ into $T_q = \lceil n / B_q \rceil$ blocks of size $B_q \times G \times d_h$ each.
2. Divide the coarse compressed keys $K_{C1}$ into $T_1 = \lceil n / s_{C1} / B_k \rceil$ blocks.
3. Divide the fine compressed keys $K_{C2}$ into $T_2 = \lceil n / s_{C2} / B_k \rceil$ blocks.
4. For each query block (in parallel):
   • Load the query block from HBM to on-chip SRAM.
   • Initialize the online-softmax statistics (log-sum-exp lse).
   • First pass: Sequentially load $K_{C2}$ blocks from HBM to SRAM, compute the coarse attention scores $S_{C2}$, and update the lse statistics for each head.
   • Second pass: Sequentially load $K_{C1}$ blocks from HBM to SRAM, compute the fine-grained attention scores $S_{C1}$, normalize using the stored lse, sum across the head group dimension within SRAM to produce $S_{shared}^{ij}$, and write only this reduced block to HBM.

The sparse attention kernel (Algorithm 2): The paper also provides pseudocode for the sparse attention computation itself, which closely follows FlashAttention-2 (Dao, 2024) with three modifications:

1. The query blocks contain a group of attention heads (size $G$) for a single token, rather than a single head for multiple tokens, so that all heads in the group share the same sparse mask.
2. The inner loop over key blocks iterates only over the visible blocks — those in the set $I(i)$ determined by block selection — rather than over all blocks.
3. The key block size $B_k$ must divide the sparse attention block size $B$ (i.e., $B$ is a multiple of $B_k$), ensuring that sparse block boundaries align with FlashAttention's tiling structure.

Why the two-pass approach is the right trade-off: The alternative — writing all per-head scores to HBM and performing the group summation in a separate kernel — would be simpler to implement but would incur the full I/O cost of materializing the attention score matrix. The two-pass approach doubles the computation (mitigated to 1.25× by the LSE approximation) but reduces HBM writes by a factor of $G = 16$. In modern GPU architectures, HBM bandwidth is the primary bottleneck for attention computation — arithmetic throughput is abundant. Trading additional arithmetic for reduced memory traffic is almost always the right decision, which is the same insight that motivated FlashAttention's original design. The paper's profiling results (Table 5) confirm this: on A100 at 128K sequence length, the LSE Approximation reduces block selection time from 75.36 ms to 56.59 ms, a 25% reduction that matches the expected $1.25\times$ overhead reduction.

Future work noted but not implemented: The paper acknowledges that the max-pooling and top-k operations could also be fused into the kernel to further reduce HBM traffic, but leaves this implementation for future work. Currently, the max pooling from $S_{shared}$ to $S_{cmp}$ and the top-k selection over $S_{cmp}$ happen in a separate kernel, requiring an additional read and write of the intermediate data.

---

Dense-Sparse Switch: Sequence-Length-Gated Routing

The switch between dense and sparse attention is implemented as a simple conditional: if the input sequence length $n$ is below a threshold, run the standard dense FlashAttention kernel (Algorithm 3); if above, run the sparse attention pipeline (Algorithm 1 + Algorithm 2).

Why this is zero-cost: Because the parameters are shared, there is no state to swap, no weights to reload, and no architectural reconfiguration needed. The same $W_Q, W_K, W_V, W_O$ are used in both paths. The only difference is which CUDA kernel is launched. This means the switch can be made on a per-sequence basis with essentially no overhead — the model can process a batch containing both short and long sequences, with each sequence taking the appropriate path.

Implication for short-sequence efficiency: In NSA, even a short sequence of 1K tokens must execute all three attention modules (Compressed, Selected, Sliding) and the gating mechanism, because the architecture provides no fallback. This imposes a constant-factor overhead on every forward pass, regardless of sequence length. InfLLM-V2 avoids this entirely — for short sequences, it uses the same dense FlashAttention that the model was pretrained with, achieving identical speed. This is especially important for production deployments where the sequence length distribution is heavily skewed toward short inputs but occasionally includes very long ones. The model should not pay a speed penalty on the common case to enable the rare case.

The switch threshold: The paper does not explicitly specify the sequence length threshold for switching, but the experimental setup provides context. The model is pretrained on 4K-length sequences, and the sparse attention is configured to select 96 blocks of size 64, for 6,144 visible tokens. This suggests that the switch threshold is around the pretraining length (4K) — for sequences shorter than this, dense attention is both faster (because the $O(n^2)$ cost is still manageable and the block selection overhead would dominate) and more accurate (because the model was trained with full attention at this length). For sequences longer than this, sparse attention becomes faster (because the $O(n \cdot |I|)$ cost of sparse attention grows more slowly than the $O(n^2)$ cost of dense attention) with minimal accuracy loss.

Dense mode after long-context finetuning: A critical validation is that the model can switch back to dense mode even after finetuning on long sequences with sparse attention, without performance degradation on short-sequence tasks. Table 4 demonstrates this: InfLLM-V2 in Dense mode achieves an average of 66.76% across seven general benchmarks, compared to 67.41% for the full-attention baseline that was also finetuned on long sequences (both dense during finetuning), and 67.73% for the original pretrained model. The fact that dense mode maintains performance means the KV representations learned during sparse finetuning remain compatible with dense attention — another consequence of the shared-parameter design.

Design choice: why not always use sparse attention? For very short sequences (e.g., $n < 1K$), the block selection overhead (computing compressed attention scores, max pooling, top-k selection) can exceed the cost of simply computing full dense attention. Dense FlashAttention for 1K tokens is extremely fast — the $O(n^2)$ term is small in absolute terms, and the constant-factor overhead of block selection would dominate. The switch allows the model to use the asymptotically optimal algorithm at each sequence length: dense for short, sparse for long.

---

Overall Training and Configuration Summary

The paper's experimental instantiation of InfLLM-V2 uses the following configurations, which are the result of design choices described above:

Model architecture: 8B parameters, GQA with $d = 4096$, $h_q = 32$, $h_{kv} = 2$, $d_h = 128$, group size $G = 16$.

Sparse attention hyperparameters:

• Compression block size $l_{C1} = 32$, stride $s_{C1} = 16$
• Attention block size $B = 64$
• LSE approximation block size $l_{C2} = 128$, stride $s_{C2} = 64$
• Selected block count $|I| = 96$ (with $|I_{init}| = 1$, $|I_{topk}| = 63$, $|I_{local}| = 32$)
• Total visible tokens: $96 \times 64 = 6,144$

Training:

• Pretraining: 8T tokens of 4K-length sequences (FineWeb-Edu + Stack-v2), batch size 8M tokens, WSD learning rate scheduler with 2000 warmup steps to initial lr of $7.5 \times 10^{-3}$, 27,000 decay steps to final lr of $3 \times 10^{-4}$
• Long-context finetuning: 5B tokens, sequences from four length intervals (0-4K, 4-12K, 12-24K, 24-32K) in 1:1:1:1 ratio, initial lr $3 \times 10^{-4}$ linearly decaying to $2.75 \times 10^{-4}$

4. Key Insights and Innovations

Innovation 1: Parameter Identity as the Necessary Condition for Trainable Sparse Attention in the Pretrain-on-Short, Finetune-on-Long Paradigm

This is not an incremental engineering improvement over NSA — it is a diagnosis of a fundamental architectural constraint that the field had not articulated. Before InfLLM-V2, the design space for trainable sparse attention was explored primarily through the lens of pattern design: what sparsity patterns (block-sparse, sliding window, token-level selection) and what training recipes (self-distillation, router training, multi-module architectures) yield the best accuracy-efficiency tradeoff. NSA represents the culmination of this line of thinking — a sophisticated three-module design with learned compression, token-level sparsity, and gated output aggregation that achieves strong results when trained from scratch on long sequences.

The paper's pivotal insight is that this entire framing misses a constraint that is decisive for the dominant deployment workflow. In the pretrain-on-short, finetune-on-long paradigm — which is how virtually all practical long-context LLMs are built, because pretraining from scratch on 32K+ sequences is prohibitively expensive — the sparse architecture must satisfy a condition that has nothing to do with its stand-alone performance: it must preserve the parameter state and computational semantics of the pretrained dense model so that the transition from short to long finetuning does not reset the model's capabilities.

This insight is a reframing of the problem rather than a new solution to the old problem. The old problem was: design a sparse attention that achieves high accuracy and high speedup. The new problem is: design a sparse attention that can be swapped into a pretrained dense model without disrupting what the model has already learned. These are different optimization criteria, and the paper shows that they lead to radically different architectural choices.

What the field assumed before this paper: That trainable sparse attention mechanisms could be evaluated solely on their final performance after training, and that any architectural mismatch could be overcome with sufficient finetuning data and compute. This assumption is implicit in NSA's design — the three KV projections, the gating MLP, and the compression module represent a substantial departure from dense attention, but the original NSA paper (training from scratch on long sequences with a large budget) did not surface the mismatch as a problem because there was no short-pretraining stage to disrupt.

What this paper demonstrates instead: The architectural mismatch is not recoverable within practical finetuning budgets. The evidence is stark — NSA and InfLLM-V2 receive identical pretrained models, identical long-context finetuning data (5B tokens), and identical hyperparameters. Yet NSA's RULER performance (59.92%) is barely above the training-free MInference baseline (73.22%) and far below full attention (84.26%), while InfLLM-V2 achieves 82.62%. The training loss curve (Figure 5) provides the mechanistic explanation: NSA experiences a loss spike at the architectural transition that it never fully recovers from, while InfLLM-V2 continues smoothly from the pretrained loss.

This is a negative result with positive implications: the failure mode of NSA is not that sparse attention is inherently lossy, but that parameter proliferation during architectural transition is the specific mechanism of failure. By eliminating this mechanism — zero new parameters, single output, shared KV projections — InfLLM-V2 recovers near-parity with full attention. The implication for future research is that trainable sparse attention designs should be evaluated not just on their asymptotic performance, but on their transition cost — the amount of training compute needed to recover from the architectural switch, and whether full recovery is even possible within realistic budgets.

The paper also provides a theoretical vocabulary for discussing this constraint: the concepts of "parameter alignment," "computational semantics," and "distributional shift at architectural transition" give the field language for reasoning about why some sparse architectures succeed in the finetuning paradigm and others fail. This is a conceptual contribution that extends beyond the specific InfLLM-V2 implementation — it provides a design principle (maintain parameter identity with the pretrained model) that any future trainable sparse attention method must satisfy if it is to be used in the pretrain-on-short, finetune-on-long workflow.

---

Innovation 2: The Architectural Mismatch as a Measurable Phenomenon with a Specific Mechanism

This is a diagnostic contribution rather than a methodological one. The paper does not just show that NSA underperforms — it isolates why it underperforms, and in doing so, introduces a new type of empirical analysis that future work on efficient attention should adopt.

The diagnostic chain works as follows:

First, the paper identifies that NSA introduces three distinct forms of architectural deviation from dense attention: parameter proliferation (three KV projection sets), output multiplicity (three attention outputs gated together), and learned auxiliary components (compression MLP, gating MLP). These are not merely quantitative differences (e.g., "NSA has more parameters") — they are qualitative differences in the structure of computation. During dense pretraining, the model's residual stream, attention outputs, and KV representations evolved under a specific computational regime (single query-to-key attention, single output per head). NSA changes the type signature of the attention layer: from (Q, K, V) → O to (Q, K_cmp, V_cmp, K_slc, V_slc, K_win, V_win) → g_cmp·O_cmp + g_slc·O_slc + g_win·O_win.

Second, the paper provides direct evidence that this type-signature change causes training instability. Figure 5 is the key exhibit — it shows the training loss over time for three conditions: FullAttn (dense attention, long-context finetuning), InfLLM-V2 (sparse attention, same parameter structure as pretraining), and NSA (sparse attention, modified parameter structure). FullAttn continues decreasing smoothly from the pretrained loss. InfLLM-V2 continues with no visible discontinuity. NSA jumps upward by approximately 0.2 loss units — a substantial disruption equivalent to erasing hundreds of steps of pretraining progress — and remains elevated throughout finetuning.

Third, the paper demonstrates that this training instability translates to permanent capability regression on tasks unrelated to long-context processing. Table 4 shows NSA's performance on general benchmarks: MMLU drops from 73.38% (FullAttn) to 68.27%; MATH-500 drops from 54.60% to 44.40%; HumanEval drops from 71.34% to 62.20%. These are short-sequence tasks where attention sparsity is irrelevant — the degradation is purely from the architectural disruption damaging representations that were already learned. InfLLM-V2 in Dense mode, by contrast, achieves 71.29%, 54.80%, and 73.17% respectively — essentially at parity with the full-attention baseline.

Why this is a conceptual innovation rather than just a benchmarking result: The paper provides a causal model for why trainable sparse attention methods succeed or fail in the finetuning paradigm: success depends not on the sparsity pattern itself, but on whether the architectural transition introduces a distributional shift in the residual stream that downstream layers were not trained to handle. This model generates testable predictions: any sparse attention method that maintains parameter identity and single-output semantics should transition smoothly; any method that introduces new parameters or restructures the attention output should experience a loss spike proportional to the magnitude of the perturbation.

This diagnostic framework is valuable beyond the NSA comparison. It suggests that future work on trainable sparse attention should report transition curves (loss as a function of training steps immediately before and after the architectural switch) as a standard diagnostic, alongside final task performance. It also suggests that the initiation protocol for converting dense models to sparse architectures — how the new parameters are initialized, whether the model is gradually transitioned or abruptly switched — is an underexplored research direction that may be as important as the sparse architecture itself.

The paper's identification of NSA's LongPPL (4.24 vs. 2.06 for FullAttn) as evidence that "NSA has not adequately learned long-range dependencies" (Section 4.2) adds a further diagnostic layer: the loss spike does not just delay convergence; it permanently impairs the model's ability to learn the very long-range dependencies that sparse attention is meant to enable. This suggests a catastrophic interference mechanism — the gradients from the new architectural components overwrite the pretrained representations before the model can learn to coordinate them, and the damage is irreversible within the finetuning budget.

---

Innovation 3: Dense-Sparse Switchability as an Architectural Property Rather than an Optimization Target

Most work on efficient attention treats the choice between dense and sparse attention as an optimization problem: find the sparsity pattern, compression ratio, or token selection strategy that maximizes the accuracy-efficiency Pareto frontier. InfLLM-V2 reframes this choice as an architectural property: a well-designed sparse attention mechanism should make the dense/sparse decision trivial — a zero-cost routing operation — by ensuring that both modes use identical parameters and produce structurally identical outputs.

What prior work did: NSA executes all three attention modules regardless of sequence length because the gating mechanism requires all three outputs to produce the final result. There is no "dense mode" in NSA — the architecture is always sparse, always gated, always multi-output. This means that for short sequences, where sparse attention provides no benefit, the model pays the full architectural overhead (three attention computations, gating) on top of the already-sufficient dense attention cost.

Training-free methods like InfLLM and MInference are applied post-hoc to dense models — they can conceptually be turned off for short sequences, but they were never trained with sparsity, so the model's parameters are not optimized to produce informative block selection scores. This creates a different kind of switchability problem: the model can run dense or sparse, but the sparse mode is inherently lossy because the model never learned to operate under the sparsity constraint.

What InfLLM-V2 does differently: By sharing parameters and maintaining single-output semantics, the model learns to produce KV representations that work well for both dense attention and the block selection mechanism. The switch is implemented as a conditional kernel launch, not a learned decision. This has two practical consequences:

1. Short-sequence performance is preserved (Table 4): After long-context finetuning with sparse attention, switching back to dense mode recovers near-identical performance on short-sequence benchmarks. The model has not "forgotten" how to do dense attention because the parameters were updated by gradients from both modes — the sparse finetuning gradients improved the KV representations, and those improved representations benefit dense attention as well.

2. The efficiency sweet spot is accessible (Figure 6, Figure 7): Because the model can use dense attention for short sequences and sparse attention for long ones, the average-case speedup in a mixed-length deployment is much closer to the theoretical maximum than a method that imposes sparse overhead on all sequences. This is not just a convenience — it is what makes the method practically deployable, because production workloads are mixtures of lengths.

Why this is a conceptual shift: The paper is arguing that switchability is a first-class design objective for sparse attention, not an afterthought. This changes the evaluation criteria: a sparse attention method should be judged not just on its long-sequence accuracy and speedup, but also on its short-sequence performance (does it match dense attention?), its switching cost (is there latency or memory overhead for the transition?), and its mixed-length throughput (how does it handle batches with heterogeneous sequence lengths?). These criteria are absent from most prior work, which reports only long-sequence benchmarks.

The paper's demonstration that InfLLM-V2 in Dense mode sometimes outperforms the FullAttn baseline on long-context tasks (Table 1: 88.32% vs. 84.26% on RULER) adds an additional dimension to switchability: the sparse finetuning process may actually improve the model's dense attention capabilities, perhaps because the sparsity constraint forces the KV representations to be more discriminative (a form of regularization). This is a speculative interpretation — the paper does not investigate the mechanism — but it suggests that dense-sparse switchability is not merely a zero-sum tradeoff between modes, but potentially a synergistic training dynamic.

---

Innovation 4: The Block Selection Bottleneck as a First-Class Systems Problem Requiring Algorithm-Hardware Co-Design

The paper's work on the efficient CUDA kernel implementation (Section 3.4) is not just an engineering contribution — it identifies a structural bottleneck in block-sparse attention that prior work had not systematically characterized, and introduces a design pattern (two-pass computation with LSE approximation) that future implementations can adopt.

The bottleneck diagnosis: In block-sparse attention methods, there is a hidden cost that the asymptotic complexity analysis obscures. The attention computation itself is $O(n \cdot |I| \cdot d_h)$ where $|I|$ is the number of selected blocks — this scales linearly with sequence length, which is the desired efficiency gain. But the block selection step requires computing approximate attention scores between all queries and some compressed representation of all keys. This step scales as $O(n \cdot (n/s) \cdot d_h)$, where $s$ is the compression stride. If $s$ is small (to preserve selection granularity), this approaches the $O(n^2)$ cost that sparse attention was designed to avoid.

Prior work either ignored this cost (treating it as negligible constant-factor overhead), absorbed it into the training cost (in methods like SeerAttention where the router is trained offline), or accepted it as necessary (NSA's compression attention is one of the three attention modules and its output is used in the final result, so its cost is partially amortized). InfLLM-V2 identifies that in the finetuning paradigm, where the block selection uses a parameter-free compression pipeline, the memory I/O from materializing the compressed attention scores can dominate the total attention time, because writing $h_q \cdot n^2 / s$ values to HBM is bandwidth-intensive and $s \ll n$.

The solution as a design pattern: The two-pass approach with LSE approximation is not just an optimization trick — it is a general strategy for fusing reductions into attention score computation under the constraint that the reduction (summation over heads) and the normalization (softmax over sequence positions) do not commute. The pattern is:

1. Compute the normalization statistics (log-sum-exp) on a coarsened representation to minimize compute.
2. Recompute the attention scores on the fine-grained representation, apply the precomputed normalization, perform the reduction in SRAM, and write only the reduced result.

This pattern applies to any block-sparse attention method that needs to aggregate per-head attention scores into a shared block selection mask. The LSE approximation (using coarser keys for the first pass) trades a small amount of selection accuracy for a large reduction in compute, and the paper's ablation (Table 1, with/without LSE Approx row) shows that this tradeoff is essentially free in practice.

Why this is an innovation rather than a routine optimization: The paper connects the algorithmic design (how to compute block selection scores) to the hardware constraint (HBM bandwidth is the bottleneck, not FLOPs) in a way that yields a non-obvious design choice: it is better to compute the attention scores twice (once coarse, once fine) than to compute them once and store the intermediate results. This is counterintuitive from a pure algorithmic perspective ("why do the same work twice?"), but it is the correct decision from a hardware-aware perspective ("writing to HBM is ~10× more expensive than doing extra arithmetic in SRAM"). The paper provides quantitative evidence for this tradeoff (Table 5: 75.36 ms → 56.59 ms at 128K on A100 with LSE Approximation) and gives sufficient detail (Algorithm 1, the specific block sizes and strides) that future implementations can replicate and extend the approach.

This contribution also opens a research direction: what other reductions can be fused into FlashAttention-style kernels for sparse attention? The paper notes that max-pooling and top-k selection are not yet fused, suggesting there is further headroom. More broadly, the paper demonstrates that the block selection stage should be treated as a first-class systems problem in sparse attention research, with its own scaling analysis, optimization targets, and hardware-algorithm tradeoffs, rather than being dismissed as a constant-factor overhead.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper evaluates on three categories of benchmarks: long-context understanding (RULER, LongBench, LongPPL), long chain-of-thought reasoning (MATH-500, AIME 24, AIME 25, LiveCodeBench v5 and v6), and general short-sequence tasks (MMLU, MMLU-Redux, CEval, MATH-500, HumanEval, MBPP, BBH). RULER (Hsieh et al., 2024) is a synthetic benchmark with configurable average length, used here at 32K length across 13 subtasks spanning retrieval, multi-hop tracing, and variable tracking. LongBench (Bai et al., 2024) is a bilingual real-world benchmark covering single-document QA, multi-document QA, summarization, few-shot learning, synthetic tasks, and code. LongPPL (Fang et al., 2025) evaluates perplexity on long sequences. Reasoning benchmarks are standard: MATH-500 (Hendrycks et al., 2021b) for competition math, AIME (MAA) for challenging math competition problems, and LiveCodeBench (Jain et al., 2025) for competitive programming. General task benchmarks include MMLU (Hendrycks et al., 2021a), MMLU-Redux (Gema et al., 2025), CEval (Huang et al., 2023), HumanEval (Chen et al., 2021), MBPP (Austin et al., 2021), and BBH (Suzgun et al., 2023).

• Base model(s). All experiments use an 8B-parameter GQA model with hidden size d = 4096, number of query heads h_q = 32, number of KV heads h_kv = 2, and head dimension d_h = 128, yielding a group size G = 16. The model is pretrained from scratch by the authors using full dense attention on 8T tokens of 4K-length sequences, primarily drawn from FineWeb-Edu (Penedo et al., 2024) and Stack-v2 (Lozhkov et al., 2024). This model is marked as SHORT. A WSD learning rate scheduler (Hu et al., 2024) is used with 2000 warmup steps to an initial learning rate of 7.5 × 10^{-3}, followed by 27,000 decay steps to a final learning rate of 3 × 10^{-4}, with a batch size of 8M tokens. The choice of an 8B GQA model is deliberate: it is large enough to exhibit meaningful long-context capabilities but small enough that pretraining and finetuning experiments are computationally tractable, and the GQA structure enables the block-sparse kernel optimizations described in Section 3.4.

• Metrics. For RULER and LongBench, the primary metric is accuracy (percentage of questions answered correctly), with LongBench reporting macro-average across six task categories. For LongPPL, the metric is perplexity (lower is better). For reasoning tasks (MATH-500, AIME, LiveCodeBench), accuracy is reported, with LiveCodeBench using pass@1 on generated code solutions. For general tasks, each benchmark uses its standard metric: accuracy for MMLU, MMLU-Redux, CEval, MATH-500, and BBH; pass@1 for HumanEval and MBPP. Average across all seven general benchmarks is also reported. For efficiency, the metrics are kernel execution time (milliseconds) for the attention computation alone and end-to-end inference time (seconds for time-to-first-token TTFT during prefilling, milliseconds for time-per-output-token TPOT during decoding).

• Baselines. The paper compares against five baselines organized into two categories. Full-attention baselines: FULLATTN is the pretrained SHORT model finetuned on long sequences using full dense attention with identical training data and hyperparameters as InfLLM-V2 (SPARSE). SHORT+YARN applies the YaRN position embedding extension method (Peng et al., 2023) to the pretrained SHORT model without long-context finetuning. Training-free sparse attention baselines: InfLLM (Xiao et al., 2024a) and MInference (Jiang et al., 2024) are applied post-hoc to the FULLATTN model — they use the same dense-pretrained, dense-finetuned weights but with training-free sparse attention masks applied at inference time. Trainable sparse attention baseline: NSA (Yuan et al., 2025) is implemented by converting the pretrained dense model to NSA's architecture (initializing its three sets of KV parameters by replicating the original dense attention parameters) and finetuning with the same long-context data as InfLLM-V2. Because NSA has not released its code, the paper uses an open-source Triton implementation. All sparse attention methods maintain the same sparsity level (the same number of selected blocks |I| = 96, block size B = 64, yielding 6,144 visible tokens) to ensure fair comparison.

• Generation budget / compute accounting. The primary unit of test-time compute for efficiency comparisons is kernel execution time, measured in milliseconds, for the attention computation alone (excluding FFN layers and other components). This isolates the attention mechanism's performance from confounding factors like differing linear layer implementations. For end-to-end measurements, TTFT and TPOT are reported with W4A16 quantization (Frantar et al., 2025) applied, and the time is decomposed into attention time and "other time" (FFN, embeddings, etc.) to show the fraction of total time attributable to attention. All timing measurements are conducted at batch size 1 on NVIDIA A100 and NVIDIA 4090 GPUs. For fair comparison with NSA, the paper ignores NSA's sliding attention component in kernel benchmarks and compares solely on the compression and sparse attention parts by selecting an equal number of blocks |I|.

• Cross-validation / statistical protocol. The paper does not employ cross-validation or statistical significance testing. Results are reported as single-point measurements on the standard test sets of each benchmark. For RULER, the specific evaluation length is 32K tokens. For LongPPL, the evaluation uses sequences up to the model's maximum context length. The long-context finetuning data uses sequences from four length intervals (0-4K, 4-12K, 12-24K, 24-32K) with token counts in a 1:1:1:1 ratio. The paper does not report standard deviations, confidence intervals, or multiple training runs with different random seeds for any result. This is a meaningful limitation: for benchmarks like RULER with 13 subtasks, some of which may have small sample sizes, variance across runs could affect the reliability of the comparisons, particularly the narrow gap between InfLLM-V2 (82.62%) and FULLATTN (84.26%).

Main Quantitative Results

Long-Context Understanding

RULER at 32K (Table 1). InfLLM-V2 (SPARSE) achieves 82.62% average accuracy across the 13 RULER subtasks, compared to FULLATTN's 84.26% — a gap of only 1.64 percentage points. This represents 98.1% retention of the full-attention baseline's performance, the headline efficiency-performance tradeoff figure claimed in the abstract. The breakdown across subtasks reveals where InfLLM-V2 succeeds and where it loses ground:

• On single-key retrieval tasks (SG1, SG2, SG3), InfLLM-V2 achieves a perfect 100.00% across all three, matching FULLATTN exactly. This indicates that the block selection mechanism reliably identifies and attends to blocks containing the target key, even when the key must be retrieved from among thousands of distractor tokens.

• On multi-key retrieval tasks (MK1, MK2, MK3), InfLLM-V2 achieves 94.00%, 82.00%, and 62.00% respectively, compared to FULLATTN's 96.00%, 94.00%, and 92.00%. The degradation is mild on MK1 and MK2 but substantial on MK3 — a 30-point drop that suggests InfLLM-V2 struggles when the task requires simultaneously tracking multiple widely-separated pieces of information. MK3, which likely involves the longest-range dependencies among the multi-key subtasks, exposes the fundamental limitation of selecting only 6K visible tokens from a 32K context: if the required keys fall in blocks that are not selected, the model cannot retrieve them.

• On multi-value tracking (MV), InfLLM-V2 achieves 98.50%, dramatically outperforming FULLATTN's 82.00% — a 16.5-point improvement. This is a striking anomaly: how can a sparse attention mechanism that sees less than 20% of the context outperform full attention? The paper offers no explanation, but a plausible hypothesis is that the sparsity constraint acts as a regularizer that prevents the model from being distracted by irrelevant tokens during value tracking, forcing it to focus on the most salient positions.

• On multi-query (MQ) and variable tracking (VT), InfLLM-V2 achieves 94.50% and 98.00%, close to FULLATTN's 98.50% and 93.20%. The VT result is another case where sparse attention outperforms dense, albeit by a smaller margin (4.8 points).

• On the common word extraction (CWE) and frequent word extraction (FWE) subtasks, InfLLM-V2 achieves 50.40% and 82.67%, compared to FULLATTN's 44.40% and 91.33%. CWE shows a 6-point improvement under sparsity, while FWE shows an 8.66-point degradation.

• On QA subtasks (QA1, QA2), InfLLM-V2 achieves 72.00% and 40.00%, compared to FULLATTN's 48.00% and 56.00% — a substantial improvement on QA1 (+24 points) and a substantial degradation on QA2 (−16 points).

This pattern — InfLLM-V2 outperforms full attention on some subtasks while underperforming on others — suggests that the sparsity constraint is not uniformly lossy; rather, it changes the model's attention allocation in ways that can be either beneficial or detrimental depending on the task structure. The paper does not analyze this heterogeneity, which is a missed opportunity.

Comparison to sparse baselines (Table 1). InfLLM-V2 (SPARSE) substantially outperforms all other sparse attention methods:

• NSA achieves only 59.92%, a 22.7-point gap from InfLLM-V2. NSA's failures are concentrated in the multi-key (54%, 38%, 30%), multi-value (59.00%), and variable tracking (56.00%) subtasks — precisely the tasks requiring coordinated attention across dispersed context positions. This is consistent with the paper's architectural mismatch diagnosis: NSA's disrupted pretrained representations prevent it from learning effective long-range attention patterns during finetuning.

• MInference achieves 73.22%, trailing InfLLM-V2 by 9.4 points. As a training-free method, MInference is limited by the quality of attention patterns learned during dense training — it can select relevant blocks but the model's KV representations were never optimized to be informative under block-level selection.

• InfLLM achieves only 27.94%, performing catastrophically on multi-key and multi-query subtasks (4-10% range). This is surprising given that InfLLM is the direct predecessor of InfLLM-V2's block selection mechanism, but the training-free application to a model not optimized for block-sparse attention proves inadequate at this sparsity level.

• SHORT+YARN achieves 40.63%, confirming that position embedding extension alone (without long-context finetuning) is insufficient for 32K-length tasks.

InfLLM-V2 (DENSE) vs. FULLATTN. A noteworthy finding is that InfLLM-V2 in Dense mode achieves 88.32%, outperforming FULLATTN's 84.26% by 4.06 points. This means the model trained with sparse attention during long-context finetuning, when switched back to dense attention for evaluation, actually performs better than a model trained with dense attention throughout. The paper notes this observation but does not investigate the mechanism. Possible explanations include: (1) the sparsity constraint during finetuning acts as a regularizer that prevents overfitting to the finetuning data distribution; (2) the model learns more robust KV representations because it must produce informative scores under aggressive pooling during block selection; (3) the sparse finetuning effectively trains the model on a curriculum of incomplete attention, and switching to dense attention at test time provides additional information that the model can exploit.

LongBench and LongPPL (Table 2). InfLLM-V2 (SPARSE) achieves 42.54% on LongBench (macro-average over six categories), nearly identical to FULLATTN's 42.30% and actually slightly higher. The per-category breakdown (Table 6, Appendix B) shows that InfLLM-V2 matches or exceeds FULLATTN on most categories, with notable improvements on HotpotQA (54.11% vs. 50.13%), LCC (44.73% vs. 35.72%), and RepoBench-P (44.62% vs. 35.00%). The code-related improvements (LCC, RepoBench-P) are particularly interesting because they suggest that block-sparse attention may be well-suited to code understanding tasks where relevant context is often structurally localized (function definitions, class declarations) rather than uniformly distributed.

On LongPPL, InfLLM-V2 (SPARSE) achieves 2.12 perplexity, close to FULLATTN's 2.06 and dramatically better than NSA's 4.24. The low perplexity confirms that InfLLM-V2 learns coherent long-range language modeling despite seeing only 6K tokens per attention operation, validating that the block selection mechanism successfully identifies the most predictive context.

InfLLM-V2 with vs. without LSE Approximation. The w/ LSE Approx variant achieves 82.62% on RULER versus 82.09% without — a 0.53-point improvement despite using approximate normalization statistics. This counterintuitive result (the approximation should be strictly lossy in theory) suggests that the approximation noise is below the threshold where it affects block selection decisions, and the slight difference is within the range of evaluation noise. On LongBench, the w/ LSE Approx variant achieves 42.54%; the without-LSE-Approx variant is not reported for LongBench.

Long Chain-of-Thought Reasoning

Reasoning benchmarks (Table 3). The paper evaluates on long-output scenarios by finetuning InfLLM-V2 and baselines on OpenMathReasoning (Moshkov et al., 2025) and OpenCodeReasoning (Ahmad et al., 2025) and evaluating on MATH-500, AIME 24, AIME 25, and LiveCodeBench. The key results:

• InfLLM-V2 (SPARSE) achieves an average of 42.66% across the five benchmarks, compared to FULLATTN's 42.79% — a difference of only 0.13 percentage points, representing 99.7% retention of the full-attention baseline's performance.

• On individual benchmarks, the performance is remarkably close: MATH-500: 87.80% vs. 86.00% (InfLLM-V2 higher by 1.8 points); AIME 24: 38.33% vs. 37.50% (higher by 0.83 points); AIME 25: 29.38% vs. 30.63% (lower by 1.25 points); LCB v5: 29.94% vs. 30.67% (lower by 0.73 points); LCB v6: 27.83% vs. 29.14% (lower by 1.31 points). None of the differences exceed 2 points, which is likely within the variance of single-evaluation measurements.

• NSA achieves 37.28% average, substantially below both InfLLM-V2 and FULLATTN. The gap is particularly large on AIME 24 (28.75% vs. 37.50%) and AIME 25 (23.54% vs. 30.63%), suggesting that NSA's architectural disruption impairs the model's mathematical reasoning capability — a capability that depends on precise multi-step attention to intermediate reasoning steps.

• InfLLM-V2 (DENSE) achieves 40.53%, lower than the Sparse variant (42.66%) and FULLATTN (42.79%). This is a different pattern from the long-input tasks, where Dense mode outperformed FULLATTN. On reasoning tasks, being trained with sparse attention and evaluated with dense attention appears to provide no benefit — possibly because the long chain-of-thought generation inherently involves attending to very recent context (the immediately preceding reasoning steps), and the sparse training has optimized the model to focus on a limited set of blocks, which is better matched to the sparse evaluation mode.

What these results demonstrate: The near-parity between InfLLM-V2 and FULLATTN on reasoning tasks is arguably the stronger result than the RULER numbers, because chain-of-thought reasoning involves generating long sequences autoregressively, where the model must attend to its own previously generated tokens. This is a more dynamic attention pattern than the static long-input tasks in RULER — the model's keys and values are constantly shifting as new tokens are generated, and the attention pattern must adapt online. That InfLLM-V2 maintains performance in this setting suggests that the block selection mechanism works robustly even when the context is self-generated rather than provided as input.

General Short-Sequence Tasks

General benchmarks (Table 4). This evaluation is critical for validating the paper's claim that InfLLM-V2 can switch back to dense mode without performance degradation on short-sequence tasks after long-context finetuning. The results:

• InfLLM-V2 (DENSE) achieves an average of 66.76% across seven benchmarks, compared to SHORT's 67.73% (the original pretrained model before any long-context finetuning) and FULLATTN's 67.41% (the model finetuned on long sequences with dense attention). The 0.97-point drop from SHORT and 0.65-point drop from FULLATTN are modest, confirming that the long-context finetuning with sparse attention does not substantially degrade short-sequence capabilities.

• NSA achieves only 60.63%, substantially below all other methods. The degradation is most severe on MATH-500 (44.40% vs. 54.60% for FULLATTN), HumanEval (62.20% vs. 71.34%), and MBPP (65.00% vs. 75.10%). This confirms that NSA's architectural disruption during long-context finetuning causes permanent regression on capabilities that were already well-learned during pretraining.

• The InfLLM-V2 (DENSE) results are particularly notable because they demonstrate that the model has not "forgotten" dense attention patterns despite being finetuned exclusively with sparse attention on long sequences. The shared KV projection design ensures that improvements to KV representations learned during sparse finetuning (to make block selection scores informative) transfer back to dense attention, and the single-output design ensures that the residual stream distribution remains compatible with downstream layers.

Efficiency

Kernel-level attention speed (Figure 6). The paper measures the execution time of the attention computation alone (excluding FFN and other layers) for FullAttn (FlashAttention-2), NSA, and InfLLM-V2, breaking InfLLM-V2's time into Block Selection and Sparse Attention components. Measurements are taken on NVIDIA A100 and NVIDIA 4090 at sequence lengths 32K, 64K, 96K, and 128K, with the number of visible tokens (selected blocks × block size) swept from 1K to 6K.

On A100 at 128K sequence length with 1K visible tokens (highest sparsity):

• FullAttn (FlashAttention): baseline dense attention time
• NSA: reported speedup of 2.9× over FlashAttention
• InfLLM-V2: 7.4× speedup over FlashAttention — 2.55× faster than NSA at the same sparsity level

With 6K visible tokens (the configuration used in accuracy experiments, |I| = 96, B = 64):

• FullAttn (FlashAttention): baseline
• NSA: speedup of 1.4×
• InfLLM-V2: speedup of 3.1× — 2.2× faster than NSA

The speedup increases with sequence length (7.4× at 128K vs. 3.0× at 32K on A100 with 1K visible tokens) because the relative cost of dense attention's O(n^2) scaling grows faster than InfLLM-V2's O(n × |I|) scaling. The speedup decreases with more visible tokens (7.4× at 1K vs. 3.1× at 6K at 128K) because the sparse attention computation itself becomes more expensive as the number of selected blocks increases.

Block Selection vs. Sparse Attention breakdown (Figure 6, stacked bars). A critical finding from the breakdown is that the Block Selection overhead in InfLLM-V2 has been substantially reduced compared to NSA. In NSA, the block selection (compressed attention) computation represents a significant fraction of total time. In InfLLM-V2, the Block Selection time (shown in the stacked bars) is small relative to the Sparse Attention time, particularly at higher sequence lengths. This validates the effectiveness of the fused head group summation kernel (Algorithm 1) and the LSE approximation described in Section 3.4.

LSE Approximation ablation (Table 5). The paper measures Block Selection time with and without the LSE approximation at various sequence lengths on both A100 and NVIDIA 4090, with the number of selected blocks fixed at 16:

On A100:

• 32K: 3.93 ms vs. 4.67 ms (15.9% reduction)
• 64K: 14.07 ms vs. 18.20 ms (22.7% reduction)
• 96K: 32.44 ms vs. 42.46 ms (23.6% reduction)
• 128K: 56.59 ms vs. 75.36 ms (24.9% reduction)

On NVIDIA 4090:

• 32K: 3.70 ms vs. 4.89 ms (24.3% reduction)
• 64K: 14.39 ms vs. 19.95 ms (27.9% reduction)
• 96K: 33.16 ms vs. 46.51 ms (28.7% reduction)
• 128K: 59.04 ms vs. 83.26 ms (29.1% reduction)

The reduction is approximately 25% at longer sequence lengths, consistent with the expected reduction from $2\times$ to $1.25\times$ overhead (saving $0.75\times$ out of $2\times$, or 37.5% of the original overhead). The slightly lower-than-theoretical reduction suggests that the LSE approximation introduces some additional I/O or compute overhead beyond the simple factor-of-4 coarsening, but the achieved speedup is substantial. The NVIDIA 4090 shows slightly better relative reduction than the A100, likely due to differences in memory bandwidth and compute throughput ratios.

End-to-end inference speed (Figure 7). The end-to-end measurements include all model components (attention, FFN, embeddings, etc.) with W4A16 quantization, using the production configuration of |I| = 96 (6K visible tokens). The results are reported as TTFT (time-to-first-token, measuring prefilling speed) and TPOT (time-per-output-token, measuring decoding speed):

On A100:

• Prefilling TTFT: speedup of 1.09× at 32K, 1.39× at 64K, 1.67× at 96K, 1.99× at 128K
• Decoding TPOT: speedup of 1.17× at 32K, 1.43× at 64K, 1.63× at 96K, 1.82× at 128K

On NVIDIA 4090:

• Prefilling TTFT: speedup of 1.16× at 32K, 1.49× at 64K, 1.79× at 96K, 2.13× at 128K
• Decoding TPOT: speedup of 1.43× at 32K, 1.78× at 64K, 2.03× at 96K, 2.32× at 128K

The end-to-end speedups are substantially lower than the kernel-level speedups (e.g., 2.13× vs. 7.4× at 128K on 4090) because the attention mechanism is only one component of the Transformer layer — the FFN layers, layer norm, embeddings, and other operations are not accelerated by sparse attention. The paper explicitly notes this: "Since InfLLM-V2 does not accelerate the Feed-Forward Network (FFN) layers, a higher speedup ratio can be achieved by incorporating FFN-specific acceleration techniques in future work" (Section 4.3). The stacked bars in Figure 7 decompose the total time into attention time and other time, showing that attention dominates at longer sequences but FFN and other components become the bottleneck as attention is accelerated.

The NVIDIA 4090 shows consistently better speedups than the A100, likely due to the 4090's higher memory bandwidth relative to its compute throughput, making the I/O reductions from sparse attention more impactful. This hardware-dependence is not discussed in the paper but is an important practical consideration for deployment.

Ablation Studies and Robustness Checks

InfLLM-V2 with vs. without LSE Approximation on RULER (Table 1, bottom two rows): The variant without LSE Approximation achieves 82.09% on RULER versus 82.62% with the approximation — a 0.53-point difference. The approximation actually yields slightly higher average accuracy. Per-subtask inspection shows that the differences are small and bidirectional (some subtasks improve, others degrade), consistent with the approximation noise being uncorrelated with block selection quality. This ablation validates that the LSE Approximation, which reduces block selection overhead by approximately 25% (Table 5), does not meaningfully degrade the quality of block selection decisions.

InfLLM-V2 (SPARSE) vs. InfLLM-V2 (DENSE) on all benchmarks (Tables 1, 2, 3, 4): The paper provides Dense mode results throughout, which serves as an implicit ablation of the sparsity constraint itself. On RULER, Dense mode (88.32%) outperforms Sparse mode (82.62%) by 5.7 points, indicating that sparsity does impose an upper bound on long-input task performance — even with trainable block selection, losing access to 81% of the context limits achievable accuracy. On LongBench, Dense (42.49%) and Sparse (42.54%) are essentially identical, suggesting that the real-world long-context tasks in LongBench are less sensitive to the sparsity constraint than the synthetic RULER tasks. On reasoning benchmarks, Sparse (42.66%) slightly outperforms Dense (40.53%), a reversal that the paper does not explain but which may relate to the autoregressive generation setting. On general benchmarks, Dense (66.76%) performance confirms that sparse finetuning does not degrade short-sequence capabilities.

NSA comparison as an architectural ablation (Tables 1-4, Figure 5): The NSA baseline serves as an ablation of the "parameter-free, single-output" design principles. NSA differs from InfLLM-V2 in three ways: it uses three separate KV projections (vs. one shared), three attention outputs with gating (vs. one unified), and a learned MLP for compression (vs. parameter-free pooling). The large performance gap — 59.92% vs. 82.62% on RULER, 37.10 vs. 42.54 on LongBench, 4.24 vs. 2.12 on LongPPL, 37.28% vs. 42.66% on reasoning, 60.63% vs. 66.76% on general tasks — constitutes strong evidence that these architectural differences collectively cause the training instability and capability regression. However, this is a bundled ablation: the paper does not isolate which of NSA's three deviations (multiple KV projections, multi-output gating, or learned compression) is primarily responsible for the degradation. An ablation that, for example, gave NSA shared KV projections but kept the multi-output gating would help disambiguate the mechanisms.

Training loss curves (Figure 5) as a mechanistic ablation: The training loss trajectories for FullAttn, InfLLM-V2, and NSA during the long-context finetuning phase show that InfLLM-V2's loss continues smoothly from the pretrained value (approximately 1.2), while NSA's loss spikes to above 1.4 at the transition and remains elevated. This figure is the paper's most direct evidence that architectural mismatch — not the sparsity constraint itself — causes NSA's underperformance. The fact that InfLLM-V2's loss closely tracks FullAttn's throughout finetuning suggests that the sparsity constraint alone (seeing only 6K out of 32K tokens) does not significantly impair the model's ability to learn from long-context data.

InfLLM vs. MInference as training-free vs. trainable ablation (Tables 1, 2): The comparison between training-free InfLLM (27.94% on RULER) and trainable InfLLM-V2 (82.62%) using the same block-sparse attention pattern demonstrates the value of finetuning with the sparsity constraint. The training-free InfLLM applies the block selection mechanism to a model that was never optimized to produce informative block selection scores, resulting in poor block choices. The ablated variable is whether the model receives gradient signals that encourage its KV representations to be discriminative under the compression and selection operations — the 54.68-point gap quantifies the importance of this training signal.

SHORT+YARN as a position extension ablation (Tables 1, 2): SHORT+YARN (40.63% on RULER) demonstrates that extending the position embeddings to support longer sequences, without any long-context finetuning data, is insufficient. The gap from FULLATTN (84.26%) quantifies the value of long-context finetuning data, while the gap from InfLLM-V2 (82.62%) confirms that sparse attention with finetuning can recover most of the benefit of dense attention with finetuning.

Block size and sparsity configuration: The paper does not ablate the block size B, the number of selected blocks |I|, the number of initial blocks |I_init|, the number of local blocks |I_local|, or the compression block size l_C1. All experiments use a single configuration (B = 64, |I| = 96, |I_init| = 1, |I_local| = 32, |I_topk| = 63, l_C1 = 32). This is a significant gap: the tradeoff between sparsity ratio and accuracy is a central design axis for any sparse attention method, and the paper provides no guidance on how sensitive InfLLM-V2's performance is to these hyperparameters. Would a larger |I| (more visible tokens) close the remaining gap to full attention? Would a smaller |I| (fewer visible tokens) still maintain acceptable accuracy while providing greater speedup? These questions are left unanswered.

Group size G for GQA: The paper uses G = 16 throughout and does not ablate this choice. The block-sparse kernel design relies on the group size to share sparse masks across heads, and the efficiency of the fused head group summation depends on G. Smaller G would reduce the I/O savings from the fusion; larger G would increase them but might constrain the model's expressivity. The paper does not explore this tradeoff.

Finetuning data scale: The paper uses 5B tokens for long-context finetuning and does not ablate this quantity. Given that NSA's underperformance is attributed to insufficient recovery from architectural disruption, it would be informative to know whether additional finetuning data (10B, 20B tokens) allows NSA to eventually catch up, or whether the damage is permanent. This is particularly relevant because the paper's claim that NSA is "unsuitable for the pretrain-on-short, finetune-on-long paradigm" depends on the finetuning budget being realistically constrained — if unlimited finetuning could recover NSA's performance, the claim would be about training efficiency rather than fundamental architectural compatibility.

Critical Assessment

Claim 1: "InfLLM-V2 is 4× faster than dense attention while retaining 98.1% and 99.7% of the performance"

The abstract's headline claim requires careful deconstruction because the "4×" and the "98.1% / 99.7%" come from different experimental configurations and are not simultaneously achievable in a single deployment.

The 4× speedup is measured at the kernel level on A100 at 128K sequence length with 1K visible tokens (Figure 6, top row, leftmost bar group). This is the most aggressive sparsity setting — only 1K tokens visible out of 128K, a 0.78% density. However, the accuracy experiments (Tables 1-3) use 6K visible tokens (|I| = 96), where the kernel-level speedup at 128K is 3.1× on A100 and the end-to-end speedup is 1.99× on A100 (prefilling) and 1.82× (decoding). The 4× figure does not correspond to any configuration for which accuracy is reported.

The 98.1% retention refers to RULER accuracy (Table 1): InfLLM-V2 (SPARSE) at 82.62% vs. FULLATTN at 84.26%. The 82.62 / 84.26 = 0.981 ratio is correct arithmetic. However, this comparison is at the 6K-visible-token configuration, where the end-to-end speedup on A100 at 32K (RULER's evaluation length) is only 1.09× for prefilling and 1.17× for decoding (Figure 7) — nowhere near 4×. The 98.1% retention and the 4× speedup are achieved at different sparsity levels, at different sequence lengths, and at different measurement granularities (kernel vs. end-to-end).

The 99.7% retention refers to reasoning tasks (Table 3): InfLLM-V2 (SPARSE) at 42.66% average vs. FULLATTN at 42.79%. The 42.66 / 42.79 = 0.997 ratio is correct. But reasoning tasks involve autoregressive generation where the sequence length grows incrementally — the attention cost during early tokens (short sequences) is dominated by operations other than attention, and the sparse attention benefit only materializes for later tokens when the context has grown long. The paper reports no speedup measurements specific to the reasoning setting.

The paper would be more precise to state: "At a sparsity level that retains 98.1% of RULER accuracy, InfLLM-V2 achieves 1.09× end-to-end prefilling speedup at 32K on A100 and scales to 3.1× kernel-level speedup at 128K. At more aggressive sparsity, kernel-level speedup reaches 7.4× (128K, 1K visible tokens) but accuracy at this sparsity level is not reported." This would accurately characterize the accuracy-efficiency tradeoff without conflating different operating points into a single headline.

Claim 2: "InfLLM-V2 seamlessly adapts models from short to long sequences ... maintaining consistency between short and long sequence processing"

The evidence for seamless adaptation is strong. Figure 5 shows no loss discontinuity when switching from dense pretraining to sparse finetuning. Tables 1-3 show that the sparse-finetuned model achieves near-parity with the dense-finetuned model on long-context tasks. Table 4 shows that switching back to dense mode recovers short-sequence performance (66.76% vs. 67.41% for the dense-finetuned baseline).

However, the claim of "consistency" requires qualification. The model's attention pattern in sparse mode is fundamentally different from dense mode — it sees only 6K out of 32K tokens. This is not "consistent" processing; it is an approximation that happens to work well for the evaluated benchmarks. The RULER subtask breakdown shows that this approximation is not uniformly successful: MK3 drops from 92.00% to 62.00%, QA2 drops from 56.00% to 40.00%. For tasks that require attending to information distributed across more than 6K tokens of the 32K context, the sparse model will necessarily fail in ways the dense model would not. The paper does not characterize which types of long-context reasoning are fragile under sparsity, limiting the practical utility of the "seamless" claim.

Additionally, the "consistency" between short and long sequence processing is demonstrated only for the specific finetuning recipe (5B tokens, 1:1:1:1 length ratio). The paper does not explore whether other finetuning recipes (different length distributions, different data mixtures) would preserve this consistency or whether the model's short-sequence performance is robust to these choices.

Claim 3: "NSA introduces excessive extra parameters and disrupts the conventional pretrain-on-short, finetune-on-long workflow, resulting in slow convergence and difficulty in acceleration"

The evidence for this claim is largely convincing but has important limitations.

The "excessive extra parameters" claim is supported by the architectural description (Section 3.1) but the paper never reports the actual parameter count overhead of NSA relative to InfLLM-V2 or dense attention. How many extra parameters does NSA introduce? What fraction of the total 8B model parameters do they represent? Without these numbers, "excessive" is a qualitative judgment.

The "disrupts the workflow" claim is supported by Figure 5 (loss spike) and Tables 1-4 (performance degradation). However, the NSA implementation used in the paper may not be representative of what a well-optimized NSA adaptation could achieve. The paper notes that "NSA does not publish their code, we adopt an open-source Triton implementation of NSA for experiments" (Section 4.1). This introduces two potential confounds: (1) the Triton implementation may have bugs or performance issues not present in NSA's original CUDA kernels; (2) the initialization strategy (replicating the original KV parameters three ways) and the training hyperparameters (identical to InfLLM-V2's) may not be optimal for NSA's more complex architecture. The original NSA paper demonstrated strong performance when training from scratch on long sequences; the paper's claim that NSA fails specifically in the pretrain-on-short, finetune-on-long setting would be strengthened by showing that NSA can succeed in the train-from-scratch setting under the authors' own experimental setup, confirming that the implementation is functional and the failure is genuinely due to the finetuning paradigm rather than implementation issues.

The "difficulty in acceleration" claim is supported by Figure 6, which shows NSA's kernel speedup is limited to 3.5× at 128K with 1K visible tokens on A100, compared to InfLLM-V2's 7.4×. The stacked bar breakdown shows that NSA's block selection overhead is substantially larger than InfLLM-V2's. However, this comparison ignores NSA's sliding attention component (the paper explicitly states: "For a fair efficiency comparison with NSA, we ignore its sliding attention component"), which may change the speedup calculation. The paper is transparent about this but the comparison is not truly "fair" in the sense of comparing complete methods — it compares InfLLM-V2's complete pipeline against a subset of NSA's pipeline.

Claim 4: "InfLLM-V2 ensures computational efficiency across all sequence lengths, by using dense attention for short inputs and smoothly transitioning to sparse attention for long sequences"

This claim is well-supported for the configurations tested. Figure 7 shows that at 32K, the end-to-end speedup is modest (1.09× prefilling on A100) but positive, meaning the switch to sparse attention does not impose overhead relative to dense attention at the crossover point. The stacked bar decomposition shows that the block selection overhead is small relative to the attention time saved. Table 4 confirms that short-sequence performance in dense mode is preserved.

However, the paper does not measure what happens at the exact crossover point between dense and sparse modes. The threshold for switching is not specified, and there is no measurement of the latency discontinuity (if any) when crossing the threshold. In a production system serving variable-length requests, a latency spike at the dense-to-sparse transition could cause tail-latency problems even if average throughput improves.

Missing Evaluations That Would Strengthen the Paper

1. Scaling behavior of accuracy with number of visible tokens: The paper uses a single configuration (|I| = 96, 6K visible tokens) for all accuracy experiments. A sweep over |I| values (e.g., 32, 64, 96, 128, 192) at fixed sequence length would reveal the accuracy-efficiency Pareto frontier and allow practitioners to choose an operating point based on their accuracy requirements and latency budgets. This is a standard analysis for sparse attention methods that is conspicuously absent.

2. Scaling behavior with sequence length: All accuracy evaluations are at a single sequence length (32K for RULER, up to 32K for LongBench). How does the accuracy gap between InfLLM-V2 and FULLATTN change as sequence length increases to 64K, 96K, 128K? The fixed 6K visible tokens means the sparsity ratio becomes more aggressive at longer lengths (6K/64K = 9.4%, 6K/128K = 4.7%), which may cause accuracy to degrade. Without these measurements, the paper cannot claim that InfLLM-V2 "scales" to arbitrary long contexts — it demonstrates performance at 32K only.

3. Ablation of block selection quality: The paper does not compare InfLLM-V2's block selection against an oracle that selects the blocks containing the highest true attention scores (computed by a full-attention forward pass). Such a comparison would decompose the performance gap to full attention into two components: (a) error from the block selection mechanism choosing suboptimal blocks, and (b) error from the model being unable to use the selected blocks as effectively as full attention. This decomposition would guide future improvement efforts.

4. Multiple random seeds or statistical significance: All results are single-point measurements. For benchmarks where the gap between InfLLM-V2 and FULLATTN is small (LongBench: 42.54 vs. 42.30; reasoning: 42.66 vs. 42.79), it is unclear whether these differences are statistically significant or within the noise of a single evaluation run.

5. Performance on the original NSA training-from-scratch setting: To isolate whether the architectural mismatch is genuinely the cause of NSA's underperformance (rather than implementation issues with the Triton port), the paper could train NSA from scratch on long sequences (no short pretraining) and verify that it achieves strong performance, as reported in the original NSA paper. This would confirm that the NSA implementation is functional and that the degradation is specific to the finetuning paradigm.

6. Memory consumption analysis: The paper focuses entirely on speed (latency, throughput) and does not report GPU memory consumption for InfLLM-V2 versus dense attention. For long-context inference, memory is often the binding constraint (KV cache size scales with sequence length), and one of the motivations for sparse attention is reducing memory pressure. Does InfLLM-V2 reduce peak memory usage during prefilling? Does it reduce KV cache size during decoding (since not all KV pairs need to be cached if only a subset will be attended to)?

7. Performance on longer-than-training context lengths: The model is finetuned on sequences up to 32K. How does InfLLM-V2 perform when evaluated at 64K or 128K — lengths it was never trained on? The original InfLLM paper claimed training-free length extrapolation; does the trainable InfLLM-V2 retain any extrapolation capability, or is it strictly limited to the maximum training length?

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted For, Making Headline Efficiency Gains Partially Theoretical

The compute-optimal allocation framework requires estimating each prompt's difficulty before deciding how to allocate the inference budget. The paper's method for doing so — generating 2048 samples per question and averaging either ground-truth correctness (oracle) or the verifier's 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"

This is not a minor accounting oversight — it undermines the practical interpretation of the paper's central efficiency claim. The reported 4× efficiency gains over best-of-N (Figures 4 and 8) are computed after difficulty is known, without amortizing the cost of learning it. In a realistic deployment, the total cost would be difficulty estimation + strategy execution. The difficulty estimation step alone — 2048 samples per question — is 4–8× larger than the largest test-time budgets studied (256–512 generations). This means the actual total compute for a single question would be dominated by the difficulty estimation step, potentially making the overall process less efficient than simply running a uniform best-of-N with a large budget.

The paper frames this as an exploration-exploitation tradeoff (Section 3.2): compute spent assessing difficulty versus compute spent solving the problem. But without a method for amortizing difficulty estimation across many similar questions, or a cheaper estimator, this tradeoff is hypothetical. The paper suggests future work on "pretraining or finetuning models to directly predict difficulty of a question" (Section 8) but provides no such model or even a feasibility analysis. Until a cheap difficulty estimator is developed and demonstrated, the 4× figure should be understood as an upper bound on achievable efficiency in a deployment where difficulty is known in advance (e.g., from historical query patterns or metadata), not as a realized gain for a cold-start system.

The predicted-difficulty variant partially addresses this by removing the need for ground-truth labels, but it does not remove the 2048-sample cost — it still requires generating and scoring 2048 samples per question using the PRM. The curves in Figures 4 and 8 show predicted difficulty tracking oracle difficulty closely, confirming that the PRM score distribution is a sufficient proxy for ground-truth difficulty, but this is a verification of the proxy's quality, not a reduction in its cost.

---

The Method Provides Zero Benefit on the Hardest Problems, Revealing a Fundamental Capability Ceiling

Across every experiment — search, revisions, and their compute-optimal combinations — the hardest questions (difficulty bin 5) show near-zero improvement regardless of how much test-time compute is allocated. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets (4, 16, 64, 256 generations). 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%, well below the 14× larger model's performance even at the most favorable inference-to-pretraining ratio.

This is not a failure of the specific allocation strategy — it reflects a hard boundary on what test-time compute can achieve. If the base model's pass@1 on a problem class is near zero, no amount of search or revision can help, because there are simply no correct solutions in the proposal distribution to find or refine. Test-time compute amplifies existing capability; it does not create capability from nothing. The paper is transparent about this, stating in the Section 7 takeaway that on hard problems, "pretraining is almost always more effective." But the implications are more severe than the paper emphasizes:

• For any problem distribution with a non-trivial fraction of genuinely hard examples, the average-case benefit of compute-optimal test-time scaling will be diluted. If 20% of user queries fall in difficulty bin 5 (a plausible scenario for challenging real-world applications), the method provides zero lift on those queries while still incurring the difficulty estimation cost.

• There is no diagnostic to determine in advance whether a problem is too hard for test-time compute, because difficulty estimation itself requires sampling. The system would need to spend the full difficulty estimation budget (2048 samples) only to discover that no allocation strategy will help — a worst-case outcome where costs are incurred without benefit.

• The boundary between "recoverable" and "hopeless" problems is model-specific and task-specific. The paper's five-quintile difficulty bins are computed relative to PaLM 2-S* on MATH. A different model on a different task would have different difficulty cutoffs, and there is no general principle for predicting where the boundary lies without running experiments.

The authors acknowledge this limitation in Section 7, noting that for the hardest questions, "no method makes meaningful progress — the base model simply lacks the capability to produce correct solutions." They do not propose a solution, because the limitation is fundamental: test-time compute cannot substitute for capabilities the base model never acquired during pretraining.

---

The PRM Verifier Over-Optimization Problem Is Documented but Not Solved, Placing a Hard Ceiling on Scaling

The paper provides clear evidence that verifier over-optimization — the phenomenon where search finds solutions that score highly under the PRM but are actually incorrect — is the primary bottleneck preventing unbounded improvements from additional test-time compute. The evidence is multi-faceted:

• In Figure 3 (right), beam search degrades performance on easy questions (bins 1–2) at high budgets, a hall-mark of verifier exploitation — the search finds adversarial examples that the PRM erroneously rates highly.
• Lookahead search, the most powerful optimizer (it uses extra computation to get better step-level scores), paradoxically performs worst overall (Figure 3, left), because its more accurate optimization amplifies the PRM's systematic errors rather than correcting them.
• Qualitative examples in Appendix M show search producing degenerate outputs — repetitive low-information steps and overly short 1–2 step solutions — that nevertheless score highly under the PRM.

The compute-optimal allocation policy mitigates but does not solve this problem. By routing easy problems away from aggressive search (using best-of-N instead of beam search), the policy avoids the worst over-optimization regime. But on medium-difficulty problems (bins 3–4), where beam search is deployed because it genuinely helps, over-optimization still limits the scaling ceiling — the beam search curves flatten and sometimes decline before the budget is exhausted (Figure 3, right, bin 3 at 256 generations).

This means the compute-optimal approach is fundamentally bounded by verifier quality, and the paper provides no mechanism for improving verifier robustness beyond the specific Monte Carlo rollout training procedure described in Appendix D. The current results are specific to the PRM quality achievable with that training recipe and that base model (PaLM 2-S*). A model with different output characteristics, or a PRM trained with less on-policy data, could exhibit different over-optimization thresholds, changing the optimal allocation policy. The paper does not explore how verifier improvements (e.g., adversarial training, ensemble methods, or better calibration) would shift the scaling landscape.

The authors acknowledge this implicitly in Section 8, identifying "development of robust verifiers" as a key area for future work. But the current system offers no defense against verifier over-optimization beyond simply avoiding aggressive search on easy problems — a strategy that works only because the paper can identify easy problems via difficulty estimation. In a deployment where difficulty is unknown or miscalibrated, the system might inadvertently apply aggressive search to problems where the verifier is unreliable, degrading performance.

---

The Revision Model's Correct-to-Incorrect Reversion Rate (~38%) Is a Systemic Problem with Only Palliative Fixes

Section 6.1 reports a significant practical problem with the iterative revision approach: approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect answers in the subsequent step. This is a direct and predictable consequence of the training data construction: the revision model was trained exclusively on sequences where all in-context answers are incorrect (followed by a correct target). It has never seen a training example where the current answer is correct and should be preserved. When it encounters a correct answer in its own revision chain at test time, it has no learned behavior for recognizing and preserving it — the model's training objective was "given incorrect answers, produce a correct one," not "given an answer, determine whether it needs revision."

The paper mitigates this with within-chain selection: rather than taking the final revision output, the system uses majority voting or verifier-based selection across the entire chain of revisions to pick the best answer from any step. This is acknowledged as a patch, not a solution:

"To mitigate this, the system uses a selection mechanism (majority voting or verifier-based selection) across the entire chain of revisions, picking the best answer from any point in the chain rather than always taking the last revision."

The consequence is that the revision model's effective yield — the fraction of revision chains that ultimately produce a correct answer — is lower than what an ideal revision model would achieve, because the model periodically destroys correct answers it has already produced. The within-chain selection recovers some of these (if a correct answer appears at step 3 and is revised to incorrect at step 4, selection can still retrieve the step-3 answer), but it may miss cases where the correct answer appears early and the verifier or majority vote prefers a later incorrect answer.

A more principled solution — such as training the model on trajectories that include "no revision needed" steps, or training a separate stopping criterion — is not explored. The ReST^EM experiment (Appendix K, Figure 16) further highlights the fragility of revision training: attempting to optimize the revision model with reinforcement learning caused performance to degrade substantially with sequential revisions, likely because on-policy data collection amplified spurious correlations. This suggests the revision training procedure is sensitive to design choices in ways the paper does not fully characterize, and the positive results depend on specific implementation details (offline data construction, edit-distance-based pairing) that may not transfer robustly.

---

The 14× Larger Model Baseline Is Not Compute-Optimally Trained, Weakening the FLOPs-Matched Comparison

The FLOPs-matched comparison in Section 7 scales model parameters by ~14× while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023) rather than the Chinchilla-optimal paradigm where both data and parameters are scaled equally (Hoffmann et al., 2022). The paper acknowledges this explicitly:

"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."

This matters because a Chinchilla-optimal model trained with 14× more total FLOPs (scaling both model size and training tokens appropriately) would likely outperform a parameter-only-scaled model, making the pretraining baseline weaker than it could be. The reported advantages of test-time compute over pretraining — e.g., +27.8% relative improvement on easy questions in the revisions setting at R<1 (Figure 1, top-right bar chart) — may shrink or reverse against a properly compute-optimal larger model.

The paper argues that the LLaMA-style scaling is "representative" of common practice, and this is a defensible choice — many deployed models are trained with parameter scaling rather than Chinchilla-optimal scaling. But the paper's central claim about the pretraining-inference tradeoff ("test-time compute can substitute for pretraining") depends on the specific baseline. Against a weaker baseline, the substitution appears more favorable. The paper does not bound how much of the observed advantage is attributable to the baseline's suboptimal pretraining rather than the intrinsic power of test-time compute.

Additionally, the 14× larger model uses only greedy decoding — no majority voting, no best-of-N, no search. This is an even weaker baseline than necessary, because it compares a test-time-optimized smaller model against a larger model with no test-time optimization at all. A fairer comparison would give the larger model a modest test-time compute budget (e.g., best-of-8 or best-of-16) proportional to its higher per-token cost. The paper does not include such a comparison, which means the reported advantages of test-time compute conflate two effects: the benefit of test-time optimization and the benefit of the smaller model being given any test-time budget at all.

---

All Results Are on a Single Benchmark (MATH) with a Single Model Family (PaLM 2-S*), Leaving Generality Unverified

Every experiment in the paper — search, revisions, FLOPs-matched comparison — uses the MATH benchmark (Hendrycks et al., 2021) as the evaluation dataset and PaLM 2-S* (Anil et al., 2023) as the base model family. The test set is 500 questions, split into five difficulty quintiles of ~100 each, further halved by cross-validation for strategy selection. The authors state they "believe this model is representative of the capabilities of many contemporary LLMs" (Section 4), but this claim is an assertion, not an empirically supported statement.

Several aspects of the findings could be model- or domain-specific:

• The PRM's quality and over-optimization behavior depend on PaLM 2-S*'s output distribution — its calibration, its error patterns, and the statistical structure of its attention and generation. A model with different properties (e.g., better calibrated probabilities, different beam search behavior) might exhibit different difficulty-dependent scaling curves.

• 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. PaLM 2-S* was chosen partly because it has non-trivial baseline performance (10–19% pass@1) with room for improvement — a model with higher baseline performance might show different revision dynamics, while a model with lower baseline performance might see smaller revision benefits.

• The MATH benchmark consists exclusively of competition-level math problems requiring multi-step symbolic reasoning with clean, verifiable correct answers. It is unclear whether the difficulty-dependent patterns — beam search hurting easy problems but helping medium ones, sequential revisions dominating on easy problems but balanced ratios being optimal on hard ones — generalize to other reasoning domains (code generation, logical reasoning, scientific QA) or to tasks requiring factual knowledge rather than step-by-step inference.

• The test set of 500 questions is small relative to the complexity of the compute-optimal policy. With five difficulty bins and two-fold cross-validation, strategy selection is based on ~50 questions per fold per bin. The paper does not report confidence intervals on the compute-optimal scaling curves, making it impossible to assess whether the observed gains are statistically reliable or whether the selected strategies would generalize to a fresh test set.

The paper's findings about the relationship between difficulty and optimal strategy are its most robust contribution — the qualitative patterns (beam search degrades easy problems, revisions help easy problems) are consistent and striking. But the quantitative claims about efficiency gains (4×, specific accuracy levels) may not transfer to other models, other datasets, or even other random splits of the MATH test set without additional validation.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper does not introduce a fundamentally new sparsity pattern or attention mechanism — the block-sparse structure, the use of compressed key representations for block selection, and the combination of global and local attention are all established ideas. Rather, it introduces a diagnostic reframing of the trainable sparse attention problem. The core contribution is the argument that, for the dominant pretrain-on-short, finetune-on-long workflow, the success of trainable sparse attention depends not on the sophistication of the sparsity pattern but on whether the sparse architecture preserves parameter identity and computational semantics with the pretrained dense model.

This is a shift in design philosophy rather than a new algorithm. Before this paper, the research trajectory for trainable sparse attention was toward increasing architectural complexity — more attention modules, learned routing, gating mechanisms, auxiliary losses — under the implicit assumption that the sparsity problem requires dedicated architectural components to solve. NSA represents the culmination of that trajectory: three parallel attention pathways, learned compression, per-token gating. This paper argues, through empirical evidence, that this trajectory leads to methods that are incompatible with how long-context models are actually built. Most practitioners pretrain on short sequences (4K tokens) and finetune on longer ones, because pretraining from scratch on long sequences is cost-prohibitive. An architecture that cannot survive the transition from dense to sparse without disrupting the pretrained representations is, for practical purposes, unusable regardless of its asymptotic performance.

The paper's diagnosis — that architectural mismatch causes a training discontinuity that practical finetuning budgets cannot fully recover from — identifies a new constraint on the design of efficient attention mechanisms. This constraint is not captured by existing evaluation protocols, which typically measure final task performance after training from scratch or after full finetuning, without examining the transition dynamics. The training loss curve at the architectural switch (Figure 5) becomes a diagnostic that future work in this area should adopt as a standard reporting requirement. If a proposed sparse attention method cannot show a smooth loss transition from the dense pretrained model, it is not suitable for the finetuning paradigm regardless of its other merits.

The paper also reconciles a tension that was latent in the literature but not previously articulated. The original NSA paper demonstrated strong performance when training from scratch on long sequences, yet this paper shows NSA substantially underperforms in the finetuning setting. These are not contradictory results — they reflect the same architecture succeeding in one regime and failing in another — but the field lacked the conceptual vocabulary to distinguish these regimes. This paper provides that vocabulary: transition cost, parameter alignment, and computational semantics as design criteria for sparse attention in the finetuning paradigm. Future work on trainable sparse attention should specify which paradigm it targets (train-from-scratch vs. pretrain-on-short, finetune-on-long) and evaluate accordingly.

A less obvious but potentially significant consequence is that this work redirects attention away from pattern design and toward training pipeline design. The paper's strongest result is not that InfLLM-V2's specific block-sparse pattern is optimal — the pattern is essentially inherited from InfLLM and NSA — but rather that a simple pattern, when deployed with zero architectural disruption during the finetuning transition, achieves near-parity with full attention while complex patterns that disrupt the transition fail catastrophically. This suggests that the field may be over-investing in attention pattern engineering and under-investing in understanding how architectural transitions affect pretrained representations. A research program that systematically characterizes the transition dynamics — which types of architectural changes cause recoverable vs. permanent disruption, how the disruption scales with model size and finetuning budget, whether gradual architectural annealing can reduce transition cost — could yield larger practical gains than further iterations on sparsity pattern design.

The block selection bottleneck analysis (Section 3.4) and the two-pass fused kernel with LSE approximation contribute a design pattern for hardware-efficient sparse attention that extends beyond InfLLM-V2. The insight that reductions over head groups and softmax normalization are non-commutative, and that a two-pass approach with coarse lse approximation can resolve this at acceptable accuracy cost, is a reusable algorithmic strategy. Any block-sparse attention method that shares selection masks across query heads — which includes most GQA-based approaches — will face the same I/O bottleneck and can apply the same solution. This is an incremental systems contribution but a practically important one, because the gap between theoretical savings from sparsity and realized wall-clock speedup is often where promising methods fail to deliver in deployment.

Follow-Up Research This Work Enables

Characterizing the accuracy-efficiency Pareto frontier across sparsity ratios. The paper evaluates a single operating point: 96 selected blocks, 6K visible tokens, yielding approximately 19% density at 32K. This is a reasonable choice but provides no guidance on how the accuracy-efficiency tradeoff behaves across different sparsity levels. A natural follow-up would sweep |I| across a wide range (e.g., 32, 48, 64, 96, 128, 192 blocks) at multiple sequence lengths (32K, 64K, 96K, 128K) on RULER and LongBench, measuring both accuracy and kernel-level latency. This would produce a Pareto curve showing exactly how much accuracy is sacrificed for each increment of speedup, enabling practitioners to select an operating point based on application-specific accuracy requirements and latency budgets. Critically, such a study would also reveal whether the accuracy degradation is smooth (suggesting the block selection is gracefully degrading) or exhibits phase transitions at particular sparsity thresholds (suggesting critical blocks are being dropped). The paper hints at such a threshold in the RULER MK3 result (62.00% at 6K visible vs. 92.00% for full attention), which suggests that some tasks require a minimum number of attended blocks and collapse when that threshold is crossed. A systematic sweep would map these thresholds across task types.

Isolating the mechanisms of NSA's failure through component-wise ablation. The paper bundles NSA's three architectural deviations — separate KV projections, multi-output gating, and learned compression — into a single comparison against InfLLM-V2, so it is impossible to determine which specific deviation causes the training instability. A minimal-component ablation study could disentangle these factors. The key experiment: start with InfLLM-V2 (shared KV, single output, parameter-free compression) and add one NSA feature at a time. Add only the multi-output gating (keep shared KV and parameter-free compression). Add only the learned compression MLP (keep shared KV and single output). Add only the separate KV projections (keep single output and parameter-free compression). Measure the training loss discontinuity (Figure 5-style curves) and final RULER performance for each variant. The hypothesis, based on the paper's architectural mismatch argument, is that multiple KV projections are the primary culprit because they introduce parameters that diverge during finetuning, creating a representational rift between the dense and sparse attention paths. If this hypothesis is confirmed, it implies that future sparse attention methods can safely incorporate multi-output gating or learned compression as long as they maintain a single set of KV projections. If the hypothesis is disproven (e.g., multi-output gating alone causes substantial degradation), it would suggest the problem is more fundamentally about the attention layer's type signature than about parameter count.

Testing whether InfLLM-V2 generalizes to code generation and long-document reasoning. All accuracy evaluations in the paper are on synthetic retrieval tasks (RULER), relatively short real-world QA (LongBench, where most sequences are well under 32K), and math/code reasoning (where the long context is self-generated chain-of-thought). These are important benchmarks but they do not capture the full diversity of long-context workloads that motivate the method. A strong follow-up would evaluate on: (a) repository-level code understanding tasks such as SWE-Bench (Jimenez et al., 2023), where the context includes entire codebase files and the model must localize and fix bugs — a task that combines retrieval, multi-hop reasoning, and precise attention to specific code regions; (b) long-document question answering on benchmarks like NarrativeQA or Qasper at their full document lengths (often 30K–100K tokens); (c) few-shot in-context learning with hundreds of examples, where the model must attend across many independent demonstrations. The hypothesis is that InfLLM-V2 will perform well on tasks where relevant information is spatially clustered (as in code, where function definitions and their uses are often proximal) but may struggle on tasks requiring uniform attention across the entire context (as in retrieving a single fact from a 500-page book). Testing this hypothesis would refine our understanding of which types of long-context reasoning are fragile under block-sparse attention and which are robust.

Developing and evaluating cheap difficulty estimators for the compute-optimal allocation framework. The paper's compute-optimal test-time scaling paper (referenced in the prior sections) demonstrates large efficiency gains from difficulty-conditioned allocation, but its difficulty estimation method costs 2048 samples per question — more than the inference budget itself. This paper's InfLLM-V2 architecture is orthogonal to that problem, but the techniques are complementary: a model equipped with InfLLM-V2 for efficient attention could use a separate lightweight module for difficulty estimation. A practical follow-up would train a small classifier — perhaps a single linear layer or a lightweight MLP on top of the model's final hidden state — to predict the difficulty bin (1–5) from the question text alone, using the 2048-sample PRM-based difficulty estimates as training labels. The evaluation would measure: (a) classification accuracy of the difficulty predictor against the oracle bins; (b) end-to-end accuracy when using the predicted difficulty to select the allocation strategy, compared to using oracle difficulty; (c) total compute cost including difficulty estimation, compared to a uniform best-of-N baseline. The key question is whether a cheap estimator can preserve enough of the compute-optimal gains to be practically worthwhile, or whether the estimation noise degrades the allocation so much that uniform strategies become preferable. The paper's finding that predicted difficulty bins closely track oracle bins (Figures 4, 8 in the prior work) is encouraging but does not address the cost of the prediction.

Combining InfLLM-V2's sparse attention with KV cache eviction for memory-constrained deployment. The paper focuses exclusively on speed (latency, throughput) and does not address memory consumption. In many real-world deployments, particularly on consumer GPUs or edge devices, GPU memory is the binding constraint — the KV cache for a 128K sequence of an 8B model can exceed available VRAM even if the attention computation is fast. KV cache eviction methods (H2O, SnapKV, and related work cited in Section 2.1) discard or compress KV pairs based on attention score heuristics, reducing memory footprint at the cost of potentially discarding tokens the model might later need. InfLLM-V2's block selection mechanism provides a natural integration point: if the model has already determined which blocks are important during prefilling, the KV pairs for non-selected blocks could be evicted from the cache entirely during decoding, rather than being retained and compressed. A practical study would measure peak memory usage and decoding latency when combining InfLLM-V2 with aggressive KV cache eviction, compared to either technique alone, on sequences at the model's maximum supported length. The hypothesis is that the combination enables larger effective context windows on memory-constrained hardware than either technique separately, because the sparse attention reduces the number of blocks that must be retained while the KV eviction reduces the per-block storage cost.

Training the model to dynamically adjust sparsity during autoregressive generation. InfLLM-V2 uses a fixed sparsity level (96 blocks, 6K tokens) throughout generation, determined at the start of each forward pass. But during autoregressive decoding for chain-of-thought reasoning or agent trajectories, the attention pattern evolves dynamically — early tokens may need broad context to plan, while later tokens may need only local context to refine or execute. A fixed sparsity level must be conservative enough to handle the most demanding generation steps, which limits average-case speedup. A natural extension would train the model to produce its own block selection budget: add a small learned component that, given the query and compressed key representation, outputs a scalar k(i) representing the number of top-k blocks to select for token i, and train this component with a loss that penalizes both attention cost (number of selected blocks) and task error. This is analogous to adaptive computation time in recurrent networks, applied to attention sparsity. The evaluation would compare the total attention FLOPs and task accuracy of the adaptive method against the fixed-|I| baseline on long reasoning benchmarks (MATH-500, AIME), where the variance in per-token attention needs is likely highest. The key risk is that learning a dynamic budget introduces instability — the model might learn to select minimal blocks and produce degenerate outputs — and the training procedure would need careful regularization.

Practical Applications and Downstream Use Cases

Cost-efficient long-context API serving with variable-length workloads. Cloud LLM APIs serve a mixture of short and long requests — from single-sentence queries to 100K-token document analyses. A deployment using full dense attention must provision compute for the worst-case quadratic cost, over-provisioning for short requests. A deployment using InfLLM-V2 with dense-sparse switching can serve short requests at full dense-attention speed (no sparse overhead, as validated by Table 4's short-sequence results) and automatically transition to sparse attention for long requests, with the speedup growing with sequence length. Based on Figure 7's end-to-end measurements: at 32K (RULER length), the prefilling speedup on A100 is 1.09× — modest but already positive. At 128K, the prefilling speedup is 1.99×. In a workload where 80% of requests are under 4K (processed with dense attention at zero overhead) and 20% are 32K–128K (processed with sparse attention at 1.1×–2.0× prefilling speedup and 1.2×–1.8× decoding speedup), the overall throughput improvement is a weighted average that depends on the exact length distribution but will be strictly positive — there is no regime where InfLLM-V2 is slower than dense attention. This property — "never slower, sometimes much faster" — is what makes the method immediately deployable without per-workload tuning.

On-device or consumer-GPU deployment for personal AI assistants with long-term memory. A personal AI assistant that maintains context across weeks of conversation, referencing past interactions while responding to current queries, faces a severe memory bottleneck on consumer hardware (e.g., NVIDIA 4090 with 24GB VRAM). An 8B model with full attention at 128K context requires a KV cache of approximately 128K × 32 layers × 2 KV heads × 128 dimensions × 2 (K and V) × 2 bytes (FP16) ≈ 4GB for the KV cache alone, plus model weights, activations, and overhead. InfLLM-V2 with 6K visible tokens reduces the attention FLOPs by approximately 3.1× at the kernel level at 128K on the 4090 (Figure 6), and if combined with KV cache eviction for non-selected blocks, could reduce the memory footprint proportionally. A concrete deployment scenario: a 4090-equipped desktop running a locally-hosted assistant that processes a user's entire email and chat history (potentially 100K+ tokens) while responding with latency under 1 second per token. Without sparse attention, this is infeasible on 24GB VRAM; with InfLLM-V2, the attention component of latency drops from being the bottleneck to being competitive with FFN time (Figure 7), and the memory footprint becomes manageable with aggressive eviction. The paper does not provide the memory measurements needed to validate this use case fully, but the speed improvements alone shift the bottleneck from attention to FFN, making the deployment more tractable for further optimization.

Training data generation for self-improving long-context models. Self-improvement pipelines (such as STaR, ReST^EM, or rejection sampling fine-tuning) require generating high-quality outputs from a model on a large corpus of inputs, which are then used to fine-tune the model itself. For long-context tasks — generating summaries of long documents, producing reasoning chains for complex problems, executing multi-step agent trajectories — the generation phase dominates the total compute budget. InfLLM-V2 enables generation with 2–3× throughput improvement at sequence lengths above 64K (Figure 7 on 4090) while maintaining output quality within 1–2% of full attention (Table 3 reasoning results: 42.66% vs. 42.79%; Table 2 LongBench: 42.54 vs. 42.30). This directly translates to 2–3× faster self-improvement iterations, enabling either faster experimentation or larger-scale data generation within a fixed compute budget. The caveat is that the quality of generated data must be carefully monitored — the paper's RULER subtask breakdown shows that some task types (MK3, QA2) degrade meaningfully under sparsity, and if these task types are overrepresented in the generation corpus, the self-improvement loop could amplify the errors rather than correcting them. A practical deployment would apply InfLLM-V2 with a verification step that uses the Dense mode for any generated outputs flagged as low-confidence.