中文

QUOKA:面向查询的高效 LLM 前缀注意力 KV 选择

机器学习 2026-02-10 v1 人工智能

摘要

我们提出了 QUOKA:一种面向查询的 KV 选择方法,用于加速 transformer 推理中的注意力计算,这是一种无需训练且与硬件无关的稀疏注意力算法,可加速分块前缀注意力。虽然许多查询在注意力算子中聚焦于较小的一组键值,但我们观察到与平均查询向量余弦相似性较低的查询与更多键值之间的相互作用更强,对最终注意力 logits 贡献最大。通过优先处理这些余弦相似性较低的查询,可以更准确地近似前向注意力的行为。QUOKA 利用这一观察,通过 (1) 首先保留一小组代表性查询,(2) 然后选择与这些查询最匹配的键值,来加速注意力计算。在 Needle-In-A-Haystack、LongBench、RULER 和 Math500 等实验中,我们表明 QUOKA 在保持接近基准准确率的同时,实现了对 Nvidia GPU 上的 3 倍时间到第一个 token 的减少,5 倍注意力加速,以及在 Intel Xeon CPU 上可达近 7 倍的加速,仅使用 88% 的键值对即可完成注意力计算。

关键词

引用

@article{arxiv.2602.08722,
  title  = {QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill},
  author = {Dalton Jones and Junyoung Park and Matthew Morse and Mingu Lee and Chris Lott and Harper Langston},
  journal= {arXiv preprint arXiv:2602.08722},
  year   = {2026}
}