Sparse attention method speeds up long context language model use
CompKV: Compensation-Aware KV Selection for Long-Context LLM Inference
Machine LearningArtificial Intelligence
Summary
Large language models struggle with slow memory use when they read very long text. The authors propose CompKV, a new way to pick important parts of the text to focus on that also considers how to best make up for ignored parts. This approach works by grouping text tokens and choosing groups to minimize errors from ignoring others, based on math that predicts how well compensation will work. Tests showed CompKV speeds up attention calculation by nearly seven times compared to looking at everything.
What this means in practice
- •For ai system engineers: Optimize large language model deployments to handle longer inputs faster by using compensation-aware token selection.
- •For cloud infrastructure teams: Reduce memory and compute costs for long text inference in hosted AI services by deploying efficient sparse attention.
Authors
Zhen Huang, Ruizhe Yao, Danyi Liu, Xinrui Chen, Shuwei Li, Siru Zhong, Zijian Cao, Yushan Lai, Mingming Guo, Weijie Zheng, Haohuan Fu
Abstract
Despite their strong performance, large language models (LLMs) are bottlenecked by KV cache memory traffic during long-context inference. Sparse attention is widely used to accelerate LLM inference by computing exact attention over a selected subset of tokens. To recover the contribution of tokens excluded from exact attention, recent methods apply coarse-grained compensation to the omitted attention tail. However, existing methods typically select tokens based on attention mass and only then compensate for the unselected tokens. This decoupled design overlooks their interaction: selection should prioritize tokens that would leave the largest compensation error if omitted. To address this limitation, we introduce CompKV, the first compensation-aware sparse attention framework that divides tokens into blocks and explicitly optimizes selection for the downstream compensation mechanism. Our theoretical analysis shows that the residual left by block-level mean compensation is governed by both block attention mass and within-block logit variation. We approximate this residual using compact block-level statistics, yielding a deployable selection criterion. We further develop an efficient asynchronous implementation. Experiments on RULER and LongBench-Pro show that CompKV performs best among the evaluated sparse baselines while delivering up to a $6.85\times$ self-attention speedup over full attention.