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

TL;DR

The paper tackles the long-context inference bottleneck from quadratic attention costs $O(L^2)$ by introducing Token Sparse Attention, which for each head selects a token subset of size $L' < L$, compresses Q/K/V to this subset, computes attention, and then reconstructs outputs to the full length $L$ for subsequent layers. This interleaved Compress-then-Decompress design preserves full sequence information and enables head-specific sparsity, while dynamic token coverage governed by a layer-selection heuristic allows on-the-fly adaptation of sparsity budgets. The method is fully compatible with dense kernels like Flash Attention and can complement existing sparse attention kernels, yielding substantial speedups with minimal accuracy loss (up to about 3.23× at 128K context). The approach is supported by a drift-based layer selection mechanism and a dynamic, per-head token scoring scheme, demonstrating practical scalability for long-context inference across multiple models and benchmarks.

Abstract

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.

Figures (7)

• Figure 1: Speedups with Token Sparse Attention. Attention acceleration ratios obtained by applying the proposed Token Sparse Attention (ours) to existing attention acceleration methods. $\tau$ denotes the sparsity level (i.e., higher the sparser).
• Figure 2: Dynamics of Token Importance.(a) Layer-wise overlap of top-k important tokens, showing that token importance shifts significantly across layers. (b) Head-wise token importance rankings within the same layer, illustrating that different attention heads prioritize different subsets of tokens.
• Figure 3: Overview of the proposed Token Sparse Attention. Stage 1 compresses $Q$, $K$, and $V$ by selecting a per-head token subset $S_{H=0}$, yielding compact $\hat{Q}$, $\hat{K}$, and $\hat{V}$ that remain compatible with standard attention kernels. Stage 2 performs attention on the compressed tensors and scatters the resulting outputs back into the full sequence layout before adding the residual connection.
• Figure 4: Sparse Layer Selection.(a) Layer-wise normalized drift measured across different tasks and context lengths. Task A and B correspond to a retrieval task and a summarization task, respectively. (b) Relationship between accuracy and drift under random 3-layer sparsification over 200 runs.
• Figure 5: Accuracy-Speedup Trade-offs.(a) Accuracy-speedup Pareto frontier obtained by sweeping FlexPrefill hyperparameters, comparing with the Token Sparse Attention at token coverage of $\tau$=0.005. (b) Accuracy-speedup trade-off achieved by applying Token Sparse Attention with varying token coverage to FlashAttention (left) and FlexPrefill (right).