稀疏注意力作为范围搜索问题：面向 KV 缓存的推理高效索引

稀疏注意力通过选择键值条目的子集来提高 LLM 推理效率，但代价是潜在的准确性下降。特别是，省略关键 KV 条目可能会导致模型输出出现重大错误。现有方法通常在固定或自适应 词元预算 下运行，并提供经验鲁棒性或部分理论保证，但它们不能确保解码步骤中的零假阴性，特别是因为相关标记集既依赖于查询又依赖于步骤。我们的经验观察证实，即使缺少一个关键密钥也会导致严重的错误峰值，特别是在长推理任务中，其中一组重要标记在解码过程中会发生变化。这一观察结果激发了对索引方法的需求，该方法能够动态适应解码步骤中的这些变化，同时保证完全召回高于特定阈值的相关密钥。我们通过将稀疏注意力重新表述为半空间范围搜索问题来解决这一挑战。然而，现有的范围搜索索引由于其计算和实现开销而不适合现代 LLM 推理。为了克服这个问题，我们引入了 Louver，一种专为高效 KV 缓存检索而定制的新颖索引结构。 Louver (i) 在理论和实践中保证指定阈值的零误报，(ii) 重量轻，可以集成到现有的 LLM 管道中，(iii) 结合了针对 CPU 和 GPU 执行的硬件感知优化。我们的实验表明，Louvre 在准确性和运行时间方面都优于先前的稀疏注意力方法，并且比高度优化的密集注意力（例如 FlashAttention）更快。这些结果凸显了召回保证是稀疏注意力的一个关键且被忽视的维度，并为构建有理论依据的高效 KV 缓存索引开辟了新方向。