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

Problem Statement

Long-context inference in Large Language Models (LLMs) is currently bottlenecked by the quadratic computational and memory costs associated with full attention mechanisms. While existing efficient alternatives exist, they often rely on either native sparse training (which is expensive) or heuristic token eviction strategies (which create a trade-off between efficiency, training cost, and accuracy). The central question addressed by this work is whether full-attention models, already trained, possess intrinsic sparsity that can be exploited to transform them into highly sparse models with minimal adaptation, thereby avoiding the high costs of native sparse pretraining.

Methodology: RTPurbo

The authors propose RTPurbo, a head-wise sparse attention framework designed to convert full-attention LLMs into efficient sparse inference systems through minimal surgical adaptation. The method is grounded in three core observations and corresponding technical components:

1. Head Specialization and Offline Calibration

The authors observe that attention heads in pretrained LLMs are not homogeneous; they specialize into distinct functional roles. Specifically, a small subset of heads (retrieval heads) is responsible for retrieving distant, semantically relevant content, while the majority (local heads) primarily process local information.

• Implementation: RTPurbo employs an offline calibration phase where a lightweight sequence with a "needle" inserted at the beginning and end is used to score heads based on their ability to retrieve the earlier needle from the later one.
• Action: Based on these scores, heads are partitioned into a retrieval set ($H_{ret}$) and a local set ($H_{loc}$). During inference, local heads use a sliding window with attention sinks, while retrieval heads retain the full KV cache during the prefill phase but switch to dynamic sparse selection during decoding.

2. Low-Dimensional Token Indexing for Retrieval

For retrieval heads, the challenge is to identify relevant tokens efficiently without full-dimensional scoring.

• Theoretical Insight: The authors analyze the impact of Rotary Position Embeddings (RoPE), finding that high-frequency components vary rapidly with distance and degrade long-range affinity, whereas low-frequency components preserve retrieval signals.
• Implementation: A low-dimensional projector ($r \ll d$, specifically 16 dimensions) is applied to query and key features before RoPE injection. This projects the high-dimensional space into a compact subspace where relevance is estimated. This indexer allows for efficient token selection with over 90% recall using only 16 dimensions.

3. Dynamic Top-p Selection vs. Static Top-k

The authors argue that the required token budget for retrieval is strongly query-dependent. Some queries trigger broad retrieval (needing many tokens), while others are highly concentrated (needing few).

• Implementation: Instead of a fixed top-$k$ budget, RTPurbo uses a dynamic top-$p$ selector. It accumulates attention mass in the low-dimensional space until a cumulative threshold $p$ (e.g., 0.9) is reached. This adapts the sparsity level to the specific query complexity, avoiding the information loss of small $k$ or the wasted computation of large $k$.

4. Two-Stage Training Pipeline

To restore model performance after sparsification, RTPurbo utilizes a lightweight two-stage training process:

• Stage 1 (Projection Training): The backbone LLM is frozen. Only the low-dimensional projection weights for the identified retrieval heads are trained to minimize the KL divergence between the original full attention distribution and the projected low-dimensional distribution.
• Stage 2 (Self-Distillation): The sparse model acts as a student learning to match the dense teacher's next-token predictions. Crucially, this stage uses self-distillation (aligning top-10 logits) rather than standard supervised fine-tuning, which avoids the need for complex data mixture tuning and prevents overfitting. This stage requires only a few hundred training steps (approx. 1M label tokens).

5. Hardware-Aware Kernels

The authors implement custom GPU kernels to handle the dynamic top-$p$ selection efficiently. These kernels utilize a sort-free histogram-based approach for thresholding and bandwidth-optimized sparse decoding to manage long contexts without excessive memory overhead.

Key Results

Experiments were conducted on NVIDIA H20 GPUs using Qwen3-Coder-30B-A3B and Qwen3-30B-A3B-Think models.

• Accuracy Preservation: RTPurbo achieves near-lossless accuracy on long-context benchmarks (LongBench, RULER) and reasoning tasks (AIME, MMLU-PRO).
  • On RULER at 64K context, RTPurbo (top-p) achieves an average score of 90.06%, outperforming all sparse baselines (e.g., Minference at 65.61%, FlexPrefill at 77.77%) and matching the full attention baseline (86.23% at 64K, though note the baseline drops significantly at longer lengths in some tasks, while RTPurbo sustains high accuracy).
  • On AIME24/25, RTPurbo matches the full attention baseline exactly (86.67).
• Efficiency Gains:
  • Prefill Speedup: RTPurbo delivers up to 9.36× speedup at 1M context length compared to FlashAttention-2, with speedups increasing linearly with context length (2.83× at 32K).
  • Decode Speedup: RTPurbo achieves up to 2.01× speedup in decoding at 1M context.
  • Sparsity: The method achieves high sparsity (up to 89.2% at 64K and over 97% at 512K) while maintaining high attention mass (>0.93).
• Training Efficiency: The entire adaptation process requires only hundreds of training steps (approx. 1M label tokens), demonstrating that native sparse pretraining is not necessary for effective sparse inference.

Significance and Claims

The paper claims that full-attention models are intrinsically sparse and can be transformed into highly efficient sparse inference systems with minimal adaptation. The significance of this work lies in several key points:

1. Redefining the Path to Efficiency: The authors challenge the notion that native sparse pretraining is the only path to efficient long-context inference. They demonstrate that standard full-attention training remains a highly competitive choice, as the resulting models can be "sparsified" post-hoc with very small additional cost.
2. Interpretability and Minimal Surgery: The approach is highly interpretable, relying on the natural specialization of heads (retrieval vs. local) and low-dimensional subspaces. It avoids the "black box" nature of heuristic token eviction by grounding the sparsity in the model's learned attention patterns.
3. Dynamic Adaptation: By moving away from static top-$k$ constraints to dynamic top-$p$ selection, the method effectively handles the varying complexity of different queries, a limitation the authors identify in existing fixed-budget approaches.
4. Practicality: The ability to achieve near-lossless accuracy with only a few hundred training steps and lightweight hardware kernels suggests that this paradigm is immediately applicable to existing large-scale models without the prohibitive costs of retraining from scratch.

In summary, RTPurbo provides a practical, efficient, and theoretically grounded method to unlock the sparse inference potential inherent in standard full-attention LLMs.