CritiPrefill: A Segment-wise Criticality-based Approach for Prefilling Acceleration in LLMs

📄 arXiv: 2409.12490v2 📥 PDF

作者: Junlin Lv, Yuan Feng, Xike Xie, Xin Jia, Qirong Peng, Guiming Xie

分类: cs.CL, cs.AI, cs.LG

发布日期: 2024-09-19 (更新: 2024-09-23)

💡 一句话要点

CritiPrefill：基于分段关键性的LLM预填充加速方法

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 预填充加速 长文本处理 注意力机制 查询关键性 分段剪枝 模型推理

📋 核心要点

  1. 长文本处理中，LLM预填充阶段计算量大，效率低，成为推理瓶颈。
  2. CritiPrefill利用查询关键性的局部性，通过分段剪枝非关键计算加速预填充。
  3. 实验表明，CritiPrefill在长文本任务上显著加速，同时保持了模型性能。

📝 摘要（中文）

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

🔬 方法详解

问题定义：论文旨在解决大型语言模型（LLMs）在长上下文推理中预填充阶段效率低下的问题。现有的方法在解码阶段做了很多优化，但预填充阶段的计算复杂度仍然很高，特别是对于长文本输入，这严重限制了LLMs的应用。

核心思路：论文的核心思路是观察到预填充阶段查询关键性的局部性。这意味着相邻的查询token倾向于关注相似的KV缓存子集。基于此，可以通过识别和剪枝不重要的计算来加速预填充过程。

技术框架：CritiPrefill方法主要包含以下几个阶段：1) 将输入序列的查询和KV缓存划分为段和块；2) 使用分段算法估计每个查询段的关键性；3) 基于关键性评估结果，剪枝查询段和缓存块之间非关键的自注意力计算。整体框架旨在减少冗余计算，提高预填充效率。

关键创新：该方法最重要的创新点在于利用了查询关键性的局部性，并提出了分段剪枝策略。与传统的全局剪枝方法相比，CritiPrefill能够更精细地控制剪枝粒度，从而在加速的同时保持模型性能。此外，分段算法的设计也是一个关键创新，它能够在保证效率的前提下，准确地估计查询关键性。

关键设计：CritiPrefill的关键设计包括：1) 段和块的大小设置，需要在计算效率和关键性估计的准确性之间进行权衡；2) 关键性评估算法，需要能够快速准确地识别出重要的KV缓存块；3) 剪枝策略，需要避免过度剪枝导致模型性能下降。具体的参数设置和算法细节在论文中进行了详细描述，但此处未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，CritiPrefill在Llama3-8B和Yi-9B模型上实现了显著的加速效果。在128K上下文长度下，Llama3-8B加速高达2.7倍，Yi-9B加速高达3.0倍，且模型性能下降很小。这些结果证明了CritiPrefill在长文本处理中的有效性。

🎯 应用场景

CritiPrefill具有广泛的应用前景，尤其是在需要处理长文本的场景中，例如长篇文档摘要、长对话生成、代码生成等。该方法可以显著降低LLM的推理成本，使其更容易部署在资源受限的设备上。未来，CritiPrefill可以与其他加速技术相结合，进一步提高LLM的推理效率。

📄 摘要（原文）

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.