PaperlayerView exact revision on arXiv

Source-linked AI summary

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps

Yanke Zhou, Yiduo Li, Hanlin Tang, Maohua Li, Kan Liu, Tao Lan, Lin Qu, Yuan Yao, Xiaoxing Ma

arXiv:2605.16928v2cs.CLcs.AI

TL;DR

Long-context inference is bottlenecked by full attention’s cost, while sparse alternatives often require costly training or heuristic eviction. RTPurbo sparsifies full-attention models through specialized heads, low-dimensional retrieval indexing, and dynamic selection, preserving near-lossless accuracy while achieving up to 9.36× prefill and 2.01× decode speedups.

  • Problem

    Full attention makes long-context inference inefficient, while existing sparse approaches face trade-offs among efficiency, training cost, and accuracy.

  • Method

    RTPurbo uses retrieval/streaming head specialization, low-dimensional token indexing, and dynamic top-p selection to adapt full-attention models into sparse inference systems.

  • Results

    RTPurbo preserves near-lossless accuracy on long-context and reasoning tasks, with up to 9.36× prefill and 2.01× decoding speedups.

  • Takeaways & Limitations

    Full-attention models can support effective sparse execution through lightweight post hoc adaptation, without native sparse pretraining.

  • Takeaways & Limitations

    Reported ultra-long-context sparsity is averaged across query heads, so actual sparsity varies among individual heads.

Abstract

from arXiv · show

Long-context inference in large language models is bottlenecked by the quadratic cost of full attention. Existing efficient alternatives often rely either on native sparse training or on heuristic token eviction, creating an undesirable trade-off among efficiency, training cost, and accuracy. In this work, we show that full-attention LLMs are already intrinsically sparse and can be transformed into highly sparse models with only minimal adaptation. Our approach is built on three observations: (1) only a small subset of attention heads truly requires full long-context processing; (2) long-range retrieval is governed primarily by a low-dimensional subspace, allowing relevant tokens to be retrieved efficiently with a 16-dimensional indexer; and (3) the useful token budget is strongly query-dependent, making dynamic top-$p$ selection more suitable than fixed top-$k$ sparsification. Based on these insights, we propose RTPurbo, which retains the full KV cache only for retrieval heads and introduces a lightweight token indexer for sparse attention. By exploiting the model's intrinsic sparsity, RTPurbo achieves sparsification with only a few hundred training steps. Experiments on long-context benchmarks and reasoning tasks show that RTPurbo preserves near-lossless accuracy while delivering substantial efficiency gains, including up to a 9.36$\times$ prefill speedup at 1M context and about a 2.01$\times$ decode speedup. These results suggest that strong sparse inference can be obtained from standard full-attention training without expensive native sparse pretraining.

1 Introduction

Long-context inference is bottlenecked by the rapidly increasing cost of full attention, yet full-attention LLMs exhibit intrinsic sparsity that RTPurbo exploits through selective heads, efficient token indexing, adaptive sparsity, and minimal adaptation. RTPurbo achieves near-lossless compression with substantial prefill and decoding speedups.

  • Full-attention cost grows rapidly with context length, making long-context inference a major efficiency bottleneck and motivating sparse attention.
  • Full-attention models exhibit sparsity at both the head and token levels, with most heads relying primarily on local information.
  • RTPurbo addresses head selection, efficient token indexing, and adaptive sparsity, because different queries require different numbers of attended tokens.
  • Over 90% recall is achieved using only 16 dimensions, supporting low-dimensional retrieval indexing for relevant tokens.The analysis finds that high-frequency components contribute little to long-range retrieval and can interfere with it.
  • 9.36× prefill speedup and 2.01× decoding speedup are achieved with custom sparse kernels, while self-distillation requires only a few hundred training steps.The alignment stage uses about 1M label tokens to recover sparsified-model performance and reduce overfitting risk.

2 Insight Behind RTPurbo

RTPurbo exploits three intrinsic properties of full-attention LLMs: retrieval heads handle distant semantic retrieval, low-frequency RoPE components preserve long-range signals in a compact subspace, and token requirements vary by query.

  • Retrieval heads: Retrieval heads attend to semantically related tokens across distant context regions, unlike most heads that focus mainly on local information.This retrieval behavior motivates selectively preserving long-context information for retrieval heads.
  • Retrieval heads: Only a small subset of heads retrieves distant relevant content, while most heads primarily process local information.RTPurbo retains the full KV cache only for retrieval heads and discards remote tokens for the remaining intrinsically sparse heads.
  • Compact retrieval subspace: Low-frequency RoPE components vary smoothly with relative distance and better preserve long-range retrieval signals than high-frequency components.RTPurbo uses this low-frequency structure as a compact retrieval subspace for low-cost token selection without full-dimensional scoring.
  • Query-dependent token budgets: Top-16k recovers only 3.8% more attention mass than dynamic top-p but requires computing about 8k additional tokens.Fixed top-k can under-retrieve diffuse queries or waste computation on concentrated queries, motivating dynamic selection.

3 Method

RTPurbo combines offline retrieval-head identification, query-aware sparse attention, two-stage self-distillation, and hardware-aware decoding. It preserves exact attention features while routing sparsity through low-rank pre-RoPE projections and dynamic top-p selection.

  • 3.2 Sparse computation pattern: Top-16k computes about 8k extra tokens than top-p but recovers only 3.8% more attention mass.This motivates dynamic top-p selection over fixed top-k sparsification.
  • 3.1 Retrieval-head identification: RTPurbo scores retrieval heads by measuring attention mass from a later needle to an earlier identical needle in a calibrated long document.A single long sequence is sufficient for offline partitioning into retrieval and local head sets, performed only once.
  • 3.2 Sparse computation pattern: Local heads use sliding-window attention with sinks, while retrieval heads use dense prefill followed by query-aware dynamic sparse decoding.Low-rank pre-RoPE projections route relevant tokens; final attention retains the complete feature space and exact relative positional geometry.
  • 3.3 Two-stage training: The first training stage freezes the backbone and trains low-dimensional projection weights to match exact attention distributions using KL divergence.The second stage enables sparse attention and uses end-to-end self-distillation from dense teacher next-token predictions.
  • 3.4 Hardware-aware decoding kernel: A custom GPU kernel implements block-wise top-p decoding while addressing sort-free thresholding and memory-efficient sparse decoding over long contexts.Histogram-based thresholding avoids O(Nb log Nb) sorting and O(Nb) per-head binary-search memory, while register-resident single-warp CTAs optimize bandwidth.

4 Experiments

RTPurbo is evaluated on long-context and reasoning benchmarks against representative sparse-attention baselines, showing near-lossless accuracy, query-dependent sparsity, and substantial prefill acceleration. Its dynamic top-p mechanism remains effective at extreme context lengths, while the specialized decode kernel outperforms dense and native PyTorch baselines.

  • Benchmarks and Models: RTPurbo is evaluated on LongBench and RULER for long-context processing, and AIME24, AIME25, and MMLU-PRO for long-decode and general reasoning.The long-context models include Qwen3-Coder-30B-A3B; the supplied passage truncates the remaining model details.
  • Baselines: RTPurbo is compared with RazorAttn, Minference, FlexPrefill, Quest, and SnapKV under aligned evaluation settings.The comparison uses a 15% retrieval-head ratio for RazorAttn and γ=0.9 for FlexPrefill.
  • Sparsity Analysis: 5× variation in active tokens appears at 32K, with 468.8 for niah-S and 2462.1 for multi-K under query-dependent dynamic top-p selection.During prefill, 15% of heads perform dense attention, while 85% attend only to 4 sink tokens and an 8192-token window.
  • Sparsity Analysis: 97.1% sparsity is achieved at 512K while robust accuracy is sustained on ultra-long multi-hop tasks.The result is attributed to dynamic thresholding and query-aware token selection.
  • Runtime Analysis: 9.36× prefill speedup is reached at 1M context versus FlashAttention-2, increasing from 2.83× at 32K; the top-p decode kernel also outperforms FA2 and native PyTorch.RTPurbo outperforms the other sparse baselines in long-context prefill and delivers stable decode speedups over FA2.

5 Conclusion

RTPurbo shows that full-attention LLMs can become efficient sparse inference systems through minimal post hoc adaptation. It preserves near-lossless accuracy while providing substantial speedups, indicating that native sparse pretraining is not the only route to efficient long-context inference.

  • Core contribution: Full-attention LLMs are intrinsically sparse and can be transformed into efficient sparse inference systems with minimal adaptation.RTPurbo implements this transformation through head-wise specialization, low-dimensional retrieval indexing, and dynamic top-p selection.
  • Empirical findings: RTPurbo preserves near-lossless accuracy on long-context and reasoning tasks while delivering substantial prefill and decode speedups.
  • Broader implication: Full-attention models can support effective sparse execution with lightweight post hoc adaptation, rather than requiring native sparse pretraining.
  • Caveat: Reported ultra-long-context sparsity averages compute sparsity across query heads, while individual heads naturally exhibit different sparsity levels.Per-head behavior is analyzed in Appendix A.

6 Related work … A.1 Headwise Distribution.

The paper situates RTPurbo among block-, token-, and pattern-based sparse-attention methods, then shows that retrieval ability is concentrated in a small, stable subset of later-layer heads across two Qwen3 models.

  • 6 Related work: Block-sparse methods select subsets of key–value blocks using query-aware ranking, routing, softmax-contribution estimates, training-free refinement, or spectral criteria.QUEST, MoBA, BLASST, SpargeAttention, and Prism differ mainly in how they estimate block importance.
  • 6 Related work: Token-wise methods estimate token relevance before exact attention, using learned indexing, RoPE frequency structure, or recent-query relevance for KV-cache compression.The passage identifies DSA, FASA, and SnapKV as representative approaches.
  • 6 Related work: Pattern-based methods adapt sparsity to head behavior through offline patterns, context-aware selection, or retrieval–streaming head partitioning; RTPurbo adds low-dimensional indexing and dynamic top-p selection.This places RTPurbo closest to head-partitioning approaches such as MInference, FlexPrefill, DuoAttention, and RazorAttention.
  • Appendix — A Headwise Analysis of Local/Retrieval Patterns and Retrieval Sparsity: The appendix analyzes local and retrieval head-pattern distributions in Qwen3-Coder-30B-A3B and Qwen3-30B-A3B-Think.The analysis targets multi-head attention behavior in the two experimental models.
  • A.1 Headwise Distribution.: Retrieval-score heatmaps are analyzed across all query heads in Qwen3-Coder-30B-A3B and Qwen3-30B-A3B-Think.Figure 8 presents the per-head retrieval-score distributions used for the subsequent observations.
  • A.1 Headwise Distribution.: Only a small subset of heads shows strong retrieval ability, while most heads have low retrieval scores and rely primarily on local context or sink tokens.Long-range token-recall capacity is concentrated in a minority of heads.
  • A.1 Headwise Distribution.: High-scoring retrieval heads appear almost exclusively in the latter half of the model, where representations are more stable and semantically rich.Early layers primarily provide local contextualization while later layers support reliable long-range retrieval.
  • A.1 Headwise Distribution.: Individual-head retrieval scores remain highly consistent across documents, supporting offline calibration on one long sequence to partition retrieval and local heads.The observed stability supports the assumption that head behavior is highly input-agnostic.

A.2 Diversity of Sparsity between Retrieval Heads. · B Ablations on RTPurbo Design Choices

Retrieval heads exhibit highly diverse, stable sparsity patterns, with token budgets varying by more than three orders of magnitude across heads. RTPurbo’s ablations examine several key design settings.

  • A.2 Diversity of Sparsity between Retrieval Heads.: Retrieval heads retain substantially different numbers of active tokens under the same top-p threshold and input sequence.Table 7 reports recalled-token counts for three representative retrieval heads across different sequence lengths.
  • A.2 Diversity of Sparsity between Retrieval Heads.: At 64K context, L43H31 retains 21 tokens while L24H25 requires 24,621 tokens to cover 90% of attention mass.Both heads use top-p (p = 0.9), demonstrating over three orders of magnitude in token-budget variation.
  • A.2 Diversity of Sparsity between Retrieval Heads.: Across 32K, 64K, and 128K contexts, L43H31 consistently retains fewer tokens than L2H19 and L24H25.Absolute token counts scale with sequence length, but the relative ordering of heads by sparsity remains stable.
  • A.2 Diversity of Sparsity between Retrieval Heads.: Each retrieval head’s tendency toward concentrated or diffuse retrieval is an intrinsic property rather than a transient artifact of a specific input.This conclusion follows from the consistency of head-level sparsity patterns across context lengths.
  • A.2 Diversity of Sparsity between Retrieval Heads.: RTPurbo uses independent, query-aware top-p thresholds so concentrated heads use minimal token budgets while diffuse heads expand their active sets when necessary.Per-head dynamic thresholding accommodates sparsity diversity that static top-k cannot simultaneously serve.
  • B Ablations on RTPurbo Design Choices: The paper conducts ablation studies on several key design settings in RTPurbo.These experiments comprise the section’s analysis of RTPurbo design choices.

B.1 Ablation on Retrieval Head Ratio

The ablation finds that increasing RTPurbo’s retrieval-head ratio from 15% to 30% yields almost no accuracy improvement while reducing sparsity and increasing training cost. Reducing the ratio to 10% instead causes substantial accuracy drops, indicating insufficient retrieval heads for robust long-range recall.

  • 15% and 30% retrieval-head ratios were compared on representative MMLU-PRO and RULER subcategories.
  • 30% brings almost no accuracy improvement over 15% on either MMLU-PRO or RULER.
  • 30% directly reduces overall sparsity and increases training cost by requiring roughly twice as many low-dimensional projection parameters in stage one.
  • 10% causes substantial accuracy drops on several representative benchmarks, indicating too few retrieval heads for robust long-range recall.

B.2 Ablation on Low-dimension Size · C Details of the Two-Stage Training Pipeline

The ablation compares projection dimensions of 4, 16, and 32, finding that dimension 16 balances fitting quality, sparsity, and efficiency despite dimension 4’s higher end-to-end accuracy. The following section introduces the two-stage training pipeline for Qwen3-Coder-30B-A3B.

  • B.2 Ablation on Low-dimension Size: Dimensions 16 and 32 achieve nearly identical end-to-end accuracy across the evaluated representative benchmarks.
  • B.2 Ablation on Low-dimension Size: The low-dimension ablation evaluates projection sizes of dim = 4, dim = 16, and dim = 32.These settings are assessed on a subset of representative benchmarks for end-to-end accuracy and relevance-space fitting quality.
  • B.2 Ablation on Low-dimension Size: Dimension 4 records the highest end-to-end accuracy but requires recalling many more tokens to cover the same attention mass.Its weaker fitting ability reduces actual sparsity, so the accuracy advantage does not indicate a better relevance-space approximation.
  • B.2 Ablation on Low-dimension Size: With top-p p = 0.9, overly small projection dimensions force the selector to retain substantially more tokens for equivalent attention-mass recovery.The recalled-token analysis uses retrieval head L24H25 and exposes a trade-off between relevance-space expressiveness and sparsity.
  • B.2 Ablation on Low-dimension Size: Dimension 16 is selected as the default because it best balances fitting accuracy, sparsity, and computational efficiency.Larger dimensions add flexibility without improving token-selection quality.
  • C Details of the Two-Stage Training Pipeline: Section C details the two-stage training pipeline and analyzes its overall workflow for Qwen3-Coder-30B-A3B.

C.1 Details of Low-dimension Projection Training · C.2 Details of End-to-End Self-distillation Training

RTPurbo’s two-stage adaptation trains low-dimensional retrieval projections on long sequences, then uses end-to-end self-distillation to preserve the original model’s predictions. Both stages converge in roughly 600 steps while using compact supervision relative to their long-context inputs.

  • C.1 Details of Low-dimension Projection Training: 210 retrieval heads, roughly 15% of Qwen3-Coder-30B-A3B’s 1536 query heads, are selected after offline calibration.Each head has dimension d_h = 128.
  • C.1 Details of Low-dimension Projection Training: 8,000 FineWeb sequences between 32K and 80K tokens form the stage-1 training set for aligning retrieval-head attention-score distributions.The supplied passage truncates the remainder of the alignment objective.
  • C.1 Details of Low-dimension Projection Training: 10−3 is the stage-1 peak learning rate, reached through 100-step linear warmup before cosine annealing decay.The training loss for Layer 24 Head 25 converges well within about 600 steps, with similar behavior across other retrieval heads.
  • C.2 Details of End-to-End Self-distillation Training: 8,000 Dolma 3 Longmimo Mix reasoning dialogues longer than 32K tokens provide stage-2 data, averaging 48K tokens and 300 label tokens.The original model’s top-10 next-token logits are cached at each position as distillation targets.
  • C.2 Details of End-to-End Self-distillation Training: 600 steps is the approximate convergence time for stage-2 training, which uses a relatively small learning rate to limit capability drift.The passage attributes this training behavior to Figure 9b and the stated finetuning setting.
  • C.2 Details of End-to-End Self-distillation Training: 180M tokens comprise the full stage-2 corpus, while only about 1.2M label tokens participate directly in learning.This contrast reflects long-context inputs paired with sparse labeled supervision.

D Limitation

RTPurbo’s limitations concern its reliance on stable attention-head specialization, incomplete prefill sparsification, and limited evaluation breadth. Future work targets adaptive routing, stronger prefill sparsification, and broader cross-model validation.

  • Head specialization: RTPurbo relies on offline calibration to partition attention heads into retrieval and local groups, which may degrade with weaker specialization or substantial domain shift.The observed partition is stable in the studied models but may not generalize under these conditions.
  • Sparsification: Retrieval heads still use full dense attention during prefill, leaving the current design incompletely sparsified.This limitation specifically affects prefill rather than eliminating sparsification altogether.
  • Evaluation scope: Evaluation mainly covers the Qwen3 family and long-context and reasoning workloads, requiring broader validation across architectures and domains.The paper identifies broader cross-model and cross-domain evaluation as still needed.
  • Future work: Future work will pursue more adaptive head routing, stronger prefill sparsification, and broader cross-model evaluation.These directions are proposed to address the method’s current limitations.
Loading 2605.16928v2…