InfLLM-V2: Dense-Sparse Switchable Attention for Seamless Short-to-Long Adaptation
Authors: Weilin Zhao
Zihan Zhou
Zhou su
Chaojun Xiao
Yuxuan Li
Yanghao Li
Yudi Zhang
Weilun Zhao
Zhen Li
Yuxiang Huang
Ao Sun
Xu Han
Zhiyuan Liu
Venue: ICLR 2026 Poster (2026)
Keywords: Sparse attention
Long-context
Efficient algorithm
Primary area: foundation or frontier models, including LLMs

Abstract: Long-sequence processing is a critical capability for modern large language models. However, the self-attention mechanism in the standard Transformer architecture faces severe computational and memory bottlenecks when processing long sequences. While trainable sparse attention methods offer a promising solution, existing approaches such as NSA introduce excessive extra parameters and disrupt the conventional pretrain-on-short, finetune-on-long workflow, resulting in slow convergence and difficulty in acceleration. To overcome these limitations, we introduce dense-sparse switchable attention framework, termed as InfLLM-V2. InfLLM-V2 is a trainable sparse attention that seamlessly adapts models from short to long sequences. Specifically, InfLLM-V2 reuses dense attention parameters through parameter-free architecture modification, maintaining consistency between short and long sequence processing. Additionally, 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. To achieve practical acceleration, we further introduce an efficient implementation of InfLLM-V2 that significantly reduces the computational overhead. Our experiments on long-context understanding and chain-of-thought reasoning demonstrate that InfLLM-V2 is 4$\times$ faster than dense attention while retaining 98.1\% and 99.7\% of the performance, respectively. Based on the InfLLM-V2 framework, we have trained and open-sourced MiniCPM4.1 (https://huggingface.co/openbmb/MiniCPM4.1-8B), a hybrid reasoning model, providing a reproducible implementation for the research community.

Published as a conference paper at ICLR 2026
INFLLM-V2: DENSE-SPARSE SWITCHABLE ATTEN- TION FOR SEAMLESS SHORT-TO-LONG ADAPTATION
Weilin Zhao', Zihan Zhou?, Zhou Su², Chaojun Xiao' Yuxuan Li', Yanghao Li', Yudi Zhang', Weilun Zhao', Zhen Li², Yuxiang Huang', Ao Sun?, Xu Han Zhiyuan Liu TTsinguua University 'OpenBMB
Harrbi Institute of Technology zw1 arl13emaall. 1s. ts tsinghua. edu.en {xcj, han-xu, liuzy Hiuzy]etsinghuut edu.en
ABSTRACT
Long-sequence processing is critical capability for modern large language mod- els. However, the self-attention mechanism in the standard Transformer architec- faces severe computational and memory bottlenecks when processing long sequences. While trainable sparse attention methods offer a promising solution, existing approaches such as NSA introduce excessive extra parameters and dis- rupt the conventional pretrain-on-se on-short, finetune-on-lond workflow, resulting in slow convergence and difficulty in acceleration. To overcome these limitations, we introduce dense-sparse switchable attention framework, termed as InfL.LM- V2. InfLLM-V2 is trainable sparse attention that seamlessly adapts models from short to long sequences. Specifically, InfLLM-V2 V2 reuses dense attention param- eters through parameter-free architecture modification, maintaining consistency between short and long sequence processing. Additionally, InfLLM-V2 ensures computational efficiency across all sequence lengths, by using dense attention for short inputs and smoothly transitioning sparse attention for long sequences. To achieve practical acceleration, we further introduce an efficient implementation InfLLM-V2 V2 that significantly reduces the computational overhead. Our ex- periments on long-comtex: understanding and chain-of ofothh reasoning demon- strate that Infl.LM- V2 4x faster than dense attention while retaining 98.1 % and 99.7% of the performance, respectively. Based on the InfL.LM- V2 framework, we have trained and ing reproducible implementation for the research community.
MiniCPM4. P hybrid reasoning model, provid-
INTRODUCTION
With the rapid development of large language models (LLMs) (Brown et al. 2020 Bommasani 2021 Han et al. 2021 OpenAl 2023), the demand for long-sequence processing capabilities become increasingly critical. From long-input scenarios such as deep research (Zheng et al., Xu & Peng 2025). chatbots with long-term memory, and software issue resolution (Jimenez 2023 Yang et al. 2025). to long-output tasks including complex reasoning (OpenAl et al.
et al.
has
2025
ct al.
2024
DeepSeek et al. 2025) and LL.M-driven agents (Wang al. 2024 model's capability to
understand and generate long sequences directly determines its performance real-world applica- tions. However, the self-attention mechanism in the existing Transformer (Vaswani al. 2017) architecture faces severe computational and memory bottlenecks when processing long sequences.
To address the challenge of processing long sequences, efforts have been devoted exploring sparse attention mechanisms (Beltagy et al. 2020 Zaheer et al.] 2020 TTy et al. 2022). which restrict each token within the context attend to only a subset of tokens related to that token. Early research in this area focuses on the sriiiig-frree setting, leveraging the sparsity naturally occurring in self- attention mechanisms accelerate inference (Xiao et al., 2024a bb] Jiang et al., 2024). However, the training-free setting introduces fundamental trade-off between sparsity and model performance. To avoid significant performance degradation, the degree sparsity that can be applied often limited, which in turn restricts the potential efficiency gains.
"Corresponding Authors.
face http///ca..c.coo///////////////a CPM. 1

---

Published as a conference paper at ICLR 2026
NSA
Vanilla
InfLLM-V2
3x KV Projection
Redundant
Single KV
Reusable
Shared KV
Attention
Expensive
Single Attention
Portable
Unified Attention
Gated Output
Mismatch
Single Output
Aligned
Single Output
Figure The comparison of Vanilla Full Attention, NSA Yuan et al.
2025 and our InfLLM--V V2.
Given the limitations of training- free attention mechanisms, trainable sparse attention mechanisms
Gao et al. (024). Among them, the natively trainable sparse attention (NSA) (Yuan et al. 2025 method adopts the widely- used block-sparse attention (Child et al. 2019) structure, designing three different sparse attention modules and developing corresponding CUDA kemelss to accelerate model computation. Despite its effectiveness, we find misalignment between the sparse architecture of NSA and the stan- dard pretrain-on-short, finetune-on-long workflow. A widely used way build long L.MMs is to pretrain on short sequences and finetune on long sequences. The NSA creates an architectural mismatch with vanilla full attention, it introduces three sets of key-value parameters and three attention modules, forcing the model to abruptly switch from single-output attention to a multi- output attention architecture. As shown in Section 4 this mismatch hinders smooth adaptation, erases what the model has already learned, and introduces efficiency bottleneck for short sequences.
have garnered increasing attention from researchers (Lu et al.) 2025
To address all the above issues, we propose dense- sparse switchable attention framework (InfLLM- V2). I2 is built on InfLLM (Xiao et al. 2024a), training-free block-sparse attention mechanism, and introduces three core innovations:
1. Seamless Short-to-Long Adaptation: As depicted in Figure I different from NSA, which re- quires additional parameters and multiple attention modules, IniLLM- V2 seamlessly transitions from dense to sparse attention by directly reusing existing dense attention parameters. This de- sign naturally aligns with the standard architectural mismatches and minimizing loss fluctuations.
finetune-on-long workflow, eliminating
2. Efficiency for Both Short and Long Sequences: Because the transition from dense to sparse attention in InfLLM-V2 requires no additional parameters and introduces minimal distributional shifts, the model preserves its strong performance on short texts and easily switch back to dense attention for short sequence efficiency.
3. Accelerated Block Selection Mechanism:
The block selection step before sparse attention in- herently undermines the efficiency gains of the sparse attention itself. We propose hardware- awared efficient implementation, effectively removing the bottleneck and unlocking the full po- tential of sparse attention.
We evaluate our method on long-context understanding and long chain-offthouggh (CoT) generation benchmarks. Our InfLLM- V2 is 4x faster than dense attention while maintaining 98. 18%1% and 99.7% of the original performance in wtlhttss:// /tap::///ttttit wit.bb..tttttgggggoo com Fatom/OpenBBete in mv mV2 mww__mdd__iaaaeeeeee_ee_
these tasks, respectively. We release all associated implementations
2 RELATED WORK
As the demand for LLMs to understand and generate long sequences continues to grow, rescarch on improving attention efficiency has garnered increasing attention (Tay et al., 2022, Sun et al.) 2025 Zhang al. 2025a). In this section, we discuss the sparse attention paradigm from two perspectives: training-free and trainable sparse attention approaches.
2.1 TRAINING-FREE SPARSE ATTENTION
Training-free free sparse attention approaches aim to utilize the intrinsic sparsity of attention layers. These methods enable LLMs trained with dense attention to perform sparse attention between each token and a small subset of relevant contexts. Based on the selection strategy for relevant contexts, these algorithms can be categorized into predefined sparse patterns and dynamic sparse patterns.
Predefined Sparse Patterns. Sparse attention with a predefined pattern employs manually de- fined heuristic rules determine which contextual tokens should be selected for attention compu-
2

---

Published as a conference paper at ICLR 2026
tation (Xiao al. 2024b Han et al. 2024 Child et al. 2019 Zaheer al. 2020 Beltagy al, 2020 Xiao et al. 2025). For instance, sliding window attention restricts each token to interact only with neighboring tokens (Beltagy et al., 2020). Building upon sliding windows, some works select special tokens such as initial tokens or segment separators, requiring all tokens to attend to these special tokens (Xiao et al., 2024b) Chen al. 2024, Child et al., 2019). Furthermore, some works combine multiple predefined attention patterns (Zaheer et al.) 2020; [eltagy et al., (020). These approaches typically rely on human observations to formulate heuristic rules for selecting relevant contexts.
Dynamic Sparse Patterns. Dynamic sparse patterns incorporate the semantic information of query tokens into the context selection process by computing the relevance between query tokens and candidate contexts. Early works primarily perform similarity computation at the token level (Ki- taev et al., 2020, Roy et al. [202]; Wang et al. 2020). As sequence lengths increase, block sparse methods have gained widespread adoption, which partition contexts into contiguous block units and
perform relevance computation and context selection at the block granularity (Xiao et al. 2024a
Jiang et al. 2024 Xu et al., 2025 Tang et al. 2024, Zhang et al. 2025b; Lai et al. (025). Further- more, research attention sparsity has inspired the development of key-value (KV) eviction and compression methods, which reduce memory consumption by discarding or compressing KV caches
with low attention probabilities (Zhang et al. 2023 Li al. 2024 Huang et al. 2024 2025).
Training-free methods, while focusing on improving the inference efficiency of dense attention mod- els, are often constrained by insufficient sparsity levels in order to avoid severe performance degra- dation and finally suffer from limited acceleration benefits.
2.2 TRAINABLE SPARSE ATTENTION
To further enhance efficiency for long sequence processing, researchers incorporate sparse attention into the model training phase. SecrAttention (Gao et al. 2024) employs a self-distillation post- training algorithm to train router that selects relevant contexts for query blocks. MoBA (Lu al, 2025) employs a block sparse attention structure during the short-to-long adaptation phase, training routers between query blocks and KV blocks for context selection. These methods partition query tokens into blocks and can only accelerate the prefilling phase. NSA (Yuan et al., 2025) designs three attention components for token-level sparsity, effectively accelerating both prefilling and decoding processes. However, NSA introduces substantial additional parameters, making it unsuitable for ef- ficient short-to-long adaptation and imposing significant computational overhead on short-sequence processing. In this paper, we focus on proposing sparse attention mechanism that effectively and efficiently processes both short and long sequences, supporting both prefilling and decoding.
METHOD
3.1 BACKGROUND
Grouped-Query Attention. Attention mechanisms enable models to selectively focus on rel- evant parts of the input sequence. Among various attention variants, grouped-query attention (GQA) (Ainslie et al. 2023) has emerged as a popular method that strikes balance between model performance and computational efficiency. Given an input sequence of hidden states X E Raxd where n is the sequence length and d is the model dimension, GQA computes the queries (Q), keys (K), and values (V) via linear projections as Q = XW Q, K = XWK, V XWv. The pro- jection matrices have the shapes We E dx and W K, Wy E IR (hvudd), with the head dimension d,, These tensors are then reshaped to form hy query heads (Q:122 hkve KV heads (Kj, V, V with each head having the shape n x dy. The query heads are partitioned by group size G = h The attention scores S; and the attention output O for the i-th query head are computed by attending to its corresponding KV heads with the index [(i 1)/G] +
Q.K
S.
Softmax
O, S,V
Vdn
The final output is obtained by concatenating the attention outputs and projecting them through final linear layer Wo TR (h,dd))xd Attention X) Concat( O ..,..., Os,) Wo o.
3

---

Published as a conference paper at ICLR 2026
V
Kile Ville
Kuin
Virin
Projecsion
Comoressed Attontion
octed Mitentinn
Skding Asartion
Larger
Q
Attention Outsut
Afteetton Outbut
Aftention Outbut
Gatn
NSA In-
K.
Switch (sealen)
Atorttco
Seleceed Skting
Denow Aftention
Q
Attoriion Score Only [
Top hit More Local
Spana Alterior Output
Deraa Altwnton Output
A Exmm Parsmeter
CUUDA Comp to Kemne
L GueryKKyy'auu Tokens
Attertion Citpue
L
Vacced Token
Athenton Score
Figure 2:
The overview of NSA and InfLLM-V2. InfLLM-V2 uses shared KV for both Sparse Attention and Dense Attention. InfLLM-V2 fuses Selected Attention and Sliding Attention and eliminates the output of Compressed Attention. InfL.LM-V2 introduces no extra parameters.
NSA. NSA Yuan et al. 2025) is an enhancement of GQA designed for efficiency on long se- quences. The key insight that for long sequences, e.g., when 32k, the attention score matrix exhibits strong sparsity. This allows for approximating the attention matrix by ignoring negli- gible values, leading faster computation. As illustrated in Figure 2 NSA utilizes three distinct modules and combines them using a gating module. Based on the observation that adjacent atten- tion scores are similar (Jiang al.) 2024). NSA splits the sequences into blocks of size B. First, Compressed Attention employs compressed representation of the KV tensors to reduce the com- putational complexity. Second, Selected Attention leverages the attention scores from compressed attention to compute only the blocks with high attention scores. Finally, Sliding Attention is used to focus on local contextual information within the sequence. For these three attention modes, they introduce three sets of KV projection matrices: W W. Wale Wsle K W This final output can be mathematically represented Output g Oemp + ""0"""nn where Osmp, Oslc, and Owin are the outputs the three respective modules, and the gate scores gemp, gelc, and win derived from the input features X via an MLP and a sigmoid activation. They also train an MLP module for compressing the KV tensors. The three distinct KV projections, combined with an additional MLP and gating module, result in highly complex architecture. This complexity, in turn, makes the model poorly suited for training from scratch on short-sequence data and also complicates the process of converting pretrained dense models to sparse ones.
3.2 OVERALL FRAMEWORK
We propose InfLLM-V2, a more concise framework with zero extra parameters that more closely aligns dense and sparse attention patterns.
Shared Value Projection. We find that using three separate sets of KV projection parameters in NSA (Yuan al., 2025) is unnecessary, which not only complicates the adaptation from short to long sequences but also significantly slows down computation for short sequences. Therefore, we propose using a single shared set of projection parameters, W K and Wv. initialized with the pretrained dense attention parameters and used for finetuning long sequences.
Aligned Computation. In addition to ensuring that sparse and dense attention share the same pa- rameters, their computational processes must also be closely aligned. In NSA, the three attention modules all generate outputs that are aggregated by an extra gating module. This forces the computa- tion of all three modules we take a union of the two sparse patterns in Selected Attention and Sliding Attention and eliminate the output of Compressed Attention. forming a unified Sparse Attention module. Specifically, the original Selected Attention module identifies important token blocks based on the attention scores from the Compressed Attention module, Semp, For query token with index i, located in the block b, = i-1 + 1, attention is always granted to fixed set of initial blocks and set of local blocks:
for short sequences, leading substantial overhead. To mitigate this,
Immit (1,2.... Ninit: Irnal())
{b: Nivcal + 1,..., by 1, 1,b:).
(2)
4

---

Published as a conference paper at ICLR 2026
The top-k selection is then applied to Srmp over the set of remaining blocks, denoted as Tupk (1). The complete set of attended block indices for this query token is the union of these three sets:
I()) = Tint U Tocal (7) U Zook (1).
(3)
If we denote the set of token indices in the j-th block as T3 = {jB+1, + +...., (j 1)B), the selected attention allows a token in the block b, to attend to the union of blocks UJEzan Tj. The Sliding Attention, on the other hand, allows the i-th token to attend to a range {i - w w.......+ of window size w. Since the local blocks in Selected Atlention and the window in Sliding Attention create overlapping, we merge them by expanding the number of local blocks within our unified Sparse Attention to strictly cover the region of the Sliding Attention, that is, Noccl 2 w + 1, as illustrated in Figure
NSA Selected Attention Mask (init=1, topk=2, local=1)
Furthermore, we eliminate the output of the Compressed Attention module, only retaining its attention scores Semp for block selection in Sparse Attention. This single-output de-
NSA Sliding Attention Mask (w=6)
sign more closely mirrors dense attention and aids the training of the sparse attention model. InfLLM-V2 can thus dynamically switch be- tween dense and sparse attention patterns based on the input sequence length.
Inft LMfLLM-2 Sparse Attention Mask (Ini=11, topk=2, local-3)
Simplified and Efficient Compression Mod- ule. Since we eliminate the output of the Compression Attention, using MLP for token compression would not receive gradients. We replace it with a more intuitive parameter-free pooling function, which will be detailed in Section 3.3 Additionally, computing the attention scores Simp introduces non- non-negligible overhead, and we will reduce this overhead in Section 3.4
The illustration of the union of Selected Attention and Sliding Attention.
Figure 3:
3.3 BLOOK REPRESENTATION
Simply compressing a long sequence with large block size B in 1-stage can lead to significant loss of granular information (Yuan ett al. 2025). To address this, we implement a 3-stage, coarse-grained to fine-grained com- pression process, as shown in Figure 4 In the first stage, we process the input key sequence K produce an intermediate and coarse-grained representation k By denoting the initial compression block size (c, and the stride as sC,, we achieve this by applying a mean- pooling operation over sequential blocks:
K
Attention Score
Scmp
1-stage toker-leve S-olage group-inve
K,
Q
Head
Gooup
G+3
gahardd
Somp
KC, Mean s.s sha tle, ).
(4)
Figure 4: The illustration of the 3-stage group- level compression, compared with the I-stage token-level compression.
Then, we compute the attention scores S be- tween the query Q and K
(5) In the second stage, we employ block-wise sparse attention rather than token-level approaches for the efficiency Sparse Attention. In model utilizing GQA, we can achieve this by forcing the block selection pattern across all heads within group to be the same. We conduct summation within the head group to get the shared importance score shhred.
SC1 Softmax(((K<))T))
S shared
SCi()).
(6)
In the third stage, apply max-pooling operation, which preserve the most salient features. The aggregated score gemp are defined as follows and used for the Sparse Attention:
cmp
Max (Sshared b-2rt-s+1)
5

---

Published as a conference paper at ICLR 2026
Algorithm Computation of Sshared (Suppose luke = without loss of generality.)
Require: Q E Rnxxxds, KC* E RR(n/s0,)x dd K (nn/acy))da in HBM. Block sizes BR, Bk. Divide Q into T. [n/ B,] blocks QI, Qr of size Bg
da each.
Divide k into Tı = [n/ac,/ / Bk] blocks K ..... k size Bg x dz each.
Divide K into TN [w/sca/ [x]] blocks K1 ..... T2
C2 of size Bx X da each.
Divide S Sahored into T. x Ti blocks of size Bg x Bg each.
for i = ,, (parallel) do
Load Q, from HBM to on-chip SRAM.
On chip, initialize online-softmax related statistic log-sum-exp Ise. for 1,...,77 T2 (sequential) do
First pass (Coarse-grained)
Load K2 from HBM to on-chip SRAM.
On chip, compute attention scores S E X as in Eq. and update Ise.
for j Ti (sequential) do
> Second pass (Fine-grainede
Load K from HBM to on-chip SRAM.
On chip, compute attention scores S E &ROBBBBR in Eq. and normalize it using Ise.
On chip, compute the final block S E R BBBBB, by summing S over the head group.
Write the block Schard to its corresponding position HBM.
return the output Sthared
In our method, we set lc. = B sc, B 1 = 5, and s 4 so that it can achieve the same compression ratio as I-stage compression of block size B. Intuitively, we compute the sparse scores of the entire block based on several sliding sub-blocks within the block.
3.4 EFFICIENT IMPLEMENTATION
For efficient Sparse Attention, we follow the techniques in NSA Yuan al. 2025 to set the group size G of GQA to 16, a configuration well-suited for block sparse attention. More details can be found in Appendix [ However, our profiling reveals that the computation of the compression score, Samp, introduces a significant performance bottleneck. A primary source of this slowdown the substantial I/O required to store the first-stage attention scores SC: into the slow GPU HBM. The amount of data that needs to be written is h₂ //c////.......././// where n is the full sequence length. Given that sc, n,, materializing the full attention score matrix GPU HBM incurs a prohibitive cost.
Drawing inspiration from FlashAttention (Dao 2024), we aim to minimize this I/O by ensuring the attention scores remain within the fast GPU SRAM as much as possible. Our approach, Fused Head Group Summation, is to fuse the summation over the head group, required for the second- stage compression, directly into the SRAM-based computation loop of FlashAttention, After that, we can only store the reduced attention scores shaared into GPU HBM, whose size is hy?? /(sc, G).
Another challenge arises from the fact that summing over the head group dimension and performing the collne-softmax (Dao) 2024) along the sequence dimension are not commutative operations. This conflict prevents a straightforward fusion. To overcome this, we implement two-pass approach. In the first pass, we compute the log-sum-exp (Ise) term required for the softmax normalization within the SRAM. In the second pass, we leverage the Lse to calculate the final attention scores, perform the summation across the head group within the SRAM, and write the reduced scores to the HBM. The trade-off of this two-pass method is that it doubles the computational workload. Therefore, propose LSE Approximation to approximate the Ise computation by using coarser-grained attention score SC2, Following Eq. (4) and Eq. 5). we change them to
kCa Meaa (es,,
SC2 Softmax (?T))
(8)
By setting SC2- = 4sc, and Ic, = 4lc, , the computational overhead was reduced from 2x to 1.25% We summarize the procedure for computing shhared in Algorithm To further reduce memory I/O, the max max-pooling and top-k operations related to Semp could also be fused into the kernel; however, leave this implementation for future work.

---

Published as a conference paper at ICLR 2026
FullActt (Long Fine-tuningg)
3.5
InfLLM-V2 (Long Fine-tuning)
NSA (Long Fine-tuning)
FullAatt (Short Pre-tralning)
10592
2.5
1:2
1.5
1.0
100
200
300
400
500
600
Steps
Figure 5: The training loss of models. We only show the last few iterations of the short pretraining. 4
EXPERIMENT
We evaluate InfLLM-V2 on tasks ranging from short to long contexts, and demonstrate its efficiency.
4.1 EXPERIMENT SETUP
Pretraining Setup. We first use full attention to pretrain a model on short-sequence data, marked as SHORT. We employ standard GQA (Ainslie et al. 2023) model backbone with 8B parameters, with the hidden size = 4096, the number of heads h2 = 32, hgg = 2, and the head dimension dy, = 128. The pretraining dataset consists of 8T tokens of 4k-length sequences, primarily comprising Feb-Wee-Edu (Penedo et al. 2024) and Stack-v2 (Lozhkov et al.] 2024]. We set 8M tokens per batch, and use WSD learning rate scheduler (Hu al. (2224) with 2000 warmup steps to initial learning rate of 7,5e-3, and 27000 decay steps to the final learning rate of 3e-4.
Long-Context Adaptation. When transitioning to long-context finetuning, we switch to ) we set the compression block size lc 32, stride sc, 16. and attention block size B 64. For our efficient block selection
INFLLM-V2 (SPARSE). Following NSA (Yuan et al. 2025).
implementation in Section [3.] we additionally set the LSE Approximation block size lc. = 128
and stride SC2 = 64. We set the selected block count [ = 96 (including Taml 1, Iupp = 63, and |Zaall 32) for both training and inference. Therefore, the total number of visible tokens is [Z] B 6k. We conduct long-sequence finetuning on the pretrained model using 5B tokens, with an initial learning rate of 3e-4 and linear decay to 2.75e-4. The training batches contain sequences from four length intervals: 0-4k, 4-12k, 12-24k, and 24-32k, with token counts in 1:1:1:1 ratio.
Baselines. We finetune a baseline model with full attention, marked as FUL.LATTN, using the same training configuration as INFLLM-V2 V2 (SPARSE). We then apply several typical training-free FULLATTN as baselines, including IniLLM (Xiao et al. 2024a) and MInference (Tiang et al. 2024). In addition, we present the results of SHORT with YaRN (Peng et al. 2023 to extend the context window size. In terms of trainable sparse attention, with NSA Yuan et al. 2025). By using the same training settings as in INFLLM- V (SPARSE), we finetune our pretrained model into an NSA version. We initialize NSA's three sets of KV parameters by replicating the original KV parameters in dense attention. As NSA does not publish their code, we adopt an open-source Triton implementation of NSA for experiments
sparse attention methods
compare
Table 1: Task Performance on RULER. Best results in sparse attention are bolded.
Method
SG2
SG3
MK1
MK2
CWE
FWE
QAI
QA2
Avg.
FULLATEN
100.00
100.00
100.00
96.00
94.00
92.00
82.00
98.50
93.20
44,40
91.33
48.00
56.00
84.26
SHORT+YARN
50.00
46.00
6.00
21.40
87.33
26.00
40.63
INFLLM
98.00
6.00
4,00
10.00
10.00
10:00
9.00
7.50
70.00
16.00
80.67
18.00
24.00
27.94
MINFERENCE
100.00
100.00
100.00
76.00
36.00
46.00
79.50
93.50
88.00
64.20
92.67
32,00
44,00
73.22
NSA
100.00
85.00
54.00
35.00
30.00
59.00
34,40
86.00
56.00
34.00
59.92
INELL M-V2 INFLLM--V2 (SPAASE))
LSE Appros
100.00
100.00
100.00
82.00
82.67
w/o LSE Approx
100.00
100.00
100.00
92.00
80.00
64.00
98.50
95.50
98.00
47,80
81.33
70.00
40.00
82.09
INFLL M-V2 (DENSE)
100.00
100.00
100.00
94.00
98.00
98.00
99.00
98.00
98.40
52,80
90.00
76.00
44.00
88.12
7

---

Published as a conference paper at ICLR 2026
Table 2: Task Performance on LongBench and LongPPL. Best results in sparse attention are bolded.
Benchmark
FULLATIIN
SHRRT-YARN
INFLLM
MINFERINCE
NSA
INFLLM-V2 (SPARSE)
INrLLM-V2 (DENSE)
LongBench 42.30
37,86
32.30
41.55
37.30
42.54
42.49
LongPPL 1
2.06
5.28
12.01
2.62
4.24
2.12
2.00
For all the above sparse attention methods, we maintain the same sparsity level to ensure a fair com- parison. We provide the training curve for trainable methods in Figure NSA causes a disruption in the loss, while INFLLM-V2 is closer to FULLATTN.
4.2 TASK PERFORMANCE
In this section, we evaluate InfLLM-V2 and other baselines across various tasks. Notably, while the original NSA paper demonstrates performance comparable to full attention when training on long sequences from scratch, NSA fails to achieve satisfactory results in short-to-long adaptation settings. This indicates that the substantial parameter overhead introduced by NSA renders it unsuitable for the conventional "pretraining ineeraining-on-shorr finetuning -one-oonooooog paradigm.
Long-Context Understanding. To evaluate InfLLM-V2's performance on long-input tasks, we compare InfLLM-V2 and different baselines on RULER (Hsieh et al. 2024), LongBench (Bai et al.) 2024 and LongPPL (Fang et al., 2025). RULER is a synthetic dataset with a configurable average length. LongBench is bilingual benchmark for long-conted understanding. Compared to RULER, LongBench primarily built from existing. real-world datasets. LongPPL a perplexity evalua- tion benchmark for long sequences. The experimental results of RULER when the length is 32k are shown in [ The results on LongBench and LongPPI. are shown Table Please refer to Ap- pendix for detailed performance of the sub-tasks in LongBench. From the results, that: 1) INFLLM-V2 achieves the best performance compared to other sparse methods, with its results closely matching the strong, FULLATTN baseline. Alternative approaches, whether applying training-free sparsity or training-basse based sparsity, result in a substantial drop in performance. 2) Com- pared to NSA, INFLLM- V can achieve performance improvements through minimal finetuning on long-sequences. Although NSA has low training loss, its high perplexity on the LongPPL evalu- ations indicates that NSA has adequately learned long-range dependencies. 3) A unique advan- tage of INFLLM-VE V2 is the flexibility to seamlessly switch between dense mode and sparse mode. This not only provides an option for dense computation but can also lead further performance im- provement, surpassing even the full attention baseline. 4) Furthermore, the INFI.M-- V2 (SPARSE) variant with LSE Approximation does not lose any performance. Given its comparable performance and efficiency, we adopt the LSE Approximation setting for all subsequent tasks by default.
observe
Long Reasoning. To eval- uate the performance of InfLLM-V2 in long-output scenarios, we compared several major Long Rea- soning tasks, including MATH-500 (Hendrycks et al., 20216), AIME (MAA), and LiveCodeBench (LCB) (Jain et al.
Table 3: Task Performance on Long Reasoning Tasks.
Method
MATH-500
AIME
AIME
LCB v5
LCB
Avg. f
FULLATTN
86.00
37.50
30.63
30.67
29.14
42.79
23.54
25.14
37.28
In (Sparse)
87.80
38.33
29.38
29.94
27.83
42.66
86.40
36.67
23.33
29.94
26.29
40.53
2025). We finetune In and baselines on OpenoMannessooning (Moshkov et al.) (025) and OpenCodeReasoning (Ahmad et al. 2025). As InfLLM and MInference primarily accelerate long-input processing, we exclude them from this long-outpu evaluation. The experimental results are shown in Table 3 The results show that InfLLM-V2 attains performance on par with full attention, confirming its effectiveness for long-output scenarios.
General Tasks. We verify that the InfLLM-V'' V2 architecture can freely switch back to Dense mode without performance degradation on short-sequence tasks after long-sequence fine-tuning. Zero- shot evaluations on MMLU (Hendrycks et al. 2021a), MMLU-Redux (Gema et al.] 2025). CE- val (Huang et al., 2023), MATH-500 (Hendrycks et al. 2021b), HumanEval (Chen et al. 202T). MBPP (Austin et al. (22T)) and BBH (Suzgun et al. 2023) are shown in Table Experimental
results show that LfLLM--2 achieves performance comparable to full attention.
4.3 EFFICIENCY
We first evaluate the efficiency of our kernel implementation on NVIDIA A100 and NVIDIA 4090. We evaluate L inference efficiency the batch=1 setting. We select FlashAttention-
8

---

Published as a conference paper at ICLR 2026
Table 4: Task Performance on General Tasks.
Merhad
MMLU
MML M-luu-Redux
CEval
MATH-500
HumanEvvll
MRPP
BBH
Avg-
SHORT
72.73
72.71
76.17
51.40
70.73
73.49
51.90
67.73
FULLATTN
73.38
70.24
78.11
54.60
71.34
75.10
49.13
67.41
NSA
39
44.40
60.33
(Dense)
71.29
77.70
54
73.17
47.09
66.76
A1 00, 32q
A100, seglen=64k
A100, seglen=96k
A100, seqien=
800
200
0.7x
1,1x
1..5x
60
150
300
1.7x
40
3.0x
1x1x
1:33
5.22
200
2 6x
X.6X
2 0.
2.5x
2 SK
2.4
5x
31x
20
50-
100
1k
2k
4k
6k
1k
2k
4k
6k
1k
2k
4k
2k
4k
6k
4090, seqlen-32k
4090, seqlen
4090,
seqlen-966
4090, seqlen 128k
b.9x
600
1000
1.5x!
400
20-
40
3 5,9
1.8k
100
6,6x
4.7x
8x
2.5x
200
3:3x
6. 00
P.3x
3.8x
500
250
35x
1.0x
5.0%
Number of Visiblle
Takens
Number of Visible Tokens
Number of Visible Takens
Number Visible TB kens
FullAttn (FlashAttention)
TZI
Sparse Attention
Figure 6: Speed of the kernels NVIDIA A100 and NVIDIA 4090. "Number of Visible Tokens"
means the number of key/value tokens each query token can attend to.
2 (Dao 2024 implementation for full attention. For a fair efficiency comparison with NSA, we ignore its sliding attention component, and compare solely on the compression and sparse attention parts by selecting the number of selected blocks is 16, InfLLM-V2 achieves up to 7.4x over FlashAttention on A100 and 9.3x on 4090. In contrast, NSA's speedup is limited to 3.5 in the same setting. The breakdown of the execution time shows that the overhead from the Block Selection stage greatly optimized by our efficient implementation in Section 3.4 We further conduct ablation study the Block Selection, as shown Table3] which shows the effectiveness of our proposed LSE Approximation.
equal number of blocks [Z]. Experiment results are shown in Figure 6 When
Table 5: Ablation study Block Selection efficiency, with and without LSE Approximation. All measurements are in time (ms), and the number of selected blocks is set to 16.
Device
A
4090
Soquence Length
32k
64k
96k
128k
32k
64k
96k
128k
w/o LSB Approximation
4.67
18.20
42.46
75.36
4.89
19.95
46,51
83.26
w/ Approximation
3,93
32.44
56.59
3.70
33.16
The end-to-end inference speed (with [Z] = 96 and W4A16 quantization (Frantar et al. 2025 is shown in Figure 7 InfLLM-V2 achieve 2.13x prefilling speedup and 2.32 decoding speedup. Since InfLLM-VV V-V2 does not accelerate the ratio be achieved by incorporating FFN-specific acceleration techniques in future work.
Network (FFN) layers, a higher speedup
A100- Prefilling
A100 Decoding
4090 Prefflling
4090 Decoding
40
828
TPOT (me)
20
.55K
23x
TPOT (196)
32x
TTFT
20
TIFT
.9k
20
F8
32
64
96
128
32
64
96
128
32
64
96
128
32
64
96
128
Sealen
Segien
Segien
Sealen
InfLLM (Dense)
InfLUUM (Sparse)
Attention
Other Time
Figure 7: End-to-end inference speed of our 8B model when the number of visible tokens is 6k. TTFT means time-to-firsee firrtttkkntttt and TPOT means time-per-output-token.

---

Published as a conference paper at ICLR 2026
5 CONCLUSION
In this paper, we introduced InILLM-V2, a dense-sparse switchable attention framework designed overcome the limitations existing trainable sparse attention mechanisms. By ensuring architec- tural alignment with the standard pretrain-on-short and finetune-on-long workflow, InfLLM-V2 fa- cilitates a seamless and efficient sparse adaptation to long contexts without requiring extra parame- ters or causing disruptive distributional shifts. We believe MfLLMM-2 offers practical and power- ful solution for advancing the capabilities of large language models in the long-context era.
ACKNOWLEDGMENTS
This work is supported by the National Key Research and Development Program of China (2024YFB4505603), National Natural Science Foundation of China (No. 62576186) and the high- quality development project of MIIT.
ETHICS STATEMENT
Our method has no potential risk since we focus on developing a novel algorithm.
REPRODUCIBILITY STATEMENT
We provided our kernel implementation in the supplementary materials.
THE USE OF LARGE LANGUAGE MODELS
After writing the paper manually, the paper was polished for clarity and readability using a large language model (LLM).
REFERENCES
Wasi Uddin Ahmad, Sean Narenthiran, Somshubra Majumdar, Aleksander Ficek, Siddhartha Jain, Jocelyn Huang, Vahid Noroozi, and Boris Ginsburg. Opencoderreasoning:: Advancing data distil- lation for competitive coding, arXiv preprint arXiv:2504.01943, 2025.
Joshua Ainslie, James Lee-Thorp. Michiel de Jong, Yury Zemlyanskiy, Federico Lebron, and Sumit Sanghai. Gqa: Training generalized multi-query transformer models from multi-head check- points. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Pro- cessing, pp. 4895-4901, 2023.
Jacob Austin, Augustus Odena, Maxwell Nye, Maarten Bosma, Henryk Michalewski, David Dohan, Ellen Jiang, Carrie Cai, Michael Terry, Quoc Le, et al. Program synthesis with large language models. arXiv preprint arXiv:2108.07732 07732, 2021.
Yushi Bai, Xin Lv, Jiajic Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng, Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 3119-3137, 2024.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long- long-document transformer. arXiv preprint arXiv:2004.05/555 2020.
Rishi Bommasani, Drew A Hudson. Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportu- nities and risks of foundation models. Preprint, 2021.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models
10

---

Published as a conference paper at ICLR 2026
few-shot learners. In Advances in Neural Information Processing Systems, volume 33, pp. 1877- 1901, 2020.
Guoxuan Chen, Han Shi, Jiawei Li, Yihang Gao, Xiaozhe Ren, Yimeng Chen, Xin Jiang, Zhenguo Li, Weiyang Liu, and Chao Huang. Seplim: Accelerate large language models by compressing one segment into one separator. In Ferty- International Conference Machine Learning, 2024.
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde De Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained code. arXiv preprint arXiv:2107.03374 2021.
Rewon Child, Scott Gray, Alec Radford, and lya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509, 2019.
Tri Dao. Flashattention-22: Faster attention with better parallelism and work partitioning. In The Twelfth International Conference on Learning Representations, 2024.
Daya DeepSeek, Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseck-rl: Incentivizing reasoning capability in Ilms via reinforcement learning. arXiv preprint arXiv:2501.12948, 2025.
Lizhe Fang, Yifei Wang, Zhaoyang Liu, Chenheng Zhang, Stefanie Jegelka, Jinyang Gao, Bolin Ding, and Yisen Wang. What is wrong with perplexity for long-context language modeling? In The Thirteenth International Conference on Learning Representations, 2025.
Elias Frantar, Roberto L. Castro, Jiale Chen, Torsten Hoefler, and Dan Alistarh. Marlin: Mixed- precision auto-regressive parallel inference on large language models. In Proceedings of the 30th ACM SIGPLAN Annual Symposium on Principles and Practice of Parallel Programming, pp- 239-251, 2025.
Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Peiyuan Zhou, Jiaxing Qi, Junjie Lai, Hayden Kwok-Hay So, Ting Cao, Fan Yang, et al. Seerattention: Learning intrinsic sparse attention in your 1lms. arXiv preprint arX
2024.
Aryo Pradipta Gema, Joshua Ong Jun Leang, Giwon Hong, Alessio Devoto, Alberto Carlo Maria Mancino, Rohit Saxena, Xuanli He, Yu Zhao, Xiaotang Du, Mohammad Reza Ghasemi Madani, et al. Are we done with mmlu? In Proceedings of the 2025 Conference of the Nations of the Amer- icas Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 5069-5096, 2025.
Chi Han, Qifan Wang, Hao Peng, Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Zero-shot extreme length generalization for large language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 3991-4008, 2024.
Xu Han, Zhengyan Zhang, Ning Ding, Yuxian Gu, Xiao Liu, Yuqi Huo, Jiezhong Qiu, Yuan Yao, Ao Zhang, Liang Zhang, et al. Pre-trained models: Past, present and future. Al Open, 2:225-250, 2021.
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations, 2021a.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. In Thirty-fftth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 2), 2021b.
Cheng-Pine Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, and Boris Ginsburg. Ruler: What's the real context size of your long-context language models? In First Conference Language Modeling, 2024.
11

---

Published as a conference paper at ICLR 2026
Shengding Hu, Yuge Tu, Xu Han, Ganqu Cui, Chaoqun He, Weilin Zhao, Xiang Long, Zhi Zheng, Yewei Fang. Yuxiang Huang, et al. Minicpm: Unveiling the potential small language models with scalable training strategies. In First Conference on Language Modeling, 2024.
Yuxiang Huang, Binhang Yuan, Xu Han, Chaojun Xiao, and Zhiyuan Liu. Locret: Enhancing eviction in long-context Ilm inference with trained retaining heads on consumer-graan grade devices. arXiv preprini 4rXiv:2410.01800 2024.
Yuxiang Huang, Mingye Li, Xu Han, Chaojun Xiao, Weilin Zhao, Ao Sun, Hao Zhou, Jie Zhou, Zhiyuan Liu, and Maosong Sun. APB: Accelerating distributed long-context inference by passing compressed context blocks across GPUs. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 10708-10727, 2025.
Yuzhen Huang, Yuzhuo Bai, Zhihao Zhu, Junlei Zhang, Jinghan Zhang, Tangjun Su, Junteng Liu, Chuancheng Lv, Yikai Zhang, Yao Fu, al. C-eval: multi-level multi-discipline chinese eval- uation suite for foundation models. Advances in Neural Information Processing Systems, 36: 62991-63010) 2023.
Naman Jain, King Han, Alex Gu, Wen-Ding Li, Fanjia Yan, Tianjun Zhang, Sida Wang. Armando Solar-Lezama, Koushik Sen, and Ion Stoica. Livecodebench: Holistic and contamination free evaluation of large language models for code. In The Thirteenth International Conference on Learning Representations, 2025.
Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir Abdi, Dongsheng Li, Chin- Yew Lin, al. Minference 1.0: Accelerating pre-filling for long-context Imms via dynamic sparse attention. Advances in Neural Information Processing Systems, 37:52481-52511 2024.
Carlos Jimenez, John Yang, Alexander Wettig, Shunyu Yao, Kexin Pei, Ofir Press, and Karthik Narasimhan. Swe-bench: Can language models resolve real-world github issues? In The Twelfth International Conference on Learning Representations, 2023.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference Learning Representations, 2020.
Xunhao Lai, Jianqiao Lu, Yao Luo, Yiyuan Ma, and Xun Zhou. Flexpreñill: A context-aware sparse attention mechanism for efficient lonu-sequence inference. In The Thirteenth International Con- ference Learning Representations, 2025.
Yuhong Li, Yingbing Huang, Bowen Yang, Bharat Venkitesh, Acyr Locatelli, Hanchen Ye, Tianle Cai, Patrick Lewis, and Deming Chen. Snapkv: Llm knows what you are looking for before generation. Advances in Neural Information Processing Systems, 37:22947-22970, 2024.
Anton Lozhkov, Raymond Li. Loubna Ben Allal, Federico Cassano, Joel Lamy-Poirier, Nouamane Tazi, Ao Tang, Dmytro Pykhtar, Jiawei Liu, Yuxiang Wei, et al. Starcoder and the stack v2: The next generation. arXiv preprint arXiv:2402. 19173, 2024.
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulum Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang. et al. Moba: Mixture of block attention for long-context 1lms. arXiv preprint arXiv:2502.13 1319 2025.
MAA. American invitational mathematics examination-aine
URL https: //maa, org/
maa
Ivan Moshkov, Darragh Hanley, Ivan Sorokin, Shubham Toshniwal, Christof Henkel, Benedikt Schifferer, Wei Du, and Igor Gitman. Aimo-2 winning solution: Building state-of-the-art math- ematical reasoning models with openmathreasoning dataset. arXiv preprint arXiv:2504.16891, 2025.
OpenAI. GPT-4 technical report. Preprint, 2023.
Aaron OpenAI, Jaech, Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, et al. Openai ol system card. arXiv preprint arXiv:2412.16722 2024.
12

---

Published as a conference paper at ICLR 2026
Guilherme Penedo, Hynek Kydlíček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:308 1 1-30849, 2024.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yar:: Efficient context win- dow extension of large language models. In The Twelfth International Conference on Learning Representations, 2023.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguis- fics, 9:53-68, 2021.
Yutao Sun, Zhenyu Li, Yike Zhang, Tengyu Pan. Bowen Dong, Yuyi Guo, and Jianyong Wang. Efficient attention mechanisms for large language models: A survey. arXiv preprint arXiv:2507.19595, 2025.
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery. Quoe Le, Ed Chi, Denny Zhou, et al. Challenging big-bench tasks and whether chain-of-though guistics: ACL 2023, pp. 13003-13051, 2023.
solve them. In Findings of the Association for Computational Lin-
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikei, and Song Han. Quest: Query-aware sparsity for efficient long-context Ilm inference. In Forty-first International Confer- ence on Machine Learning, 2024.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Comput, Surv., 55(6), 2022. ISSN 0360-0300. doi: 10. 10.1145/3530811
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Lei Wang, Chen Ma, Xueyang Feng, Zeyu Zhang, Hao Yang, Jingsen Zhang, Zhiyuan Chen, Jiakai Tang, Xu Chen, Yankai Lin, et al. A survey on large language model based autonomous agents. Frontiers of Computer Science, 18(6):183345, 2024.
Sinong Wang, Belinda Z Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.00 04768. 2020.
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang. Zhiyuan Liu, and Maosong Sun. Infilm: Training-free long-context extrapolation for ilms with effi- cient context memory. Advances in Neural Information Processing Systems, 37:119638-11961 2024a.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. In The Twelfth International Conference Learning Rep- resentations, 2024b.
Guangxuan Xiao, Jiaming Tang, Jingwei Zuo, Shang Yang, Haotian Tang, Yao Fu, Song Han, al. Duoattention: Efficient long-context Ilm inference with retrieval and streaming heads. In The Thirteenth International Conference on Learning Representations, 2025.
Renjun Xu and Jingwen Peng. A comprehensive survey of deep research: Systems, methodologies, and applications. arXiv preprint arXiv:2506.1254) 2025.
Ruyi Xu, Guangxuan Xiao, Haofeng Huang, Junxian Guo, and Song Han. Xattention: Block sparse attention with antidiagonal scoring. In Forty- second International Conference Machine Learn- ing, 2025.
John Yang, Kilian Lieret, Carlos E Jimenez, Alexander Wettig, Kabir Khandpur, Yanzhe Zhang, Binyuan Hui, Ofir Press, Ludwig Schmidt, and Diyi Yang. Swe-smith: Scaling data for software engineering agents. arXiv preprint arXiv:2504.22 20795004..2....... 2025.
13

---

Published as a conference paper at ICLR 2026
Jingyang Yuan, Huazuo Gao, Damai Dai. Junyu Luo, Liang Zhao, Zhengyan Zhang. Zhenda Xie, YX Wei, Lean Wang, Zhiping Xiao, et al. Native sparse attention: Hardware- aligned and natively trainable sparse attention. arXiv preprint arXiv:2502.1108, 2025.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences, Advances in neural information processing systems, 13117283-1227,, 2020.
Jintao Zhang, Rundong Su, Chunyu Liu, Jia Wei, Ziteng Wang, Pengle Zhang, Haoxu Wang, Huiqiang Jiang, Haofeng Huang, Chendong Xiang, al. A survey of efficient attention methods: Hardware-ffficiett sparse, compact, and linear attention. 2025a.
Jintao Zhang, Chendong Xiang, Haofeng Huang. Haocheng Xi, Jun Zhu, Jianfei Chen, et al. Spargeattention: Accurate and training-free sparse attention accelerating any model inference. In Forty-second International Conference on Machine Learning, 2025b.
Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, al. H2o: Heavy-hitter oracle for efficient gen- erative inference of large language models. Advances in Neural Information Processing Systems, 36:34661-34710, 2023.
Yuxiang Zheng, Dayuan Fu, Xiangkun Hu, Xiaojie Cai, Lyumanshan Ye, Pengrui Lu, and Pengfei Liu. Deepresearcher: Scaling deep research via reinforcement learning in real-world environ- ments. arXiv preprint arXiv:2504.03/00 2025.
14

---

Published as a conference paper at ICLR 2026
A THE USE OF LARGE LANGUAGE MODELS
After writing the paper manually, the paper was polished for clarity and readability using large language model (LLM).
B IMPLEMENTATION DETAIL
We have shown the implementation Block Selection in Section 3.4 We show the implementation detail of Sparse Attention here in Algorithm 2
Algorithm Computation of Sparse Attention. (Suppose live 1 without loss of generality.)
Require: Q E Rxd K, V E Rexdd, Block sizes Bg.
Divide Q into blocks Q1..... Qn of size G dg each.
Divide K, V into Tr = / B&] blocks Kı...., KT₃ and VI,..., Vrz of size Be x da each.
Divide O xxxx,, into blocks of size G da each.
Divide the log-sum-exp (se into n blocks of size each.
for ...... n (parallel) do
Load Q from HBM on-chip SRAM.
On chip, initialize 0 (0)axa, 4 (0) m; = (-oo)
for j = 1.... Te (sequential) do
if K, visible tokens (determined by the [Z(i) in Eq. then
Load K, V, from HBM to on-chip SRAM.
On chip, compute attention scores S.j Q.K] Rx
(5-1) rowmax (s.)) E On chip, compute m,)) max(m, (5-1))
On chip, compute P. exp(S., ml (1) E R
On chip, compute e(a) exp(m (5-1) ml (1) ) rowsum Ps)) E R On chip, compute O diag diag(exp(m")
( m m(()) ()) -0:0)-10--1 P.y V-
chip, compute O, diag( e (TT)) (TT))
chip, compute Ise; = m, + log( (e)).
Write O to HBM as the i-th block of O.
Write ise, to HBM as the i-th block of Ise.
return the output and the log-sum-exp Ise.
Algorithm 3 Computation of Dense Attention. (Suppose h&v without loss of generality.)
Require: Q E n K, V Rudds, Block sizes Bq, Br.
Divide Q into T2 G [n/ Bg blocks Q ....... Qr, of size Bg da each.
Divide K, V into T5 = n/ B& blacks K Kr, and VI...., Vr, of size Bg x d, each.
Divide O Gxdd, into T, blocks of size B, x da each.
Divide the log- log-sum-exp (se into T, blocks size B, each. for ,,..., T. (parallel) do
Load Q from HBM on-chip SRAM.
On chip, initialize (,0) (0) (0) xd,, 210) (0) mp!) (-00) B.
for j 1,..,,11 T' (sequential) do
Load K,, V, from HBM to on-chip SRAM.
On chip, compute attention scores S.y
Q.K.J B.
On chip, compute m,." max(m rowmax(S,,)) (,))) R
On chip, compute P E exp(S., - (.)) B.
On chip, compute e(ā) exp(m,, (:11 m
)e!
)(( rowsum (P;;) R..
On chip, compute 0,") (exp(m (-1))
m -
(-1) P, V
On chip, compute O, (T)))) (T)) On chip, compute Ise, m;
(T)) + log Tel
Write O to HBM as the i-th block of O.
Write ise; to HBM as the i-th block of Ise. return the output O and the log-sum-exp Ise.
15

---

Published as a conference paper at ICLR 2026
Similar to FlashAttention (Dao 2024), the algorithm divides the input into blocks. The differences are: 1) The FlashAttention block size By of K, should divide the sparse attention block size B. That is, B should be multiple of Bg. 2) The FlashAttention block of Q typically contains single attention head and multiple tokens. We follow NSA Yuaan et al. 2025) to make contain a group of attention heads of single token, that they the FlashAttention iterates over all blocks of K, whereas our method's ' loop only covers the visible blocks of the sparse attention. We also show the FlashAttention implementation of Dense Attention to Algorithm for reference.
share the same sparse pattern. 3) The inner loop of
C BENCHMARK DETAILS
We provide the detailed results of the LongBench benchmark, mentioned in Table 2 in Table 6 Following LongBench (Bai et al. 2024), the "Overall" score computed by the macro-average over the six task categories.
Table 6: Task Performance on LongBench. Best results in sparse attention are bolded.
Category
PULLAFIN
SHORT YARN
INFL.LM
MINFERENCE
NSA INFLLMM-V2 V2 (SPARSE) INFL.LM--2 (DENSE)
NarQA
21.38
18.17
21.02
20.16
18.34
20.75
21.03
Single Doc
Qusper
43.80
30.98
34.92
44,51
39.96
45.29
45.29
MFRQAeen
55.07
43.81
49
54883
$1.35
$3.53
$3.54
MFQA-zh
57.26
54.51
51.75
57.00
59.06
39.64
HosporQA
50.13
48.49
44.03
48,00
46.78
$4.11
$4.07
ZWiKIQA
QA
MuSiQue
24.68
23.22
17.85
22.87
16.97
21.74
21.24
Darealer
33.54
33.00
33.01
33.94
33.62
33.39
33.29
GovReport
32.17
31.93
40
32.21
28.72
30:33
30.38
Summary
MultiNews
QMSum
24.35
22.45
20.96
25,05
23.81
24.55
24.35
VCSLM
16.37
16.55
17.81
16.17
19:12
16.17
16.20
Few-shot learning
TriviaQA
84.35
85.67
75.78
$1.93
83.95
84.22
84.22
SAMSum
40.26
42.92
37
39.81
38.47
40.69
40.51
Synthetic Task
PsgCount
PapRo-en
86.50
4.00
20.75
4.06
19:00
3:00
3,50
66.00
3.50
92.00
5.00
91.00
4.50
PagRe-zh
90.50
42,00
43.00
90.50
68.00 00.00
90.50
90.50
Code
LOC
35.72
58.65
31:35
35,91
33.83
44.73
44.73
RepoBen-P
34.95
Ovenil 1
42.30
37.86
32.30
41.55
37.10
42.54
42.49
D HYPERPARAMETERS ANALYSIS
In this section, we present analyses and discussions regarding the choice of key hyperparameters in our framework, specifically focusing on the number selected blocks (Zopp)), local blocks ([Tocaal). )). initial blocks ([Tiit)). and the block size (B). These parameters serve the core variables controlling the trade-off between model performance and computational efficiency. As shown in Table increasing [upp from 31 to 63 yields significant performance gains, whereas further creasing it to 95 offers only marginal improvements. Similarly, while larger [Tocaa| improves results, the marginal gain decreases as computational overhead grows. Consequently, we select [Zopk] = 63 and [Toaal = 32 as the optimal configuration for balancing efficiency and effectiveness.
Table 7: Hyperparameter analysis the number of selected blocks and local blocks.
Tinit
Trop
Tocal
RULER 1
LongPPL J
63
32
82.62
2.12
31
32
76.14
2.36
95
32
84.89
2.05
63
16
82.08
2.15
63
64
83.45
2.09
16

---

Published as a conference paper at ICLR 2026
Regarding other parameters, set [nnit] to minimal value following the established practice in StreamingLLM (Xiao et al. 2024b). as a small number of attention sinks is sufficient for stability. For the block size B, adopt B 64 following NSA (Yuan et al. (025). This choice is moti- vated by hardware efficiency, as Flash Attention kernels typically require block sizes of 64 or 128 to maximize Tensor Core utilization and memory bandwidth.
17


Reviews:
[{"review_id": "D5mVL93HWA", "reviewer": "Reviewer_2Rsd", "summary": "This paper presents Dense Sparse Switchable Attention, abbreviated DSSA, a trainable sparse attention framework designed for smooth short to long sequence adaptation in transformer models. The core idea is to reuse the pretrained dense key value parameters for sparse attention, to merge selected block and sliding window patterns into a unified sparse mask, and to replace learned compression with a parameter free multi stage pooling scheme. The paper also proposes engineering optimizations including a fused head group summation and a log sum exp approximation to reduce GPU memory traffic and compute during block selection. Experimental results on long context understanding and long chain of thought reasoning show that DSSA keeps model performance close to dense attention while delivering substantial speedups. Reported highlights include up to 4 times faster end to end inference relative to dense attention with about 98 percent to 99.7 percent task performance retention and kernel level speedups of up to about 7 to 9 times versus FlashAttention in some settings.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper addresses an important and timely problem for large language models namely efficient and stable adaptation from short pretraining to long finetuning. The proposed design is conceptually appealing because it removes the architectural mismatch between dense and sparse attention by reusing existing dense KV projections and producing a single sparse attention output rather than multiple gated outputs. This design choice improves compatibility with the standard pretrain on short then finetune on long workflow and appears to reduce training instability. The technical contributions combine modeling and systems innovations in a way that is practical. On the modeling side the union of selected and sliding patterns into a single mask and the three stage coarse to fine pooling are novel and justified by the need to preserve granularity while remaining block sparse. On the systems side the fusion of group summation into an SRAM friendly loop and the LSE approximation are concrete optimizations that address real GPU IO bottlenecks and yield large empirical speedups. The empirical evaluation is comprehensive in scope and includes long input comprehension, long reasoning and general short sequence tasks, which supports the claim that DSSA preserves accuracy while improving latency. The paper is generally clear and well organized with algorithmic pseudocode, explicit hyperparameter choices, and a sensible set of baselines including dense attention and several sparse methods.", "weaknesses": "1. The paper promises to \"unlock the full potential of sparse attention\" but does not discuss the limits of DSSA in terms of model scale and when the FFN or other layers dominate end to end costs.\n2. The training stability claim is interesting but under analyzed. The loss disruption observed for NSA is shown but the paper should quantify downstream performance variance, report random seed statistics, and show whether DSSA reduces catastrophic forgetting when switching modes.\n3. The NSA comparison can be fragile because NSA is reimplemented and initialized by parameter replication. Authors can provide more details that ensure a fair comparison such as exact initialization, learning rate schedules, and etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "XrfnsSfg9E", "reviewer": "Reviewer_owJK", "summary": "This paper introduces Dense-Sparse Switchable Attention (DSSA), a novel trainable sparse attention framework designed to efficiently adapt Large Language Models (LLMs) from short-sequence pretraining to long-sequence finetuning. The authors identify critical limitations in prior trainable sparse attention methods like Natively Trainable Sparse Attention (NSA), namely the introduction of excessive parameters, architectural mismatches that destabilize training, and inefficiency on short sequences.\n\nDSSA addresses these issues with three main contributions:\n\n- Seamless Short-to-Long Adaptation: DSSA reuses the key-value (KV) projection parameters from the original dense attention model, eliminating the need for extra parameters. It fuses multiple sparse patterns into a unified sparse attention module, maintaining architectural consistency with standard attention. This leads to stable training and preserves knowledge from pretraining.\n- Efficiency on All Sequence Lengths: The parameter-free design allows a model equipped with DSSA to seamlessly switch back to standard dense attention for short sequences without any overhead, while effectively using sparse attention for long sequences.\nHardware-Aware Efficient Implementation: The authors identify the block selection mechanism as a key bottleneck in sparse attention. They propose a novel, two-pass, FlashAttention-style CUDA kernel that significantly reduces I/O overhead by fusing computations in SRAM. An \"LSE Approximation\" technique further optimizes this kernel, reducing the computational cost of the two-pass approach.\n- Experiments on long-context understanding and reasoning tasks show that DSSA achieves performance nearly on par with full attention (retaining over 98%) while being up to 4x faster on the attention component and providing significant end-to-end speedups.", "questions": "1.  Block Selection Robustness: The block selection mechanism relies on attention scores from a compressed, coarse-grained representation of the keys. Could you comment on the robustness of this selection? Are there scenarios or tasks where this coarse approximation might fail to identify crucial fine-grained information, leading to performance degradation?\n2.  Training Stability Details: Figure 5 shows a compellingly smooth training curve for DSSA compared to NSA. Could you elaborate on why DSSA's fine-tuning is more stable? Is it solely due to the shared parameters, or does the unified single-output structure (as opposed to NSA's gated multi-output) also play a significant role?\n3.  On the `DSSA-Dense` Performance: In several tables (e.g., Table 1 on RULER), the `DSSA-Dense` model (which is the finetuned model switched back to dense mode) outperforms the `FULLATTN` baseline that was also finetuned on long contexts. This is an interesting result. Do you have a hypothesis for why this \"sparse-to-dense\" training regimen leads to better performance than continuous dense training? Could the sparse fine-tuning be acting as a form of regularization?\n4.  LSE Approximation Impact: The LSE Approximation is a key part of your efficiency gains, reducing the overhead of the two-pass approach from 2x to 1.25x. Table 1 shows that using the approximation does not hurt performance on RULER. Have you confirmed this holds true across all other tasks? Is there a theoretical or empirical bound on the error introduced by this approximation?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Clear Problem Formulation: The paper does an excellent job of positioning its work. It clearly articulates the shortcomings of the existing state-of-the-art trainable sparse attention method (NSA) and directly proposes solutions for each identified problem (parameter overhead, training instability, and short-sequence inefficiency). The framing around the \"pretrain-on-short, finetune-on-long\" paradigm is highly relevant and practical.\n\n2. Elegant and Simple Design: The core idea of reusing existing parameters and unifying sparse patterns is both simple and powerful. This design choice directly leads to several benefits: zero parameter overhead, stable finetuning (as evidenced by the loss curves in Figure 5), and the flexibility to switch between dense and sparse modes. This makes DSSA a much more practical and \"drop-in\" solution compared to more architecturally disruptive methods.\n\n3. Strong Technical Contribution in Implementation: The work in Section 3.4 is a significant strength. Identifying the I/O bottleneck in the block selection step and engineering a custom CUDA kernel to solve it shows a deep understanding of hardware-aware optimization. The LSE Approximation is a clever trick to mitigate the cost of their two-pass solution. This focus on practical, low-level optimization is crucial for realizing the theoretical benefits of sparse attention and is a valuable contribution.\n\n4. Comprehensive and Convincing Experiments: The empirical evaluation is thorough. The authors compare DSSA not only against full attention and NSA but also against popular training-free methods (InfLLM, MInference) and a context-extension method (YaRN).", "weaknesses": "While the paper is strong, there are a few areas that could be improved or clarified:\n1.  Direct Comparison to NSA: The paper's primary baseline for trainable sparse attention is NSA. However, the authors note they used an open-source Triton implementation of NSA, as the official code was not available. This raises a potential concern about whether the performance drop observed for NSA is due to fundamental architectural flaws (as argued) or sub-optimal implementation/hyperparameters in the third-party version. The claims against NSA would be stronger if benchmarked against an official implementation or with more extensive hyperparameter tuning for the version used.\n2.  Limited Exploration of Hyperparameters: The effectiveness of DSSA's sparse pattern depends on several hyperparameters, such as the number of initial, local, and top-k blocks (`N_init`, `N_local`, `|I_topk|`) and the block size (`B`). The paper fixes these values for the experiments but does not provide an ablation study on their impact. Understanding the sensitivity of the model's performance and efficiency to these choices would be valuable for practitioners.\n3.  Generalizability to Other Architectures: The experiments are conducted on an 8B parameter model with a specific Grouped-Query Attention (GQA) configuration. While the principles of DSSA seem general, its effectiveness and the efficiency of its custom kernels might vary with different model sizes, head dimensions, or attention variants (e.g., Multi-Query Attention). Discussing the potential challenges or necessary adaptations for other architectures would enhance the paper's scope.\n4.  Novelty of Component Ideas: The idea of merging local (sliding window) and global (top-k) attention patterns is not entirely new and exists in prior work like BigBird or Longformer. The paper's main novelty lies in how it frames this within a *trainable, zero-parameter* adaptation framework and the efficient implementation. The paper could be more precise in situating its contribution relative to these earlier sparse attention architectures.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "c3XmaTlUTm", "reviewer": "Reviewer_S847", "summary": "This paper introduced Dense-Sparse Switchable Attention, a sparse attention that adapts model from short to long sequences. DSSA starts from natively trainable sparse attention (SSA), which further modifications, including: 1. Shared Key-Value Projection. Instead of using three sets of KV, DSSA use a single shared set. 2. Remove compressed attention and only used it to get sparse selected attention and sparse sliding attention. Then replace the MLP with pooling function. They did experiments on several benchmarks and demonstrated the efficiency while maintaining the performance.", "questions": "1. In figure 1. portable is mentioned as an advantage. What is the meaning here other than reusable and aligned?\n2. Since the author examed the speed on different gpu, then what are the software configureations like cuda version and pytorch version? Different version could lead to performance gap for other baselines.\n3. What does number of visible tokens mean? Does it mean how many keys/values each query token can attend to?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The proposed architecture is simple to follow. Just a few modification to existing NSA.\n2. The performance is demonstrated across different benchmarks.\n3. The training process is quiet efficient. Figure 6 presented a clear comparision between NDA and DSSA.", "weaknesses": "1. Since DSSA is based on NSA, then the most important baseline should be NSA. Even NSA is 27b total parameters and 3b active parameters, which is larger than DSSA implementation in this paper, it is important to evaluate that on similar benchmarks. For examples, the main table, table 1 in NSA paper presented many benchmarks, which are also major benchmarks for LLM. But DSSA did not evaluate on those.\n\n2. In NSA paper, NSA achieved better performance than Full Attention as in table 2 of the paper. However, in your paper the NSA is much worse than LongBench, which leads to concerns about implementation of NSA. \n\n3. Still, in Figure 6 of NSA paper, it is around 9x faster than full attention but in your case it is only 0.9x to 2.2x. The gap is too large to be ignored.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)