CritiPrefill：一种用于LLM预填充加速的基于关键性的分段方法
CritiPrefill: A Segment-wise Criticality-based Approach for Prefilling Acceleration in LLMs
AI总结：
针对LLM长上下文预填充阶段效率低的问题，提出基于关键性的分段方法CritiPrefill，通过分段估计查询关键性并剪枝非关键计算，在128K上下文下实现最高3.0倍加速且质量损失极小。
摘要：
大型语言模型在各个领域取得了显著成功，但高效推理仍受限于注意力机制的二次计算复杂度。推理过程包含预填充和解码阶段。尽管已有多种尝试加速解码，但预填充阶段的低效性，尤其是在长上下文任务中，仍然是一个挑战。在本文中，我们观察到长上下文处理预填充阶段中查询关键性的局部性：相邻的查询token往往关注过去KV缓存中的相似子集。基于这一观察，我们提出了CritiPrefill，一种基于关键性的分段预填充方法。该方法将输入序列的查询和KV缓存划分为段和块，利用分段算法来估计查询关键性。通过在自注意力机制中剪枝查询段和缓存块之间的非关键计算，可以显著加速预填充过程。在多个长上下文数据集上的广泛评估表明，在单个A100 GPU上，对于128K上下文长度，Llama3-8B实现了高达2.7倍的加速，Yi-9B实现了3.0倍的加速，且质量损失极小。
英文摘要：
Large language models have achieved notable success across various domains, yet efficient inference is still limited by the quadratic computation complexity of the attention mechanism. The inference consists of prefilling and decoding phases. Although several attempts have been made to accelerate decoding, the inefficiency of the prefilling phase, especially for long-context tasks, remains a challenge. In this paper, we observe a locality in query criticality during the prefilling phase of long-context processing: adjacent query tokens tend to focus on similar subsets of the past Key-Value (KV) cache. Based on this observation, we propose CritiPrefill, a criticality-based segment-wise prefilling method. This method partitions the input sequence's queries and KV cache into segments and blocks, utilizing a segment-wise algorithm to estimate the query criticality. By pruning non-critical computations between query segments and cache blocks in the self-attention mechanism, the prefilling process can be significantly accelerated. Extensive evaluations on multiple long-context datasets show up to 2.7x speedup on Llama3-8B and 3.0x speedup on Yi-9B for 128K context length on a single A100 GPU, with minimal quality degradation.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
