PaperlayerView exact revision on arXiv

Source-linked AI summary

Keep the Cost Down: A Review on Methods to Optimize LLM' s KV-Cache Consumption

Luohe Shi, Hongyi Zhang, Yao Yao, Zuchao Li, Hai Zhao

arXiv:2407.18003v4cs.CL

TL;DR

KV Cache makes autoregressive generation more efficient, but its memory use grows with context length and can constrain GPU inference. This review organizes compression and system-optimization methods across training, deployment, and post-training, and surveys metrics for long-text efficiency and capability. It concludes that these approaches offer distinct ways to address KV Cache’s memory demands, while evaluation should consider both efficiency and model capabilities.

  • Problem

    KV Cache reduces per-token generation complexity to linear, but its memory grows with sequence length and can bottleneck GPU inference.

    Source: §1 · pp. 1–2

  • Method

    The review compares KV Cache optimization methods across training, deployment, and post-training, and presents long-text efficiency and capability metrics.

    Source: §1 · p. 2

  • Results

    The review synthesizes methods spanning architectural changes, deployment frameworks, and post-training strategies such as eviction and quantization, alongside metrics for evaluating efficiency and model capabilities.

    Source: §8 · p. 10

  • Takeaways & Limitations

    The review concludes that evaluating long-text optimization requires balancing efficiency with model capabilities.

    Source: §8 · p. 10

    Few-shot extensions of short tests have limited reference value because distant text serves only as reference and contains neither needed information nor inference conditions.

    Source: §6.1 · p. 9

Abstract

from arXiv · show

Large Language Models (LLMs), epitomized by ChatGPT's release in late 2022, have revolutionized various industries with their advanced language comprehension. However, their efficiency is challenged by the Transformer architecture's struggle with handling long texts. KV Cache has emerged as a pivotal solution to this issue, converting the time complexity of token generation from quadratic to linear, albeit with increased GPU memory overhead proportional to conversation length. With the development of the LLM community and academia, various KV Cache compression methods have been proposed. In this review, we dissect the various properties of KV Cache and elaborate on various methods currently used to optimize the KV Cache space usage of LLMs. These methods span the pre-training phase, deployment phase, and inference phase, and we summarize the commonalities and differences among these methods. Additionally, we list some metrics for evaluating the long-text capabilities of large language models, from both efficiency and capability perspectives. Our review thus sheds light on the evolving landscape of LLM optimization, offering insights into future advancements in this dynamic field. Links to the papers mentioned in this review can be found in our Github Repo https://github.com/zcli-charlie/Awesome-KV-Cache.

1 Introduction

KV Cache makes token generation linear in sequence length, but its memory use grows with the sequence and can bottleneck GPU inference. This review organizes optimization methods across training, deployment, and post-training, and introduces metrics for assessing efficiency and capability.

1 Introduction

  • KV Cache reduces per-token generation time complexity from quadratic to linear, while its memory use grows with sequence length and can bottleneck GPU inference.Each dialogue must retain its own cache, and GPU memory bandwidth can constrain generation speed.Source: §1 · pp. 1–2
  • The review surveys KV Cache optimization across pre-training, deployment, and post-training, focusing primarily on methods that edit or optimize the cache itself.Other approaches are briefly mentioned in Appendix C.Source: §1 · p. 2
  • Pre-training methods can be effective but are unsuitable for modifying existing models or for settings with low computational budgets.Source: §1 · p. 2
  • The review assesses long-text efficiency through generation speed and space occupancy, and capability through the impact on model performance.The stated aim is to ensure capability loss remains acceptable.Source: §1 · p. 2
  • The review aims to help readers understand KV Cache optimization and build more effective, efficient, and sustainable LLMs.Source: §1 · p. 2

2 Preliminary and Notations

The section establishes notation for token sequences, model layers, and hidden states, then introduces the decoder’s attention components and the retained key and value tensors.

2 Preliminary and Notations

  • The input X is a length-n sequence of tokens xi from vocabulary V, and the model maps it to probability distributions pi for next-token prediction.During autoregressive decoding, pn is usually used to predict the next token.Source: §2 · p. 2
  • The decoder places L Transformer blocks between an embedding layer and a linear layer with softmax, mapping token embeddings toward probability distributions.The embedding layer maps each token to a d-dimensional vector.Source: §2 · pp. 2–3
  • Each decoder block contains self-attention and a feed-forward network, with residual connections and normalization; this paper focuses on self-attention.Source: §2 · p. 3
  • In multi-head attention, queries, keys, and values are split into nh heads, while a causal mask prevents earlier tokens from attending to later tokens.The mask adds negative infinity to the upper-right triangle of the pre-softmax attention matrix.Source: §2 · p. 3
  • During autoregressive decoding, the model retains previously computed keys and values and updates them with each new token; these retained tensors are the KV Cache.Source: §2 · p. 3

3 Training Stage Optimization

Training-stage approaches reduce KV Cache use by changing attention architecture or adding context compression. Their trade-offs include model modification, added training, and unresolved memory-bandwidth constraints.

3 Training Stage Optimization

  • MQA reduces KV Cache use to 1/nh of the original by having all query heads share a single key and value head.Shazeer (2019) proposed the method, claiming satisfactory performance with one key-value head.Source: §3 · p. 4
  • GQA shares one key-value head within each of ng query-head groups, reducing stored KV Cache to ng/nh and allowing a finer efficiency-performance balance than MQA or MHA.When ng = 1 it becomes MQA; when ng = nh it becomes MHA. Ainslie et al. (2023) also report parameter savings ratio η = 0.5 + 0.5 · ng/nh.Source: §3 · p. 4
  • Cross-layer KV reuse reduces cache size by sharing KV across layers, but scattered access can leave the memory-bandwidth bottleneck unresolved.Brandon et al. (2024) reuse contents directly, while YOCO and GoldFinch use linear models to generate layers’ KV Cache.Source: §3 · p. 4
  • MLA uses an additional dimension-expanding matrix to decompress KV Cache rather than simply reusing existing content, but does not address the memory-bandwidth bottleneck.The cache still needs to be expanded into large KV vectors.Source: §3 · pp. 4–5
  • CEPE adds an encoder that compresses RAG reference text before passing it to the decoder through cross-attention, shortening the sequence whose KV Cache is stored.It can use a pre-trained LLM, but the added decoder and cross-attention layer require substantial extra training.Source: §3 · p. 5

4 Deploy-Stage Optimization

Deployment-stage methods address KV Cache allocation, bandwidth, and repeated computation through memory management, distributed serving, reuse, and CPU offloading. The approaches target different system bottlenecks and configurations.

4 Deploy-Stage Optimization

  • Repeated KV Cache allocation and release cause memory fragmentation, while limited batching worsens bandwidth bottlenecks and computational inefficiency.Source: §4 · p. 5
  • Paged Attention maps a logically continuous KV Cache onto discontinuous GPU memory and uses custom CUDA kernels, yielding almost no unused fragments and efficient inference.Kwon et al. (2023) introduced the mechanism with the vLLM framework.Source: §4 · p. 5
  • DistAttention and DistKV-LLM distribute KV Cache deployment across multiple servers to improve LLM service efficiency on large-scale cloud servers.Lin et al. (2024) developed these methods from the Paged Attention idea.Source: §4 · p. 5
  • ChunkAttention reuses KV Cache for the longest common prefix shared with a new dialogue, avoiding repeated pre-fill computation; a related extension uses multiple storage devices.Ye et al. (2024) introduce the dictionary-tree approach, and Gao et al. (2024) extend it hierarchically.Source: §4 · p. 5
  • Speculative CPU offloading keeps only critical KV Cache components on the GPU and reloads a small amount during inference, saving high-speed device memory without compromising performance.Jin et al. (2024) and Lee et al. (2024) propose this approach.Source: §4 · p. 5
  • Other system approaches include using CPU attention computations to improve GPU utilization and optimizing the inference system as a whole.He & Zhai (2024) describe CPU-assisted attention, while Qin et al. (2024) report an enterprise system-wide optimization.Source: §4 · p. 5

5 Post-Training Optimizations

Post-training methods reduce cache use by evicting or merging tokens and by quantizing KV activations or model weights. Eviction policies may be static or dynamic, while quantization methods vary in which components they compress and how they preserve information.

5.1 Eviction and Merging

  • Keeping initial and recent tokens can maintain long-context performance, whereas retaining only recent tokens can cause model collapse beyond the cache size.The review describes this as a static policy: retain initial tokens that receive consistently high attention alongside recent tokens.Source: §5.1 · p. 6
  • Dynamic eviction methods use attention history or other signals to select tokens, including TOVA’s layer-wise removal of tokens with minimal attention weights at each decoding step.Other approaches accumulate attention scores, favor recent tokens, shorten caches more in deeper layers, smooth remaining attention distributions, or select per-head policies during prompt encoding.Source: §5.1 · p. 6
  • Because attention scores are incompatible with methods such as Flash Attention, alternative eviction signals include head properties, token entropy, and key L2 norms.The review notes that these approaches can be combined with Flash Attention; Devoto et al. (2024)’s approach is described as granular and controllable.Source: §5.1 · p. 7
  • Merging extends eviction by combining KV entries rather than simply discarding them, with methods differing in how they choose regions and combine keys and values.Examples include attention-guided merging, Gaussian weighting, cross-head merging that turns MHA into GQA, and Anchor-LLM’s anchor-based self-attention.Source: §5.1 · p. 7

5.2 Quantization

  • KV Cache-only quantization targets activations while leaving model weights unchanged; KVQuant uses distinct quantization strategies for keys and values and isolates outliers.Its features include per-channel key and per-token value quantization, sensitivity-weighted non-uniform datatypes, and normalization of quantization centroids.Source: §5.2 · pp. 7–8
  • Other cache-quantization methods combine compression with selective retention or error control, including LESS, MiKV, QAQ, and GEAR.These methods use, respectively, a constant-sized low-rank cache, mixed precision for evicted and important entries, separate key/value strategies with future-attention estimates, or quantization with low-rank and sparse error correction.Source: §5.2 · p. 8
  • Full quantization compresses both model weights and attention cache; FlexGen uses 4-bit compression, and WKVQuant combines weight and KV Cache quantization strategies.WKVQuant uses Past Only Quantization, Two-dimensional Quantization, and Cross-block Reconstruction Regularization for the cache.Source: §5.2 · p. 8

6 Evaluation

The review describes datasets and metrics for evaluating long-text capability alongside efficiency. It covers long-context benchmarks, key retrieval, and extended few-shot tests, then outlines memory, speed, and perplexity measures.

6.1 Datasets

  • Long-context benchmarks test information extraction from documents typically spanning 4,000 to 20,000 tokens, with some reaching nearly 200,000 tokens.Examples include LongBench, ZeroSCROLLS, L-Eval, BAMBOO, XL2bench, InfiniteBench, and LooGLE; concerns include model-specific knowledge demands, fixed text lengths, and subjective generated content.Source: §6.1 · p. 9
  • Key-retrieval tests extract a specified item from unrelated context, offering controllability without relying on existing model knowledge, though critics consider them overly simplistic.Examples include Needle in a Haystack, Passkey Retrieval, and RULER.Source: §6.1 · p. 9
  • Few-shot testing can extend traditional test sets or simulate multi-turn dialogues, but distant text often serves only as reference, limiting the results’ value.Source: §6.1 · p. 9

6.2 Evaluation Metric

  • Per-token GPU-memory usage estimates cache space per token, but rigorous measurement should account for actual token-generated fragment space.For LLaMA2-7B, the theoretical per-entry KV Cache usage is 0.5MB.Source: §6.2 · p. 9
  • Throughput measures generated tokens per second, while decoding latency is the time per new token and is typically the reciprocal of throughput.Source: §6.2 · p. 9
  • Perplexity exponentiates the average natural-log likelihood across tokens and provides a rough reference for performance changes; a sharp rise usually signals a substantial capability decline.Source: §6.2 · pp. 9–10

7 Key Takeaways

The review frames KV Cache optimization around reducing memory consumption by compressing keys or values, while recognizing a trade-off between deletion and information retention. It also points to external storage and retrieval as a possible direction with added integration complexity.

7 Key Takeaways

  • KV Cache optimization centers on reducing memory use by compressing keys or values, with effects on memory usage and processing speed.Source: §7 · p. 10
  • Deleting less important KV pairs offers immediate memory relief but may compromise performance, whereas compression aims to retain information integrity while reducing footprint.The choice between deletion and compression remains open.Source: §7 · p. 10
  • External KV Cache storage could reduce primary-device memory use, but would turn cache management into a retrieval and reintegration challenge.The review points to storage and retrieval technologies as potentially important to future KV Cache management.Source: §7 · p. 10

8 Conclusion

The conclusion surveys KV Cache optimization across training, deployment, and post-training methods, alongside metrics that consider both efficiency and model capability. It presents the review as a roadmap for learners in this evolving field.

8 Conclusion

  • The review spans architectural changes during pre-training, framework optimizations during deployment, and post-training strategies such as eviction and quantization.It emphasizes evaluating long-text performance through both efficiency and model-capability metrics.Source: §8 · p. 10

A Popular Models With GQA

The section defines the metrics used to describe popular models and their KV-cache characteristics. It highlights lower cache-to-embedding ratios with GQA and differing FFN-to-attention ratios by architecture.

A Popular Models With GQA

  • GQA and MoE indicate whether a model uses those methods, with MoE notation giving total and activated experts.For example, Mixtral is represented as 8(2), meaning 2 of 8 experts are used per token.Source: A Popular Models With GQA · p. 17
  • The table’s metrics include total parameter count, layer count, hidden dimension, and attention-head and average KV-head counts.PC includes parameters from attention, FFN, normalization, embedding, and LM-head components; HD need not equal the qkv-vector dimension.Source: A Popular Models With GQA · p. 17
  • The metrics also report embedding size and KV-cache increment per token in bytes, using bfloat16.SHK and SKV denote these respective quantities.Source: A Popular Models With GQA · p. 17
  • F/A is typically around 2 without GQA and 3.5–5 with GQA; for MoE models, it multiplies the single-expert ratio by activated experts.F/A compares FFN-layer parameters with attention-layer parameters.Source: A Popular Models With GQA · p. 17
  • With GQA, the KV-cache-to-embedding ratio R is significantly lower than with non-GQA methods, indicating relatively compressed token information.R is the KV Cache size generated per token divided by the embedding-vector size.Source: A Popular Models With GQA · p. 17

B Dataset Examples

The section provides examples for Passkey Retrieval and Needle in a Haystack, alongside a table describing popular open-source GQA models and approximate parameter counts.

B Dataset Examples

C Other Methods

Other approaches reduce KV-cache needs by eliminating the cache or shortening the prompt or embedding. The section reviews linear-transformer and prompt- or embedding-engineering approaches.

C.1 Linear-Transformers

C.2 Prompt and Embedding Engineering

  • Prompt and embedding methods compress history with summary vectors [Chevalier et al. (2023)] or use a gist token to transfer instruction information to inputs [Mu et al. (2023)].LLMLingua uses a budget controller, iterative token-level compression, and a small model to estimate token importance; LongLLMLingua further improves it for long-context scenarios [Jiang et al., 2023a; Jiang et al., 2023b].Source: §C.2 · p. 19

Related papers

Loading 2407.18003v4…