PaperlayerView exact revision on arXiv

Source-linked AI summary

No Token Left Behind: Reliable KV Cache Compression via Importance-Aware Mixed Precision Quantization

June Yong Yang, Byeongwook Kim, Jeongin Bae, Beomseok Kwon, Gunho Park, Eunho Yang, Se Jung Kwon, Dongsoo Lee

arXiv:2402.18096v1cs.LGcs.AI

TL;DR

KV-cache memory growth makes LLM deployment difficult, while the risks of eviction-based compression remain insufficiently examined. MiKV retains evicted KVs in low precision and important KVs in high precision; across diverse benchmarks it supports compression up to 80% with minimal performance degradation.

  • Problem

    KV-cache memory grows with batch size and sequence length, while the risks of eviction-based compression and the context it discards remain insufficiently analyzed.

    Source: §1 · p. 2

  • Method

    MiKV retains eviction-bound KVs in low precision and important KVs in high precision to preserve context while maintaining generation quality.

    Source: §1 · p. 3

  • Results

    Across diverse benchmarks, MiKV achieved up to 80% cache compression with minimal performance degradation; at 20% cache it preserved complex-generation quality where uniform quantization struggled.

    Source: §1 · p. 3, §4.2 · p. 8

  • Takeaways & Limitations

    MiKV's mixed-precision strategy recovers context damage associated with eviction while maintaining generation quality, including for models equipped with GQA.

    Source: §6 · p. 8

Abstract

from arXiv · show

Key-Value (KV) Caching has become an essential technique for accelerating the inference speed and throughput of generative Large Language Models~(LLMs). However, the memory footprint of the KV cache poses a critical bottleneck in LLM deployment as the cache size grows with batch size and sequence length, often surpassing even the size of the model itself. Although recent methods were proposed to select and evict unimportant KV pairs from the cache to reduce memory consumption, the potential ramifications of eviction on the generative process are yet to be thoroughly examined. In this paper, we examine the detrimental impact of cache eviction and observe that unforeseen risks arise as the information contained in the KV pairs is exhaustively discarded, resulting in safety breaches, hallucinations, and context loss. Surprisingly, we find that preserving even a small amount of information contained in the evicted KV pairs via reduced precision quantization substantially recovers the incurred degradation. On the other hand, we observe that the important KV pairs must be kept at a relatively higher precision to safeguard the generation quality. Motivated by these observations, we propose \textit{Mixed-precision KV cache}~(MiKV), a reliable cache compression method that simultaneously preserves the context details by retaining the evicted KV pairs in low-precision and ensure generation quality by keeping the important KV pairs in high-precision. Experiments on diverse benchmarks and LLM backbones show that our proposed method offers a state-of-the-art trade-off between compression ratio and performance, compared to other baselines.

1. Introduction

KV caching accelerates autoregressive generation, but its memory footprint grows with batch size and sequence length. The introduction documents risks from eviction and motivates MiKV, which retains evicted KVs at low precision and important KVs at high precision.

1. Introduction

  • KV-cache memory grows with batch size and sequence length, creating a deployment bottleneck alongside model weights, especially as supported contexts lengthen.Autoregressive KV caching stores intermediate states from prior context to accelerate generation.Source: §1 · pp. 1–2
  • Eviction methods assume important KVs suffice, but the risks of discarded context remain insufficiently analyzed and future KV needs are hard to predict, especially in multi-turn conversations.The lost information can include system prompts or core context details.Source: §1 · p. 2
  • Empirical observations associate cache eviction with contextual incoherence, hallucinations, detail loss, and loss of safety prompts that can trigger responses bypassing safety measures.Source: §1 · p. 2
  • Retaining evicted KVs in low precision substantially recovers lost details, especially when systematic outliers in keys and queries are handled.The paper frames this as preserving information rather than exhaustively discarding evicted states.Source: §1 · pp. 2–3
  • MiKV keeps eviction-bound KVs in low precision and important KVs in high precision to preserve context details while maintaining generation quality.The authors also investigate an effective condition for quantizing the retained KVs.Source: §1 · p. 3
  • Up to 80% KV-cache compression incurred minimal performance degradation across diverse benchmarks; MiKV also preserved complex-generation quality at 20% cache where uniform quantization struggled.The introduction evaluates natural-language understanding, math, code, detail retrieval, and chatting.Source: §1 · p. 3

2. Context Damage from KV Cache Eviction

The section examines how eviction-based KV-cache compression can damage context, and whether retaining evicted information can make compression more reliable. It considers qualitative failures and line-retrieval performance under eviction.

2.1. Background

  • KV-cache memory grows with batch size and sequence length, while importance-based compression removes intermediate states and risks losing input context.Importance criteria include locality, frequency, and attention structure; predicting the importance of future information is described as inherently unfeasible.Source: §2.1 · p. 3

2.2. Qualitative Analysis

  • Eviction can compromise safety prompts, disrupt contextual coherence, and prompt hallucination of missing details or topics.The examples include lost safety instructions, partial loss of information about “Two New Sciences,” and invented topics after dialogue history is evicted.Source: §2.2 · p. 4

2.3. Quantitative Analysis

  • The section uses line retrieval, where a model retrieves a value associated with a key in randomly generated strings, to assess context damage as eviction increases.The reported failure likelihood rises as the eviction ratio increases.Source: §2.3 · p. 4
  • Eviction and even oracle eviction rapidly degrade line-retrieval performance, indicating that foreknowledge of future KV importance does not prevent performance loss.Figure 3 compares H2O, oracle eviction, and MiKV; the text reports degradation for eviction and oracle eviction.Source: Figure 3 · p. 3, §2.3 · pp. 4–5

3. Mixed-Precision KV Cache Compression

MiKV retains evicted KV pairs at low precision and important pairs at higher precision. The section examines low-bit recovery, outlier-aware quantization, precision choices for the importance cache, and an acceleration approach.

3.1. Retaining Evicted KVs with Quantization

  • Low-precision retention of evicted KVs substantially restores performance lost to eviction, but consumes memory and very low precision such as INT2 can undermine recovery.The section frames low-bit retention as a trade-off between performance recovery and compression rate.Source: §3.1 · p. 5, Table 2 · p. 5

3.2. Low-bit KV Quantization with Dynamic Outlier Awareness

  • Systematic query and key outliers occur across the model and increase quantization error; channel balancing reduces key outliers and promotes query outlier awareness.The prefill-computed balancer is applied to keys and queries; with groups of half the attention-head dimension, the measures restore line-retrieval performance at INT2 and mitigate RoPE artifacts.Source: §3.2 · pp. 5–6

3.3. Reducing the Precision of the Importance Cache

  • Reducing importance-cache precision can increase compression while keeping performance degradation minimal, but excessive reduction degrades performance.The described experiment uses a 20% importance ratio and an outlier-aware retention cache at 2-bit precision.Source: §3.3 · p. 6, Table 3 · p. 6

3.4. Accelerating the Mixed-Precision KV Cache

  • MiKV groups KV pairs by precision and lowers K and V precision while keeping Q and the attention map in floating point, enabling weight-only quantization kernels instead of batch-GEMV operations.The paper attributes the resulting speedup to addressing memory-bound generation on devices with high computational power relative to memory bandwidth.Source: §3.4 · p. 6

4. Experiments

The experiments compare MiKV with H2O and uniform-precision RTN across four benchmarks and four open-source LLMs, then assess chatbot quality and cache memory.

4.1. Experimental Setup

  • The setup evaluates MMLU, GSM8k, HumanEval, and Line Retrieval on four LLMs, comparing MiKV with H2O and uniform-precision RTN.MMLU and GSM8k use a 1-shot setting; the models are Llama-2 7b, 13b, and 70b, and Mistral-7b.Source: §4.1 · p. 7

4.2. Main Results

  • Across benchmarks and backbones, MiKV matches full-cache generation at better compression; it preserves Line Retrieval and GSM8K/HumanEval quality at 20% cache where eviction or RTN degrade.The evaluation also covers chatbot generation and memory use.Source: §4.2 · p. 8

4.3. AlpacaEval Results

  • For Llama-2-70b-chat, MiKV shows no AlpacaEval win-rate drop against full cache at cache sizes as small as 25%.Table 4 reports the AlpacaEval win rate of MiKV over full cache.Source: §4.3 · p. 8, Table 4 · p. 8

4.4. Memory Footprint Analysis

  • With batch size 8 and sequence length 4096, MiKV significantly reduces KV-cache memory across models of varying sizes and GQA availability.Table 5 compares memory footprint for full cache and MiKV under this setting.Source: §4.4 · p. 8

5. Related Work

Related work reduces KV-cache memory through attention architectures, eviction policies, and quantization, while each approach has stated trade-offs or gaps.

KV cache sharing

  • MQA and GQA reduce cache size by sharing KVs across query heads, but sacrifice performance and require substantial training costs for GQA.These approaches are associated with Shazeer (2019) and Ainslie et al. (2023).Source: §5 · p. 8

KV cache eviction

  • Eviction policies preserve local tokens, heavy-hitters, or adaptively selected tokens, but exhaustively discard context in evicted KVs.The cited approaches include Jiang et al. (2023), Xiao et al. (2023), Zhang et al. (2023), Liu et al. (2023a), and Ge et al. (2024).Source: §5 · p. 8

KV cache quantization

6. Conclusion

The conclusion presents MiKV as importance-based mixed-precision KV-cache compression that recovers context damage while maintaining generation quality, including for GQA models.

6. Conclusion

  • MiKV retains unimportant KVs in low precision and important KVs in high precision, recovering eviction-related context damage while maintaining generation quality, including for GQA models.The authors report validation through experiments.Source: §6 · p. 8

Broader Impact

The paper frames its KV-cache compression work around LLM inference memory and safety, while noting potential societal consequences.

Broader Impact

  • The work targets KV-cache memory footprint and potential safety issues, and the authors note that it may have societal consequences they do not specifically highlight.The stated context is LLM inference and preserving model performance.Source: Broader Impact · p. 9

A. Full Context for Qualitative Examinations

The appendix provides full prompts and outputs for the safety-breach and contextual-incoherency examinations, while noting that the topic-retrieval prompt is omitted because it exceeds 5,000 tokens.

A. Full Context for Qualitative Examinations

B. Additional Query, Key, Value Plots

Additional QKV plots show outliers across layer depths and backbone models, with plots identified for four LLMs.

B. Additional Query, Key, Value Plots

K Q4 V

C. Experiments on Per-Channel Quantization of Keys

A simulated per-channel key-quantization experiment reports preserved line-retrieval performance, while noting that actual deployment would require changes to eviction and kernel support.

C. Experiments on Per-Channel Quantization of Keys

  • Per-token quantization was adopted for compatibility with existing eviction strategies and kernel support despite systematic key and query outliers.Per-channel quantization can reduce error when quantization direction aligns with outlier direction, but requires implementation changes to buffering and eviction.Source: C. Experiments on Per-Channel… · p. 18
  • The per-channel result is hypothetical: actual performance may differ unless eviction is modified and suitable kernel support is provided.The simulated scheme uses group size 64 and quantizes keys without reordering or buffering KV pairs.Source: C. Experiments on Per-Channel… · p. 18
  • Per-channel key quantization is effective at preserving line-retrieval performance by isolating outliers in a simulated experiment.The setup keeps 20% of KV pairs in FP16 in the importance cache and 80% in INTx in the retained cache; Table 6 reports line-retrieval accuracy for this 20% importance ratio.Source: Table 6 · p. 18, C. Experiments on Per-Channel… · p. 18

D. Detailed Experimental Settings

The appendix specifies shared inference settings and evaluation procedures for GSM8K, HumanEval, line retrieval, MMLU, and AlpacaEval.

D. Detailed Experimental Settings

D.1. GSM8K

D.2. HumanEval

  • HumanEval uses 164 evaluation samples, each generated once with greedy decoding and scored for functional correctness.Source: §D.2 · p. 19

D.3. Line Retrieval

  • The line-retrieval evaluation set contains 500 samples with 20 key-value context lines, and Figure 15 shows an example prompt.Source: §D.3 · p. 19, Figure 15 · p. 21

D.4. MMLU

D.5. AlpacaEval

  • AlpacaEval uses official code and standard settings, comparing compressed-cache outputs against full-cache outputs with GPT-4 as judge.Source: §D.5 · p. 19
Loading 2402.18096v1…