InfLLM-V2: Dense-Sparse Switchable Attention 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

Source-linked AI summary

InfLLM-V2: Dense-Sparse Switchable Attention 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

arXiv:2509.24663v1Sep 29, 2025cs.CLcs.AIcs.LG

TL;DR

Long-sequence processing is bottlenecked by dense attention, while existing trainable sparse methods can add parameters and disrupt short-to-long adaptation. InfLLM-V2 reuses dense attention parameters with switchable dense-sparse computation and an efficient implementation. It runs 4× faster than dense attention while retaining 98.1% and 99.7% of performance on long-context understanding and long chain-of-thought generation, respectively.

Problem
Standard self-attention has severe computational and memory bottlenecks on long sequences, while NSA adds architectural complexity that misaligns with short-to-long adaptation.
Method
InfLLM-V2 reuses pretrained dense attention parameters in a dense-sparse switchable architecture and introduces efficient block selection without extra parameters.
Results
4× faster than dense attention while retaining 98.1% of original performance on long-context understanding and 99.7% on long chain-of-thought generation.
Takeaways & Limitations
InfLLM-V2 provides a practical sparse-attention framework for adapting dense models from short to long contexts while preserving efficiency across sequence lengths.
Takeaways & Limitations
Training-free sparse methods remain constrained by the sparsity they can apply without severe performance degradation, limiting their acceleration benefits.

Abstract

from arXiv · show

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 \textit{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.

1 INTRODUCTION

Long-sequence processing is increasingly important, but standard and trainable sparse attention methods face efficiency, architectural-alignment, and adaptation challenges. InfLLM-V2 addresses these issues with dense-sparse switching, parameter reuse, and efficient block selection, achieving substantial acceleration while preserving performance.

Long-sequence applications require efficient processing because standard Transformer self-attention creates severe computational and memory bottlenecks.
Existing trainable sparse attention, exemplified by NSA, introduces architectural mismatch with the pretrain-on-short, finetune-on-long workflow.
InfLLM-V2 reuses dense attention parameters to transition between dense and sparse attention without additional parameters or disruptive architectural changes.
Its hardware-aware implementation reduces block-selection overhead, enabling more effective acceleration from sparse attention.
InfLLM-V2 is 4× faster than dense attention while retaining 98.1% of original performance on long-context understanding and 99.7% on long chain-of-thought generation.

2 RELATED WORK

Sparse attention research spans training-free and trainable approaches, with methods selecting predefined or dynamically relevant context subsets. Training-free methods often limit sparsity to preserve performance, while trainable methods learn context selection during adaptation.

Training-free sparse attention: Training-free sparse attention uses predefined or dynamic patterns to restrict each token to a relevant subset of context tokens.
Predefined Sparse Patterns: Predefined methods use heuristic structures such as sliding windows, sometimes augmented with globally attended special tokens.
Dynamic Sparse Patterns: Dynamic sparse methods compute query-context relevance, increasingly using contiguous blocks as sequence lengths grow.
Training-free sparse attention: Training-free methods often cannot impose sufficient sparsity without severe performance degradation, limiting their acceleration benefits.
Trainable sparse attention: Trainable methods such as SeerAttention and MoBA learn routers during posttraining or short-to-long adaptation to select relevant contexts.

3 METHOD

InfLLM-V2 modifies dense attention into a parameter-free dense-sparse switchable framework by reusing shared KV projections and aligning sparse computation with dense attention. Its method combines selected and sliding patterns, uses multi-stage compression for block selection, and reduces compression-score I/O overhead.

Overall Framework: InfLLM-V2 reuses one shared set of pretrained KV projections for dense and sparse attention, avoiding NSA’s multiple projection sets and extra parameters.The shared projections are initialized from dense attention and then finetuned for long sequences.
Overall Framework: InfLLM-V2 merges Selected Attention and Sliding Attention into one sparse pattern, while retaining compressed attention scores only for block selection.The unified pattern expands local blocks enough to cover the sliding window.
Overall Framework: InfLLM-V2 dynamically switches between dense attention for short inputs and sparse attention for long inputs based on sequence length.Eliminating compressed-attention outputs makes the sparse computation more closely resemble dense attention.
Block Representation: Three-stage coarse-to-fine compression preserves granular information while producing group-level scores for efficient block-sparse attention.Mean pooling creates an intermediate representation, group-wise aggregation produces shared importance scores, and max pooling retains salient features.
Efficient Implementation: The efficient implementation addresses compression-score bottlenecks by avoiding materialization of first-stage scores in GPU HBM, where writing h_qn^2/s_C1 values is costly.The implementation uses on-chip processing across coarse- and fine-grained passes; the supplied passages identify memory I/O as the primary bottleneck.

4 EXPERIMENT

Experiments evaluate InfLLM-V2 across long-context understanding, long reasoning, general tasks, and inference efficiency against dense and sparse baselines. The method remains competitive with full attention while providing substantial kernel and end-to-end speedups.

Experiment Setup: NSA disrupts the training loss, whereas InfLLM-V2 stays closer to FULLATTN during trainable sparse adaptation.The comparison uses the same sparsity level across sparse attention methods.
Long-Context Understanding: InfLLM-V2 achieves the best performance among sparse methods on long-context benchmarks, closely matching the FULLATTN baseline.Evaluations cover RULER at 32k, LongBench, and LongPPL.
Long Reasoning: InfLLM-V2 attains performance on par with full attention on long-output reasoning tasks including MATH-500, AIME, and LiveCodeBench.The models are finetuned on OpenMathReasoning and OpenCodeReasoning.
General Tasks: InfLLM-V2 remains comparable to full attention on short-sequence general tasks after long-sequence fine-tuning and switching back to dense mode.The evaluation includes MMLU, MMLU-Redux, CEval, MATH-500, HumanEval, MBPP, and BBH.
Efficiency: 7.4× and 9.3× speedups over FlashAttention are achieved on A100 and 4090, respectively, when 16 blocks are selected.NSA reaches 3.5× in the same setting, while the efficient implementation reduces Block Selection overhead.
Efficiency: 2.13× prefilling and 2.32× decoding speedups are achieved end-to-end with |I| = 96 and W4A16 quantization.The reported implementation does not accelerate FFN layers.

5 CONCLUSION

The paper concludes that InfLLM-V2 enables efficient sparse adaptation to long contexts while preserving alignment with the standard short-pretraining and long-finetuning workflow.

5 CONCLUSION: InfLLM-V2 avoids extra parameters and disruptive distributional shifts through architectural alignment with the pretrain-on-short, finetune-on-long workflow.The framework is presented as a practical solution for efficient long-context adaptation.

A IMPLEMENTATION DETAIL

The implementation divides queries, keys, values, outputs, and log-sum-exp values into blocks, then computes sparse attention only over visible key-value blocks. It follows FlashAttention-style tiling while restricting the inner loop to selected blocks.

A IMPLEMENTATION DETAIL: Sparse attention partitions Q, K, V, O, and log-sum-exp values into blocks for parallel query processing and sequential key-value traversal.Each query block loads its data on chip and writes output and log-sum-exp blocks back to HBM.
A IMPLEMENTATION DETAIL: The sparse kernel loads and computes attention only when a key block belongs to the visible-token set determined by selected blocks.This distinguishes its inner loop from dense FlashAttention, which iterates over all key blocks.
A IMPLEMENTATION DETAIL: The algorithm initializes per-query-block outputs and log-sum-exp values, computes attention scores for visible blocks, and returns both outputs and log-sum-exp values.The displayed implementation uses on-chip SRAM for block computation and HBM for input and output movement.
A IMPLEMENTATION DETAIL: Sparse attention uses FlashAttention-style block computation but requires the FlashAttention key-block size to divide the sparse attention block size.The sparse attention block size B must be a multiple of Bk.

B BENCHMARK DETAILS

The benchmark details define LongBench’s reported Overall score as a macro-average across six task categories and provide detailed LongBench results in Table 6.

B BENCHMARK DETAILS: LongBench’s Overall score is computed as the macro-average over six task categories.Table 6 provides the benchmark’s detailed task-performance results.