Token Sparse Attention: Efficient Long-Context Inference with Interleaved Token Selection

1. Problem Statement

Large Language Models (LLMs) face a fundamental bottleneck in long-context inference due to the quadratic complexity ($O(L^2)$) of the attention mechanism during the prefill stage. While hardware optimizations like FlashAttention reduce memory I/O overhead, they do not alter the algorithmic complexity.

Existing acceleration methods suffer from two main limitations:

• Structured Sparsity: Methods like Minference or FlexPrefill apply block-level sparsity. This often retains irrelevant tokens if they are grouped with important tokens in the same block, leading to sub-optimal sparsity.
• Token Eviction: Methods like FastKV or GemFilter permanently discard tokens in early layers. This assumes token importance is static, ignoring the dynamic nature of token relevance across different layers and attention heads. Once evicted, a token cannot be reconsidered, leading to accuracy degradation.

2. Methodology: Token Sparse Attention

The authors propose Token Sparse Attention, a lightweight, dynamic, and reversible token-level sparsification mechanism. Unlike eviction-based methods, it does not permanently remove tokens but compresses the computation space and then restores the sequence.

Core Architecture

The method operates in two stages per attention head:

1. Compression (QKV Selection):
   • For each attention head $h$, a subset of token indices $S_h$ is selected based on importance scores.
   • The Query ($Q$), Key ($K$), and Value ($V$) tensors are gathered to form compressed tensors $\hat{Q}, \hat{K}, \hat{V}$ with a reduced sequence length $L' \ll L$.
   • Key Feature: Selection is performed independently per head, allowing different heads to focus on different subsets of tokens, addressing head-wise heterogeneity.
   • The compressed tensors remain dense and contiguous, enabling the use of optimized kernels like FlashAttention.
2. Decompression (Output Restoration):
   • The attention output $\hat{O}$ (shape $L' \times d$) is scattered back into a zero-initialized tensor of the original sequence length $L \times d$ using the head-specific indices.
   • Key Feature: The unselected positions are zeroed out, but the residual connection from the previous layer preserves the information of the skipped tokens. This allows the model to "re-visit" and re-evaluate token importance in subsequent layers.

Dynamic Token Coverage

To determine which tokens to keep and how many, the paper introduces Dynamic Token Coverage:

• Scoring: A lightweight proxy attention map is computed using recent queries against all keys. Token importance scores are derived by summing attention weights vertically.
• Budget Allocation: Instead of keeping the top-$k$ tokens, the method sorts tokens by ascending importance and prunes the "tail" (least important tokens) until the cumulative mass of pruned tokens exceeds a coverage threshold $\tau$. This treats pruning as structural regularization against attention noise.
• Head-Wise Selection: After determining the global budget, each head independently selects its top-$k$ tokens based on its specific scores.

Sparse Layer Selection

The authors observe that applying sparsity to all layers degrades performance. They introduce Inter-Layer Representation Drift ($R_\ell$) to measure the stability of token representations between layers.

• Layers with low drift (stable representations) are selected for sparsification.
• Layers with high drift (where representations change significantly) are kept dense to preserve model behavior.
• This selection is a one-time preprocessing step per model.

3. Key Contributions

1. Reversible Token-Level Sparsification: A novel "Compress and Decompress" design that reduces computational cost without permanently evicting tokens, allowing dynamic re-evaluation of token importance across layers.
2. Head-Wise Heterogeneity: The method supports independent token selection for each attention head, capturing the fact that different heads prioritize different semantic features.
3. Seamless Compatibility: The approach is fully compatible with existing dense (FlashAttention) and sparse (Minference, FlexPrefill) kernels, acting as a complementary accelerator rather than a replacement.
4. Dynamic Budgeting: A data-driven approach to determine sparsity levels based on token importance distributions and layer-wise stability, rather than fixed ratios.

4. Experimental Results

Experiments were conducted on LLaMA-3.1-8B-Instruct and Mistral-Nemo-12B-Instruct using benchmarks like RULER, InfiniteBench, and Needle-in-a-Haystack.

• Accuracy-Latency Trade-off:
  • Token Sparse Attention consistently improves the speed-accuracy trade-off.
  • At 128K context, it achieves up to $\times 3.23$ attention speedup compared to FlashAttention with less than 1% accuracy degradation.
  • When combined with FlexPrefill, it improves speedup from $\times 2.44$ to $\times 2.76$ while maintaining identical average accuracy (87.27%).
• Comparison with Eviction Methods:
  • Outperforms token eviction methods (PyramidInfer, FastKV, GemFilter) significantly. For example, at 128K, Token Sparse Attention achieves 86.84% accuracy vs. 78.49% for PyramidInfer at similar speedups.
• Scalability:
  • Speedup gains increase with context length (e.g., minimal gain at 4K, substantial at 128K) because attention sparsity increases with sequence length.
• Overhead:
  • The overhead for scoring and index management is minimal, accounting for <11% of total attention latency even at high sparsity levels.

5. Significance

This paper addresses a critical gap in long-context LLM inference by moving away from irreversible token eviction toward dynamic, interleaved sparsification.

• Theoretical Insight: It demonstrates that token importance is highly dynamic across layers and heads, necessitating flexible selection mechanisms rather than static pruning.
• Practical Impact: It provides a "plug-and-play" optimization that can be layered on top of existing state-of-the-art attention kernels (FlashAttention, Minference, etc.) to achieve massive speedups in long-context scenarios without retraining models or sacrificing significant accuracy.
• Future Direction: It establishes a new design point for scalable inference, suggesting that reversible compression is a superior strategy for handling the $O(L^2)$ bottleneck compared to permanent token removal.