PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference (2405.12532v2)

Abstract: LLMs have shown remarkable comprehension abilities but face challenges in GPU memory usage during inference, hindering their scalability for real-time applications like chatbots. To accelerate inference, we store computed keys and values (KV cache) in the GPU memory. Existing methods study the KV cache compression to reduce memory by pruning the pre-computed KV cache. However, they neglect the inter-layer dependency between layers and huge memory consumption in pre-computation. To explore these deficiencies, we find that the number of crucial keys and values that influence future generations decreases layer by layer and we can extract them by the consistency in attention weights. Based on the findings, we propose PyramidInfer, a method that compresses the KV cache by layer-wise retaining crucial context. PyramidInfer saves significant memory by computing fewer keys and values without sacrificing performance. Experimental results show PyramidInfer improves 2.2x throughput compared to Accelerate with over 54% GPU memory reduction in KV cache.

Enhancing LLM Efficiency with PyramidInfer

Introduction

If you're familiar with LLMs like GPT-3 or LLaMA, you know that while they exhibit strong capabilities in NLP, they also come with significant constraints, particularly around inference efficiency. These constraints often stem from GPU memory demands during real-time applications, such as chatbots. A recent research paper introduces PyramidInfer, an approach designed to make LLMs more memory-efficient without compromising performance. Let's break down the key concepts and findings from this paper.

The Challenge of GPU Memory in LLM Inference

Inference in LLMs is less about training but more about handling two main components: model parameters and the KV cache. Model parameters are basically the learned weights, but the KV cache (keys and values previously computed in the attention mechanism) takes care of reusing these computations to avoid repeating them. The catch? The KV cache can consume a massive amount of memory.

Example: For a model with 7 billion parameters, the parameters might take up 14 GB, but the KV cache could demand around 72 GB!

This leads to an issue where the throughput and scalability of LLMs are restricted by how much KV cache can fit in your GPU memory.

What is PyramidInfer?

PyramidInfer is an innovative method that tackles GPU memory consumption by compressing the KV cache more efficiently. Instead of simply compressing after the KV cache is computed (as many existing methods do), PyramidInfer proactively reduces the KV cache during both the prefill and generation phases.

Key Concepts:

1. Inference Context Redundancy (ICR):
   • During inference, not all tokens need to predict the next token, leading to redundant information in the KV cache.
   • Conclusion: We can focus on computing less redundant key-values to save memory.
2. Recent Attention Consistency (RAC):
   • Recent tokens tend to pay attention to the same context, defined as Pivotal Context (PvC).
   • Conclusion: Leveraging this consistency can help in selecting crucial parts of the KV cache, further aiding compression.

How PyramidInfer Works

PyramidInfer applies a layered approach to retaining essential context in the KV cache, implemented in two main phases:

1. Prefill Phase:
   • It computes only the significant keys and values (PvCs) from the prompt.
   • By averaging attention from recent tokens, it identifies essential context tokens layer-wise, forming a "pyramid" of key-values—denser at the base and thinner at the top.
2. Generation Phase:
   • Uses a sliding window approach for recent tokens to update the PvCs continually.
   • This ensures more GPU memory is saved, while still maintaining high inference quality.

Experimental Results

Strong Numerical Results:

• Throughput: PyramidInfer boosts throughput by 2.2 times compared to Accelerate, a widely-used method.
• Memory Savings: Shows a reduction of GPU memory usage by over 54% in the KV cache.

Versatile Application:

• Tasks: Works efficiently across a broad set of NLP tasks, including language understanding (MMLU, BBH), mathematical reasoning (GSM8K), coding (HumanEval), conversation (MT-Bench), and long-context tasks (LEval).
• Models: Demonstrates compatibility with various models like LLaMA 2, LLaMA 2-Chat, Vicuna, and CodeLLaMA.

Implications and Future Directions

Practical Implications:

• Scalability: PyramidInfer facilitates deploying LLMs in environments with stringent memory constraints, making technologies like chatbots more accessible and responsive.
• Cost Reduction: By significantly reducing GPU memory requirements, it enables more efficient use of existing hardware, potentially lowering infrastructure costs.

Theoretical Implications:

• Attention Mechanism Optimization: Sets the stage for deeper exploration into layer-wise importance and redundancy in attention mechanisms.
• Future Research: Opens avenues for hybrid methods combining PyramidInfer with other efficiency-boosting techniques, like Deepspeed.

Conclusion

PyramidInfer presents an effective solution to one of the biggest bottlenecks in LLM deployment—GPU memory usage. By compressing the KV cache efficiently during both crucial phases of inference, it significantly enhances throughput and maintains performance. As the demand for real-time applications grows, methods like PyramidInfer will undoubtedly play a pivotal role in optimizing the deployment of large-scale LLMs.

For more details, check out the PyramidInfer codebase and consider experimenting with it to see how it can optimize your own LLM applications!