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

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.

Modern AI assistants are increasingly asked to process very long documents, but the core computation inside these models grows dramatically with text length, creating a fundamental bottleneck. Existing methods address this by skipping text regions in fixed patterns or permanently discarding words deemed unimportant early on, but word importance shifts across layers and varies across different aspects of meaning the model tracks simultaneously. We propose Token Sparse Attention, which selects only the most relevant words at each step, processes that reduced set, and restores the full text before the next step. This reversible design lets the model continuously reconsider which words matter, rather than committing to early decisions it cannot undo. Our method plugs into existing acceleration techniques without modifying them, providing additional speed gains on top. Experiments on long document benchmarks show substantial speedups at very long contexts with less than 1\% accuracy loss.