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