STS: Efficient Sparse Attention with Speculative Token Sparsity
This work addresses the memory and computational bottlenecks of LLM inference for long-context agentic applications, offering a practical speedup without retraining.
STS introduces a sparse attention mechanism that uses a smaller draft model to predict important tokens for a larger target model, achieving 2.67x speedup at ~90% sparsity on NarrativeQA with negligible accuracy loss.
The quadratic complexity of attention imposes severe memory and computational bottlenecks on Large Language Model (LLM) inference. This challenge is particularly acute for emerging agentic applications that require processing multi-million token sequences. We propose STS, a sparse attention mechanism that requires no model retraining. STS leverages the key insight that tokens identified as important by a smaller draft model are highly predictive of important tokens for a larger target model. By integrating into speculative decoding frameworks, STS repurposes the draft model's attention scores to dynamically construct a token-and-head-wise sparsity mask. This mask effectively prunes the expensive attention computation in the target LLM. Our evaluation shows that STS achieves a 2.67x speedup operating at approximately 90% sparsity on representative benchmark NarrativeQA, maintaining negligible accuracy degradation compared to dense attention. STS establishes a new state-of-the-art on the sparsity-accuracy trade-off, outperforming prior techniques by enabling higher sparsity levels for a given accuracy budget.