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

• URL: http://arxiv.org/abs/2409.12490v2
• Date: Mon, 23 Sep 2024 02:24:33 GMT
• Title: CritiPrefill: A Segment-wise Criticality-based Approach for Prefilling Acceleration in LLMs
• Authors: Junlin Lv, Yuan Feng, Xike Xie, Xin Jia, Qirong Peng, Guiming Xie,
• Abstract summary: We propose CritiPrefill, a criticality-based segment-wise prefilling method for long-context processing. CritiPrefill partitions the input sequence's queries and KV cache into segments and blocks, utilizing a segment-wise algorithm to estimate the query criticality. 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.
• Score: 8.649971923487835
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• TokenSelect: Efficient Long-Context Inference and Length Extrapolation for LLMs via Dynamic Token-Level KV Cache Selection [23.20856449846164]
  TokenSelect is a model-agnostic, training-free method for efficient and accurate long-context inference. A comprehensive evaluation of TokenSelect demonstrates up to 23.84x speedup in attention and up to 2.28x acceleration in end-to-end latency.
  arXiv  Detail & Related papers  (2024-11-05T07:56:24Z)
• Fine-grained Attention I/O Complexity: Comprehensive Analysis for Backward Passes [24.075230599299772]
  Large Language Models (LLMs) have demonstrated remarkable capabilities in processing long-context information. The quadratic complexity of attention with respect to sequence length poses significant computational challenges. This paper presents a comprehensive analysis of the I/O complexity for attention mechanisms, focusing on backward passes.
  arXiv  Detail & Related papers  (2024-10-12T07:01:30Z)
• ThinK: Thinner Key Cache by Query-Driven Pruning [63.13363917871414]
  Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. We propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels.
  arXiv  Detail & Related papers  (2024-07-30T17:59:08Z)
• KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches [52.02764371205856]
  Long context capability is a crucial competency for large language models (LLMs) This work provides a taxonomy of current methods and evaluating 10+ state-of-the-art approaches across seven categories of long context tasks.
  arXiv  Detail & Related papers  (2024-07-01T17:59:47Z)
• Training-Free Exponential Context Extension via Cascading KV Cache [49.608367376911694]
  We introduce a novel mechanism that leverages cascading sub-cache buffers to selectively retain the most relevant tokens. Our method reduces prefill stage latency by a factor of 6.8 when compared to flash attention on 1M tokens.
  arXiv  Detail & Related papers  (2024-06-24T03:59:17Z)
• In-Context Former: Lightning-fast Compressing Context for Large Language Model [48.831304302467004]
  In this paper, we propose a new approach to compress the long input contexts of Transformer-based large language models (LLMs) We use the cross-attention mechanism and a small number of learnable digest tokens to condense information from the contextual word embeddings. Experimental results indicate that our method requires only 1/32 of the floating-point operations of the baseline during compression and improves processing speed by 68 to 112 times.
  arXiv  Detail & Related papers  (2024-06-19T15:14:55Z)
• Prepacking: A Simple Method for Fast Prefilling and Increased Throughput in Large Language Models [48.592730159983276]
  Prefilling is the computation of the key-value cache for input tokens in the prompt prior to autoregressive generation. For longer input prompt lengths, prefilling incurs a significant overhead on decoding time. We propose Prepacking, a simple yet effective method to optimize prefilling computation.
  arXiv  Detail & Related papers  (2024-04-15T07:49:10Z)
• Bifurcated Attention: Accelerating Massively Parallel Decoding with Shared Prefixes in LLMs [39.16152482491236]
  Bifurcated attention is a method designed to enhance language model inference in shared-context batch decoding scenarios. Our approach addresses the challenge of redundant memory IO costs, a critical factor contributing to latency in high batch sizes and extended context lengths.
  arXiv  Detail & Related papers  (2024-03-13T16:30:57Z)
• SubGen: Token Generation in Sublinear Time and Memory [48.35076900702408]
  Large language models (LLMs) have extensive memory requirements for token generation. In this work, we focus on developing an efficient compression technique for the KV cache. We have devised a novel caching method with sublinear complexity, employing online clustering on key tokens and online $ell$ sampling on values. Not only does this algorithm ensure a sublinear memory footprint and sublinear time complexity, but we also establish a tight error bound for our approach.
  arXiv  Detail & Related papers  (2024-02-08T22:17:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.