通过筛选关键查询，让大模型预填充更快更准。
QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill
- 只保留与平均查询差异大的关键查询，减少计算量。
- 在多个数据集上实现3倍首字延迟降低，最高7倍加速。
- 无需训练，适配不同硬件，适合部署优化场景。
我们提出QUOKA：一种面向高效注意力计算的查询导向键值选择方法，适用于分块预填充场景下的Transformer推理加速。尽管多数查询仅关注少量键，但研究发现与平均查询余弦相似度较低的查询会与更多键产生强交互，对最终注意力结果贡献最大。因此，通过优先保留此类低相似度查询，并从中筛选最匹配的键，可近似全注意力行为。QUOKA首先保留少量代表性查询，再子选择与之对齐的键。在Needle-In-A-Haystack、LongBench、RULER和Math500测试中，该方法实现了3倍首字延迟降低，Nvidia GPU上注意力速度提升5倍，Intel Xeon CPU上接近7倍加速，且仅使用88%的键值对即可保持接近基线精度。
原文摘要 · Abstract (English)
We present QUOKA: Query-oriented KV selection for efficient attention, a training-free and hardware agnostic sparse attention algorithm for accelerating transformer inference under chunked prefill. While many queries focus on a smaller group of keys in the attention operator, we observe that queries with low cosine similarity with respect to the mean query interact more strongly with more keys and have the greatest contribution to final attention logits. By prioritizing these low cosine similarity queries, the behavior of full attention during the prefill stage can be closely approximated. QUOKA leverages this observation, accelerating attention by (1) first retaining a small set of representative queries and (2) then subselectin the keys most aligned with those queries. Through experiments on Needle-In-A-Haystack, LongBench, RULER, and Math500, we show that, while realizing a 3x reduction in time-to-first-token, 5x speedup in attention on Nvidia GPUs and up to nearly a 7x speedup on Intel Xeon CPUs, QUOKA achieves near-baseline accuracy, utilizing 88% fewer key-value pairs per attention evaluation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。