Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

TL;DR

CritiPrefill, a criticality-based segment-wise prefilling method that partitions the input sequence’s queries and KV cache into segments and blocks, utilizing a segment-wise algorithm to estimate the query criticality.

Abstract

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.

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

TL;DR

CritiPrefill, a criticality-based segment-wise prefilling method that partitions the input sequence’s queries and KV cache into segments and blocks, utilizing a segment-wise algorithm to estimate the query criticality.

Abstract

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.
Paper Structure (15 sections, 3 equations, 5 figures, 1 table, 2 algorithms)

This paper contains 15 sections, 3 equations, 5 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Method
  3. Query Criticality in Attention Mechanism
  4. Locality Pattern of Query Criticality
  5. Proposed Methods
  6. Estimating segment-wise query ciriticality
  7. Pruning Self-attention Mechnisum
  8. Experiment
  9. Settings
  10. Dataset
  11. Baseline
  12. Parameter Settings
  13. Evaluations cross multiple scenarios
  14. Evaluations on Needle-in-a-Haystack Test
  15. Conclusion

Figures (5)

  • Figure 1: Acceleration of CritiPrefill
  • Figure 2: Framework of CritiPrefill
  • Figure 3: Locality Pattern
  • Figure 4: CritiPrefill on Needle-in-a-Haystack Test. This test involves inserting an answer within a large context and evaluating the retrieval ability in response to the corresponding question. The Average Score is derived from averaging scores across various context lengths (x-axis) and insertion depths (y-axis), with a maximum score of 10.
  • Figure 5: Prefilling Time on Need-in-a-Haystack Test