PaperlayerView exact revision on arXiv

Source-linked AI summary

PyramidInfer: Pyramid KV Cache Compression for High-throughput LLM Inference

Dongjie Yang, XiaoDong Han, Yan Gao, Yao Hu, Shilin Zhang, Hai Zhao

arXiv:2405.12532v2cs.CL

TL;DR

LLM inference is limited by GPU memory demands, while existing KV-cache methods overlook layer dependencies and the memory required to precompute the cache. PyramidInfer uses attention consistency to retain layer-wise crucial context during prefill and generation, achieving higher throughput and over 54% KV-cache memory reduction. Its speedup is limited at small batch sizes by additional computation, and its prefill compression is not lossless.

  • Problem

    Existing KV-cache compression methods do not account for inter-layer dependencies or reduce the substantial memory consumed while precomputing the initial cache.

  • Method

    PyramidInfer uses Recent Attention Consistency to select layer-wise Pivotal Context keys and values during both prefill and generation.

  • Results

    PyramidInfer achieves 2.2x higher throughput than Accelerate, 1.4x than Deepspeed, and 2.4x than H2O, with over 54% less KV-cache GPU memory.

  • Takeaways & Limitations

    PyramidInfer reduces KV-cache memory while maintaining model performance and supports layer-wise control of compression, including fixed PvC counts for unlimited input.

  • Takeaways & Limitations

    Additional computation limits speedup at small batch sizes, and PyramidInfer is not lossless during prefill, leaving more effective methods for future work.

Abstract

from arXiv · show

Large Language Models (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.

1 Introduction

LLM inference is constrained by substantial GPU memory use, especially from the KV cache and its prefill computation. PyramidInfer addresses this by selecting layer-dependent crucial context during prefill and generation, achieving higher throughput with substantially less KV-cache memory.

  • LLM inference requires substantial GPU memory, hindering deployment at scale for thousands of chatbot demands.
  • The KV cache stores previously computed attention keys and values for reuse, but can consume around 72 GB and constrain inference throughput.
  • Existing KV-cache compression methods prune already-computed caches, leaving the large initial prefill computation unaddressed.
  • Inference Context Redundancy motivates computing only keys and values that record context information rather than retaining all training-related information.
  • Recent Attention Consistency identifies crucial context by using consistent attention from recent tokens as an oracle for future generations.
  • PyramidInfer selects layer-wise PvCs during prefill and generation, reaching 2.2x and 1.4x higher throughput than Accelerate and Deepspeed, and 2.4x than H2O, with over 54% less KV-cache GPU memory.

2 Related Work

Prior approaches enlarge available memory, optimize GPU memory access, or prune KV caches after computation. PyramidInfer instead accounts for layer differences and compresses during both prefill and generation.

  • Inference parallelism increases available space through larger memory or multiple GPUs, including pipeline parallelism and KV-cache offload.
  • CUDA-oriented methods reduce memory traffic or allocation waste through FlashAttention 2 and PagedAttention.
  • StreamingLLM supports unlimited input by retaining recent context, while sacrificing memorization of historical context.
  • H2O and Scissorhands use attention-based KV-cache compression but treat layers uniformly and cannot compress during prefill.
  • PyramidInfer accounts for layer differences and compresses KV caches during both prefill and generation while maintaining generation quality.

3 Observation and Insight

The paper identifies redundant KV-cache information during inference and finds that crucial context can be selected from consistent attention patterns. These observations show that redundancy increases with layer depth, enabling layer-aware compression.

  • Inference Context Redundancy: Inference Context Redundancy arises because only the last token predicts the next token during inference, while earlier tokens retain training-related information in the KV cache.
  • Recent Attention Consistency: Recent tokens share attention to common context keys and values, motivating Recent Attention Consistency as an oracle for selecting future-useful KV entries.
  • PvC Consistency: 86% average PvC overlap with the last token shows shared context, but a single token at d = 25% covers only about 83% of the next-token PvC.
  • PvC Consistency: 93% average overlap after ensemble averaging indicates that recent-token groups provide a more reliable PvC oracle for predicting future tokens.At d = 20%, overlap ratios show hardly any drop, supporting prediction 20% ahead.
  • Layer-wise Redundancy: Deeper layers have lower PvC overlap because only a few keys and values receive consistently high weights while many others have similarly low weights.The resulting depth pattern is consistent with the power-law redundancy observed for ICR.

4 Layer-wise PvC Selection

PyramidInfer selects pivotal context layer by layer during both prefill and generation. It retains more context in shallow layers, progressively shortens deeper-layer caches, and updates selected entries using recent tokens.

  • Method Overview: PyramidInfer compresses KV cache in both prefill and generation by layer-wise selecting pivotal context.
  • Prefill Phase: During prefill, the method divides inputs into recent and context sequences, averages recent-token attention weights, and selects top-p context keys and values per layer.
  • Layer-wise Selection: PyramidInfer retains larger PvC ratios in shallow layers and decreases p with depth, producing shorter deeper-layer caches shaped like a pyramid.
  • Efficiency: Layer-wise PvC selection saves more GPU memory than methods that compute the whole prompt and continues improving generation efficiency through a smaller initial cache.
  • Layer-wise Selection: The forward pass computes attention weights, averages recent-context attention, gathers top-p KV entries, and reduces p by a decay ratio.
  • Generation Phase: During generation, a sliding recent window updates PvCs, while layer-specific lengths control compression and can maintain a fixed number of entries for unlimited input.

5 Evaluation

PyramidInfer is evaluated across language modeling, benchmarks, conversation, and long-context tasks, showing quality retention alongside lower KV-cache memory and higher throughput. Ablations support power-law PvC decay and a 40% recent-sequence ratio as effective settings.

  • Basic Evaluation: PyramidInfer maintains generation quality with substantially less GPU memory than full-cache and outperforms the local strategy across models and tasks.
  • Efficiency Result: 2.24x throughput and 54.6% less KV-cache GPU memory are achieved than full-cache Accelerate on LLaMA 2-13B.
  • Efficiency Result: PyramidInfer supports generation for LLaMA 2-70B, whereas H2O cannot process the prompt before running out of memory.
  • Efficiency Result: At maximum batch size, PyramidInfer enables more than 2x the batch size of other methods and achieves 2.8x and 1.7x throughput over Accelerate and Deepspeed.
  • PvC Length Decay: Under 60% compression, the power-law pattern reduces PvC length less in shallow layers and slightly improves downstream performance.
  • Recent Sequence Ratio: Increasing the recent-sequence ratio lowers KV-cache memory, while perplexity reaches a trough at 40–60%, motivating 40% as a performance-memory trade-off.

6 Conclusion

The conclusion presents PyramidInfer as a method for compressing KV caches during both prefill and generation. It reports reduced GPU memory usage while maintaining model performance for resource-constrained LLM deployment.

  • PyramidInfer compresses the KV cache during both prefill and generation phases to reduce the difficulty of deploying LLMs at scale.
  • The method is reported to reduce GPU memory usage without compromising model performance in resource-constrained environments.

Limitations

The paper identifies two limitations: selection overhead reduces speedups at small batch sizes, and prefill compression is not lossless. More effective prefill methods remain open for future work.

  • Additional computation for selecting PvCs limits PyramidInfer’s speedup when the batch size is small.
  • PyramidInfer does not compress the KV cache losslessly during prefill, leaving more effective prefill methods for future work.

A.1 Additional Computational Cost in PyramidInfer

PyramidInfer’s PvC-selection overhead can offset its KV-cache savings at small batch sizes, but becomes relatively insignificant as batch size increases. Figure 9 compares throughput against the full-cache baseline across batch sizes.

  • Figure 9 compares PyramidInfer and the full-cache baseline across batch sizes on LLaMA 2-7B with input length 512+256.
  • PvC selection adds mainly top-p sorting cost in each layer, while other additional costs are negligible.
  • PyramidInfer has limited acceleration at small batch sizes because selection overhead offsets KV-cache savings, but the overhead becomes trivial as batch size increases.

A.2 Position Encoding

PyramidInfer compares two ways to assign position encodings after layer-wise KV reduction. Gathering the scattered original encodings performs slightly better on LLaMA 2-13B downstream tasks.

  • A.2 Position Encoding: The position-encoding comparison evaluates re-encoding positions from position 0 against gathering scattered original encodings.These are the two alternatives tested after reducing keys and values.
  • A.2 Position Encoding: Gathering scattered original position encodings achieves slightly better downstream-task performance than re-encoding positions from position 0.The comparison is conducted on LLaMA 2-13B.

B Extended Discussions

The extended analysis links layer depth and attention consistency to KV-cache redundancy. It distinguishes shared PvCs, which subsequent tokens collectively attend to, from non-shared PvCs that mainly support teacher-forced next-token prediction.

  • Layer-wise Redundancy: Deeper layers exhibit lower PvC overlap ratios, consistent with a power-law distribution.The analysis attributes this pattern to rapidly deepening heatmap colors as layer depth changes.
  • Layer-wise Redundancy: High redundancy in deeper layers leaves most keys and values with low attention weights, limiting their influence on perplexity and selection as PvCs.These entries are identified as non-PvCs.
  • KV-Cache Composition: The KV cache is divided into PvCs and non-PvCs, with PvCs further divided into shared and non-shared PvCs.This decomposition supports the analysis of which entries carry context information.
  • KV-Cache Composition: 87% overlap between tokens and the last token identifies shared PvCs, while the remaining 13% are non-shared PvCs that PyramidInfer does not use.Non-shared PvCs receive high attention from the current token and are useful for predicting its successor.
  • Attention Roles: Shared PvCs are keys and values that subsequent tokens collectively attend to.The analysis selects shared PvCs using a recent sequence ratio of 20%.
  • Attention Roles: Non-shared PvCs seldom overlap with other tokens’ non-shared PvCs, indicating they mainly support current-token teacher-forced prediction rather than subsequent-token context.This distinguishes their training-time role from the context retained for future generation.
Loading 2405.12532v2…