通过预稀疏实现近似最优的KV选择以支持长上下文推理

大语言模型（LLM）推理中的核心瓶颈是随着键值（KV）缓存增长带来的计算成本。尽管近似最优的top-k KV选择可以在保持密集注意力质量的同时显著减少计算和带宽需求，但现有稀疏方法通常依赖后验启发式策略，即基于观察到的注意力或代理分数进行选择。这种条件引入了后验偏差，容易扭曲真实token的重要性并遗漏关键token，从而影响长距离推理能力。为解决这一问题，本文提出预稀疏（PrHS）方法，在注意力评分之前选择KV条目，并提供显式的精度控制。通过边际到互信息分析，推导出仅依赖丢弃质量的互信息损失上界，解释了后验启发式的失效模式，并通过提前控制丢弃质量实现可验证的保证。在PrHS中，沿时间、深度和层三个维度实例化了三种正交的预稀疏选择器。在LLaMA和Mistral系列上的大量实验验证了PrHS的有效性。在GSM8K和CoQA数据集上，PrHS将检索开销减少了90%以上，在匹配或更优精度下实现了比HShare高3倍的检索稀疏度。其在LongBench上的平均性能下降低于1%，相比先前稀疏基线减少了约15%的注意力FLOPs，并在NVIDIA A100-80GB GPU上实现了比密集基线快9.9倍的注意力操作延迟和2.8倍的吞吐量。