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

TL;DR

Token Sparse Attention introduces a dynamic, token-level sparsification method that improves long-context inference efficiency in large language models while maintaining accuracy.

Contribution

It proposes a novel, fully compatible sparsification mechanism that compresses and decompresses token representations, enhancing speed without significant accuracy loss.

Findings

Achieves up to 3.23x speedup at 128K context length

Maintains less than 1% accuracy degradation

Compatible with existing dense and sparse attention implementations

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.