PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Dongwon Jo, Beomseok Kang, Jiwon Song, Jae-Joon Kim

arXiv:2602.03216v3cs.CLcs.LG

TL;DR

Long-context inference is bottlenecked by quadratic attention and limited token-selection flexibility. Token Sparse Attention dynamically compresses and decompresses per-head token sets, consistently improving accuracy–latency trade-offs, including 87.3% accuracy with ×2.8 speedup when combined with FlexPrefill.

  • Problem

    Quadratic attention remains a bottleneck for long-context inference, while layer- or head-insensitive token eviction limits finer-grained token selection.

  • Method

    Token Sparse Attention dynamically selects informative tokens per head, computes attention in compressed space, and decompresses outputs so token relevance can be reconsidered across layers.

  • Results

    Token Sparse Attention consistently improves accuracy–latency trade-offs; combined with FlexPrefill, it reaches 87.3% accuracy with ×2.8 speedup versus ×2.4 for FlexPrefill.

  • Takeaways & Limitations

    Dynamic, interleaved token sparsification complements dense and sparse attention kernels as a strategy for efficient long-context inference.

  • Takeaways & Limitations

    Applying Token Sparse Attention across all layers causes substantial performance degradation, requiring selective sparse-layer placement.

Abstract

from arXiv · show

The quadratic complexity of attention remains the central bottleneck in long-context inference for large language models. Prior acceleration methods either sparsify the attention map with structured patterns or permanently evict tokens at specific layers, which can retain irrelevant tokens or rely on irreversible early decisions despite the layer-/head-wise dynamics of token importance. In this paper, we propose Token Sparse Attention, a lightweight and dynamic token-level sparsification mechanism that compresses per-head $Q$, $K$, $V$ to a reduced token set during attention and then decompresses the output back to the original sequence, enabling token information to be reconsidered in subsequent layers. Furthermore, Token Sparse Attention exposes a new design point at the intersection of token selection and sparse attention. Our approach is fully compatible with dense attention implementations, including Flash Attention, and can be seamlessly composed with existing sparse attention kernels. Experimental results show that Token Sparse Attention consistently improves accuracy-latency trade-off, achieving up to $\times$3.23 attention speedup at 128K context with less than 1% accuracy degradation. These results demonstrate that dynamic and interleaved token-level sparsification is a complementary and effective strategy for scalable long-context inference.

1. Introduction

Token Sparse Attention addresses the quadratic prefill bottleneck in long-context inference by selecting informative tokens dynamically at the head level. It complements existing sparse attention methods, improving the accuracy-speedup trade-off through heterogeneous token- and block-level sparsity.

  • Motivation: Attention complexity grows quadratically with context length during prefill, creating a fundamental inference bottleneck despite hardware-aware optimizations.The cited passage identifies FlashAttention as an optimization but retains the underlying long-context challenge.
  • Limitations of Prior Methods: Block-sparse methods may preserve less informative tokens, while eviction-based methods cannot capture layer- or head-specific shifts in token importance.These limitations reduce the accuracy-speedup benefits of operating at token-level granularity.
  • Token Sparse Attention: Token Sparse Attention selects a small per-head subset of informative tokens, L′<L, and performs attention in the compressed R^L′×L′ space.The selected subset can differ across attention heads, but naive selection prevents remaining tokens from being revisited in later layers.
  • Complementarity: Token Sparse Attention can prune irrelevant tokens before block-sparse or Λ-shape computations, making token sparsity complementary to existing sparse attention.This compatibility enables heterogeneous granularity rather than replacing existing sparse methods.
  • Results: 87.3% accuracy with ×2.8 speedup is achieved by combining FlexPrefill with Token Sparse Attention, versus ×2.4 speedup for standard FlexPrefill at the same accuracy.FlexPrefill is block-sparse, whereas Token Sparse Attention adds token-level sparsity.

2. Method

Token Sparse Attention addresses layer- and head-dependent token importance by selecting tokens independently per head, applying attention to compressed QKV tensors, and restoring outputs to the full sequence after each step. Dynamic Token Coverage determines the sparsity budget and selected tokens at inference time, while layer-wise sparsification is constrained because applying it everywhere substantially degrades performance.

  • Motivation: Token importance shifts substantially across layers and varies across attention heads, making permanent early-layer eviction and unified layer-level token sets unreliable.Adjacent-layer overlap decreases rapidly with layer distance, while different heads exhibit distinct token-ranking patterns.
  • Compression and Decompression: Token Sparse Attention independently selects per-head token subsets, gathers reduced Q, K, and V tensors, and applies attention at sequence length L′ ≪ L.The compressed tensors remain dense and contiguous, preserving compatibility with optimized kernels such as Flash Attention.
  • Compression and Decompression: Attention cost decreases from O(L^2d) to O(L′^2d) by operating on the compressed tensors.The reduced output contains context-aware representations for the selected tokens.
  • Compression and Decompression: The compressed output is scattered into a zero-initialized full-sequence tensor, restoring original dimensions while leaving unselected positions zero.This interleaved decompression preserves future selection space and allows token importance to be reconsidered across later layers and heads.
  • Dynamic Token Coverage: Dynamic Token Coverage estimates head-wise token scores, aggregates them across heads, and removes the least-important tokens whose cumulative mass exceeds coverage threshold τ.Scores are derived from a lightweight proxy attention map computed using recent queries and all keys, with a Triton kernel reducing memory I/O overhead.
  • Layer Selection: Applying Token Sparse Attention across all layers causes substantial performance degradation, motivating layer selection for sparsification.The paper evaluates layer stability using hidden-state drift, where lower drift indicates smaller representational changes.

3. Experiments

Experiments across two instruction-tuned models and long-context benchmarks show that Token Sparse Attention preserves baseline accuracy while improving attention efficiency and complementing dense and sparse methods. Its benefits increase with context length, remain robust under aggressive token coverage, and outperform fixed-sparsity and token-eviction alternatives at comparable efficiency.

  • Experimental Setup: Experiments use LLaMA-3.1-8B-Instruct and Mistral-Nemo-12B-Instruct, primarily evaluated on RULER and InfiniteBench for long-context understanding and retrieval.Additional LongBench and Needle-in-a-Haystack evaluations are provided in Appendix A.2.
  • Complementary Acceleration: 87.27% average accuracy is maintained when Token Sparse Attention augments FlexPrefill on LLaMA-3.1-8B-Instruct, matching vanilla FlexPrefill across context lengths.Similar preservation trends are observed with FlashAttention and Minference, while deviations on Mistral-Nemo-12B-Instruct remain consistently small.
  • Complementary Acceleration: Token Sparse Attention preserves InfiniteBench performance with FlashAttention, Minference, and FlexPrefill, producing only marginal accuracy differences while maintaining the baselines’ efficiency trends.This compatibility holds for both evaluated models.
  • Accuracy-Speedup Trade-offs: Accuracy degradation remains within 1% at higher sparsity levels, while increasing token coverage produces higher attention speedups for both FlashAttention and FlexPrefill.The method also provides complementary gains beyond aggressively tuning FlexPrefill’s sparsity parameter at comparable accuracy levels.
  • Sparsity and Latency: Attention speedup increases with context length, becoming substantially larger at 128K and 256K as attention sparsity rises and attention computation dominates latency.At 128K, Token Sparse Attention’s added overhead remains less than 11% of total attention latency across all layers, including scoring, indexing, QKV compression, and output decompression.
  • Dynamic Sparsity and Token Eviction: Dynamic sparsity consistently achieves higher RULER average accuracy than fixed sparsity at comparable 128K speedups, while Token Sparse Attention attains the highest average RULER accuracy against matched token-eviction budgets.The reported advantages are attributed to layer-wise dynamic allocation, reversible interleaving, and head-wise token selection.

4. Related Works

Related work addresses long-context inference through system-level attention optimization, token eviction and prompt compression, quantization, and KV-cache reduction during decoding. These approaches target computational, memory, and bandwidth costs across inference stages.

  • Long-Context Inference for LLMs: Long-context inference faces substantial computational and memory challenges, with quadratic attention overhead increasing as context length grows during prefill.System-level approaches include FlashAttention and FlashInfer.
  • Prefill Acceleration: Prefill acceleration methods such as FastKV, GemFilter, and PyramidInfer select important tokens at specific layers and compress corresponding hidden states.Prompt compression methods instead operate at the prompt level rather than inside the model.
  • Prefill Acceleration: ×1.49, ×1.50, ×1.53, and ×1.51 are the reported speedups for PyramidInfer, FastKV, GemFilter, and Ours, respectively, compared to FlashAttn.The comparison uses RULER accuracy on LLaMA-3.1-8B-Instruct, with methods set to similar speedup levels.
  • Long-Context Inference for LLMs: Quantization methods such as GPTQ, SmoothQuant, and AWQ reduce memory load or compress weights and activations for low-precision computation.These methods are presented as another approach to long-context inference efficiency.
  • Decoding Acceleration: During decoding, KV-cache memory and bandwidth costs become the bottleneck, motivating eviction methods that retain only important tokens.Examples include H2O, SnapKV, AdaKV, and HeadKV.

5. Conclusion

Token Sparse Attention provides dynamic, reversible token-level sparsification for efficient long-context inference. It reduces attention computation while preserving token reconsideration across layers and heads, and remains compatible with dense and sparse attention kernels.

  • Token Sparse Attention is a dynamic and reversible token-level sparsification mechanism for efficient long-context inference.
  • It reduces attention computation without permanently removing tokens, enabling token relevance to be re-evaluated across layers and heads.
  • The design is compatible with existing dense and sparse attention kernels, enabling seamless composition with prior acceleration methods.

A. Appendix · A.1. Additional Models and Datasets Details

The appendix evaluates Token Sparse Attention with two publicly available instruction-tuned models designed for 128K contexts. It uses RULER and InfiniteBench to assess long-context retrieval, reasoning, and other capabilities across synthetic and realistic tasks.

  • A.1. Additional Models and Datasets Details: The experiments use LLaMA-3.1-8B-Instruct and Mistral-Nemo-12B-Instruct, both designed to operate with 128K context windows.These are the evaluated long-context models.
  • A.1. Additional Models and Datasets Details: The evaluated model checkpoints are publicly available through the listed Hugging Face repositories.The appendix provides links for both model checkpoints.
  • A.1. Additional Models and Datasets Details: RULER is a configurable synthetic benchmark that tests long-context capabilities across varying sequence lengths and task difficulties.It extends the Needle-in-a-Haystack setup beyond a single retrieval scenario.
  • A.1. Additional Models and Datasets Details: RULER comprises 13 tasks organized into retrieval-style, aggregation, multi-hop tracing, and question-answering categories.Examples include CWE and FWE aggregation tasks and VT multi-hop tracing.
  • A.1. Additional Models and Datasets Details: InfiniteBench evaluates whether language models can understand and reason over extremely long contexts exceeding 100K tokens.The benchmark targets contexts beyond the scale of conventional long-context evaluation.
  • A.1. Additional Models and Datasets Details: InfiniteBench covers retrieval, reasoning, code understanding, mathematical computation, dialogue, and summarization across multiple domains.Its task collection includes both synthetic and realistic scenarios.

A.2. Additional Benchmarks

Additional benchmarks show that Token Sparse Attention preserves long-context accuracy across precise retrieval and diverse task settings. It also improves accuracy when combined with FlexPrefill while introducing minimal changes on LongBench.

  • Needle-in-a-Haystack: Needle-in-a-Haystack evaluation on LLaMA-3.1-8B-Instruct shows strong accuracy preservation under long-context settings.The results are presented in Figure 7.
  • Needle-in-a-Haystack: Combined with FlexPrefill, Token Sparse Attention consistently achieves higher accuracy than FlexPrefill alone.The results suggest reversible token-level sparsification removes irrelevant context while preserving critical retrieval signals.
  • LongBench: Across LongBench’s broad set of long-context tasks, Token Sparse Attention consistently preserves baseline performance with minimal accuracy changes.LongBench includes question answering, summarization, few-shot learning, synthetic reasoning, and code understanding.

A.3. Additional Ablation Study

The ablations validate recent-query scoring and drift-based layer selection as effective design choices for Token Sparse Attention. They also identify δ=0.5 as a consistent accuracy–speedup sweet spot for selecting sparse layers.

  • Token Scoring Methods: Recent-query scoring outperforms random query selection and query-only pooling for token-importance estimation.Random selection causes noticeable accuracy degradation, while query-only pooling partially recovers accuracy but remains inferior to recent-query scoring without significant computational overhead.
  • Stability of Drift-Based Layer Selection: Low-drift layers consistently achieve the highest accuracy, whereas high-drift layers cause substantial degradation when sparsified.This pattern holds for both LLaMA-3.1-8B-Instruct and Mistral-Nemo-12B-Instruct, supporting drift as a model-agnostic layer-selection indicator.
  • Delta Sweep for Sparse Layer Selection: δ=0.5 provides a consistent accuracy–speedup sweet spot for controlling sparse-layer eligibility on RULER with LLaMA-3.1-8B-Instruct.Larger δ values include more layers with higher representation drift, but values beyond 0.5 provide only marginal additional benefits.

A.4. Prefill Latency

Prefill latency captures end-to-end gains because all methods have identical decoding TPOT under dense attention with full KV cache. At 128K, prefill improvements closely track attention speedups on both GPUs, while at 8K scoring overhead is more visible but remains small for Token Sparse Attention.

  • Latency measurement: 89ms TPOT is identical across methods at 128K with 128 generation tokens because all use dense attention and full KV cache during decoding.This makes prefill latency the relevant measure for end-to-end gains.
  • Hardware results: At 128K, prefill latency improvements closely follow attention-level speedups on both A100 and A6000 GPUs.The comparison is reported in Table 10 for A100 and Table 11 for A6000.
  • Sequence-length effects: At 8K, reduced attention sparsity makes scoring overhead more visible, but Token Sparse Attention’s added overhead remains small relative to other baselines.The method prioritizes improving efficiency in long-context inference.

A.5. Comparison with Additional Baselines · A.6. Comparison with Non-Eviction Method · A.7. Limitation

Token Sparse Attention composes orthogonally with structured sparse attention methods, improving speedup while preserving comparable accuracy. Against OrthoRank, it achieves higher RULER accuracy under comparable prefill speedup, while its short-context gains are smaller and aggressive sparsification increases error risk.

  • A.5. Comparison with Additional Baselines: Token Sparse Attention operates at the Q, K, V tensor level, enabling composition with SeerAttention and X-Attention without modifying their sparse kernels.Both baselines apply structured sparsity at the attention-map level.
  • A.5. Comparison with Additional Baselines: x2.47 versus x2.19 attention speedup results from composing Token Sparse Attention with SeerAttention while maintaining comparable accuracy.With X-Attention (S=16, τ=0.9), speedup increases from x2.72 to x3.49 with negligible accuracy change.
  • A.6. Comparison with Non-Eviction Method: OrthoRank selects tokens before attention from normalized hidden states, whereas Token Sparse Attention compresses per-head Q, K, V using the attention map as its scoring signal.These methods therefore differ in both compression level and token-scoring mechanism.
  • A.6. Comparison with Non-Eviction Method: OrthoRank uses one token set across all heads in a layer, preventing it from modeling head-wise differences in attention-based token importance.The paper reports significantly different token-importance rankings across heads within the same layer.
  • A.6. Comparison with Non-Eviction Method: 87.02% and 86.84% average RULER accuracy for Token Sparse Attention exceed OrthoRank’s 79.36% and 78.38% under comparable prefill speedup.The comparison uses prefill speedup because OrthoRank also skips FFN computation for unselected tokens; at 128K, the accuracy difference is approximately 17%.
  • A.7. Limitation: Token Sparse Attention delivers smaller efficiency gains at short context lengths, where attention does not yet dominate end-to-end latency.The method is primarily designed for long-context inference, its main target regime.
  • A.7. Limitation: Overly aggressive token sparsification can increase the risk of error, limiting how far token selection can be pushed.This limitation accompanies the reduced short-context benefit described for the method.
Loading 2602.03216v3…