摘要
大型语言模型在各个领域取得了显著的成功,然而,高效的推理仍然受到注意力机制二次计算复杂度的限制。推理过程包括预填充和解码阶段。尽管已经有一些尝试来加速解码,但预填充阶段,尤其是对于长上下文任务的低效性,仍然是一个挑战。在本文中,我们观察到在长上下文处理的预填充阶段,查询关键性存在局部性:相邻的查询标记往往关注过去键值(KV)缓存的相似子集。基于这一观察,我们提出了CritiPrefill,一种基于关键性的分段预填充方法。该方法将输入序列的查询和KV缓存划分为段和块,利用分段算法来估计查询关键性。通过在自注意力机制中修剪查询段和缓存块之间的非关键计算,预填充过程可以显著加速。在多个长上下文数据集上的广泛评估显示,在单个A100 GPU上,Llama3-8B和Yi-9B在128K上下文长度下分别实现了高达2.7倍和3.0倍的加速,且质量下降最小。