通过分段评估关键性，加速大模型长文本推理的预填充阶段。
CritiPrefill: A Segment-wise Criticality-based Approach for Prefilling Acceleration in LLMs
- 将输入和缓存分段，按段估算查询重要性
- 在128K上下文下实现最高3.0倍加速，损失极小
- 适合长文本生成、高并发推理场景使用
大型语言模型在多个领域取得显著成果，但其推理效率仍受限于注意力机制的二次计算复杂度。推理过程包括预填充和解码两个阶段。尽管已有研究尝试加速解码，长上下文任务中预填充阶段的低效问题依然突出。本文观察到，在长上下文处理的预填充阶段，查询的重要性存在局部性：相邻查询令牌倾向于关注过去键值（KV）缓存的相似子集。基于此，提出CritiPrefill方法——一种基于重要性的分段式预填充策略。该方法将输入序列的查询与KV缓存划分为段和块，并采用分段算法估计查询重要性。通过在查询段与缓存块之间剪枝非关键计算，显著加速预填充过程。在多个长上下文数据集上的大量实验表明，在单张A100 GPU上，对Llama3-8B可实现最高2.7倍加速，对Yi-9B可达3.0倍加速，且质量下降极小。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。