论文

CompKV:补偿感知的稀疏注意力框架,长上下文推理最高提速 6.85 倍

CompKV: Compensation-Aware KV Selection for Long-Context LLM Inference

精选理由

一篇新论文 CompKV,重新设计了稀疏注意力里的 token 选择逻辑,长上下文推理自注意力最高快 6.85 倍,做推理优化的可以看看。

CompKV 是首个补偿感知的稀疏注意力框架,针对 KV cache 在长上下文推理中的内存瓶颈而设计。与先按注意力质量选 token 再做补偿的现有方法不同,CompKV 把 token 分块,优先选择若省略会带来最大补偿误差的部分。论文的理论分析指出块级均值补偿的残差由块注意力质量和块内 logit 方差共同决定,并给出了可部署的选择准则。在 RULER 和 LongBench-Pro 基准上,CompKV 优于所评估的稀疏基线,自注意力最高加速 6.85 倍。

原文 · arXiv cs.AI

CompKV: Compensation-Aware KV Selection for Long-Context LLM Inference

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.