How InfLLM‑V2 Achieves Seamless Short‑to‑Long Context Upgrade with Minimal Structural Changes

Background and Motivation

As large language models (LLMs) expand to more demanding reasoning and long‑document tasks, context length has become a critical bottleneck: traditional Transformer attention incurs quadratic time and memory costs as sequence length grows, making naïve extensions prohibitively expensive.

While sparse‑attention methods have been widely explored, many introduce new structures or trainable modules that misalign with the prevailing "short‑sequence pre‑training → long‑sequence fine‑tuning" paradigm, leading to degraded performance when transitioning from short to long contexts.

InfLLM‑V2 Proposal

The Tsinghua team led by Liu Zhiyuan presents InfLLM‑V2: Dense‑Sparse Switchable Attention for Seamless Short‑to‑Long Adaptation . Instead of adding new parameters or altering output forms, the method switches from dense to sparse attention only when the sequence exceeds a preset threshold, reusing the original Key and Value projections and keeping the attention output a single tensor.

This “minimal structural disturbance” aims to retain the expressive power of dense attention while gaining the computational efficiency of sparsity.

Experimental Questions

The study evaluates three core questions:

Can the sparse mode match full‑attention performance on long‑context tasks?

Does the switchable design preserve the model’s capabilities on short‑sequence tasks under the realistic "short‑pre‑train → long‑fine‑tune" workflow?

Do the theoretical attention‑kernel speedups translate into end‑to‑end inference gains?

Results on Long‑Context Understanding

On the 32k RULER benchmark, InfLLM‑V2 (Sparse) achieves performance virtually identical to Full Attention across most sub‑tasks, whereas other sparse methods (e.g., NSA, InfLLM, MInference) show noticeable drops.

On the more realistic LongBench suite—covering QA, summarization, reasoning, and multilingual tasks—InfLLM‑V2 (Sparse) slightly surpasses Full Attention, while NSA lags significantly and the length‑extrapolation method SHORT+YaRN suffers large degradations.

In the LongPPL perplexity evaluation, InfLLM‑V2’s perplexity aligns with Full Attention, while NSA’s perplexity is markedly higher, indicating poor long‑sequence language modeling after short‑to‑long transfer.

Chain‑Reasoning and Short‑Sequence Retention

For long‑chain reasoning benchmarks (MATH‑500, AIME, LiveCodeBench), InfLLM‑V2 (Sparse) remains on par with Full Attention, confirming that its sparse mechanism does not break the “thinking continuity” required for multi‑step reasoning.

When fine‑tuned on long contexts and then evaluated on short‑sequence tasks (MMLU, CEval, HumanEval), the model switched back to dense mode retains performance comparable to Full Attention, whereas NSA exhibits clear degradation.

Inference Efficiency

At a visible token count of 6k (|I|=96), InfLLM‑V2 delivers approximately 2.1× speedup in prefilling and 2.3× speedup in decoding, even without any optimizations to the feed‑forward network, demonstrating practical end‑to‑end acceleration.

Design Rationale and Ablations

The authors emphasize that any sparse‑attention scheme must avoid altering the dense‑attention output form or introducing extra KV branches, as such changes would disrupt the representations learned during short‑sequence pre‑training.

Implementation details include three key changes during long‑context training: (1) switching the attention mask to a sparse pattern, (2) reusing the original KV projection weights, and (3) preserving a single‑output attention structure without gating or multi‑head aggregation.

Ablation studies identify the block‑selection stage—particularly compression‑attention computation and explicit score materialization—as the main bottleneck. Optimizations such as head‑group fusion and LSE approximation reduce block‑selection time by 20–30% with negligible impact on model quality.

Implications

InfLLM‑V2 demonstrates that sparse attention can be integrated into existing LLMs with negligible architectural changes, enabling a “hot upgrade” to long‑context capability without retraining from scratch, expanding model utility in real‑world deployments.