PaperlayerView exact revision on arXiv

Source-linked AI summary

H$_2$O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models

Zhenyu Zhang, Ying Sheng, Tianyi Zhou, Tianlong Chen, Lianmin Zheng, Ruisi Cai, Zhao Song, Yuandong Tian, Christopher Ré, Clark Barrett, Zhangyang Wang, Beidi Chen

arXiv:2306.14048v3cs.LG

TL;DR

LLM deployment is expensive because KV caches grow with generation workloads, motivating compact cache policies that preserve useful attention information. H2O retains heavy-hitter and recent tokens through a dynamically updated eviction policy formulated as dynamic submodular maximization. Across evaluated models and tasks, it reduces cache requirements while maintaining generation quality and improving inference throughput and latency.

  • Problem

    KV caches create a prominent memory bottleneck for long-content and large-batch LLM generation, while effective eviction must balance cache size, miss rate, and computational cost.

  • Method

    H2O dynamically retains heavy-hitter and recent tokens, using local accumulated attention statistics and a dynamic submodular formulation with theoretical guarantees.

  • Results

    Across OPT, LLaMA, and GPT-NeoX tasks, H2O achieves comparable performance to full KV with less than 20% cache budget and consistently improves over the Local strategy.

  • Takeaways & Limitations

    H2O significantly improves end-to-end throughput and reduces wall-clock latency without compromising generation quality across evaluated tasks.

  • Takeaways & Limitations

    The formal eviction setting keeps cache size fixed and allows eviction of at most one KV entry between consecutive steps.

Abstract

from arXiv · show

Large Language Models (LLMs), despite their recent impressive accomplishments, are notably cost-prohibitive to deploy, particularly for applications involving long-content generation, such as dialogue systems and story writing. Often, a large amount of transient state information, referred to as the KV cache, is stored in GPU memory in addition to model parameters, scaling linearly with the sequence length and batch size. In this paper, we introduce a novel approach for implementing the KV cache which significantly reduces its memory footprint. Our approach is based on the noteworthy observation that a small portion of tokens contributes most of the value when computing attention scores. We call these tokens Heavy Hitters (H$_2$). Through a comprehensive investigation, we find that (i) the emergence of H$_2$ is natural and strongly correlates with the frequent co-occurrence of tokens in the text, and (ii) removing them results in significant performance degradation. Based on these insights, we propose Heavy Hitter Oracle (H$_2$O), a KV cache eviction policy that dynamically retains a balance of recent and H$_2$ tokens. We formulate the KV cache eviction as a dynamic submodular problem and prove (under mild assumptions) a theoretical guarantee for our novel eviction algorithm which could help guide future work. We validate the accuracy of our algorithm with OPT, LLaMA, and GPT-NeoX across a wide range of tasks. Our implementation of H$_2$O with 20% heavy hitters improves the throughput over three leading inference systems DeepSpeed Zero-Inference, Hugging Face Accelerate, and FlexGen by up to 29$\times$, 29$\times$, and 3$\times$ on OPT-6.7B and OPT-30B. With the same batch size, H2O can reduce the latency by up to 1.9$\times$. The code is available at https://github.com/FMInference/H2O.

1 Introduction

LLM deployment is constrained by the growing KV-cache memory footprint and the difficulty of choosing effective, inexpensive eviction policies. H2O addresses this by retaining influential heavy-hitter tokens alongside recent tokens, with reported accuracy and throughput benefits.

  • A 30-billion-parameter model with batch size 128 and sequence length 1024 can require 180GB of KV cache.
  • Sparsity for small cache size: Over 95% attention sparsity at inference suggests that retaining only 5% of the KV cache may enable up to a 20× reduction without an accuracy drop.
  • Heavy-Hitters for low miss rate: Heavy-hitter tokens emerge from a power-law distribution of accumulated attention scores and are critical for maintaining generation accuracy.
  • Greedy algorithm for low-cost policy: H2O retains heavy hitters and recent KV embeddings, using local accumulated attention statistics at each decoding step.

2 Related Work and Problem Setting

The paper positions KV-cache eviction as a distinct inference problem because standard sparse-attention and caching methods do not directly resolve generation-time memory, accuracy, and eviction-cost constraints. It formalizes generation with a fixed-size cache and sequential eviction updates.

  • Related work: Existing sparse-attention, low-rank, and multi-query methods reduce attention or cache costs but may not directly apply to generation with pretrained LLMs.
  • LLM inference breakdown: During generation, each step uses previously generated tokens while updating the KV cache; the paper focuses on improving this generation-phase cache.
  • Problem formulation: The formal setting fixes the cache size at k and permits eviction of at most one KV entry between consecutive decoding steps.
  • Problem formulation: The eviction policy updates the cached token set using normalized attention information, with the goal of producing outputs comparable to generation without cache-size limits.

3 Observations

The observations identify two properties that make compact KV caching plausible: inference attention is highly sparse, and a small set of heavy hitters concentrates attention value. These findings motivate retaining heavy hitters together with recent tokens.

  • 3.1 Sparsity for Small Cache Size: Attention matrices in pretrained LLMs show sparsity above 95% in almost all layers despite dense training.
  • 3.1 Sparsity for Small Cache Size: Attention sparsity suggests that accessing every previous key and value embedding is unnecessary for generating the next token.
  • 3.2 Heavy-Hitters for Low Miss Rate: The best eviction policy is difficult because evicting important keys and values can damage performance through sequential generation dependencies.
  • 3.2 Heavy-Hitters for Low Miss Rate: Accumulated attention scores follow a power-law distribution, indicating a small set of critical tokens called heavy hitters.
  • 3.2 Heavy-Hitters for Low Miss Rate: Under a submodular attention assumption, greedily constructing the retained set has a near-optimal property.

4 Heavy-Hitter Oracle

H2O formulates heavy-hitter-based KV eviction as dynamic submodular maximization and implements a greedy policy that updates the cache from accumulated attention scores. The paper provides an illustrative procedure and a theoretical approximation guarantee.

  • 4 Heavy-Hitter Oracle: H2O defines heavy-hitter cache eviction as a dynamic submodular problem and derives theoretical guarantees for the greedy policy.
  • 4.1 Greedy Algorithm for Low-Cost Policy: Local heavy hitters computed from preceding-token attention are empirically as effective as using attention to future tokens.
  • 4.1 Greedy Algorithm for Low-Cost Policy: The dynamic submodular framework treats cached words as existing elements and attention score as an example of the submodular objective.
  • 4.1 Greedy Algorithm for Low-Cost Policy: With cache budget 3, the illustrative algorithm evicts the third token after the fourth decoding step based on accumulated attention score.
  • 4.2 Theoretical Guarantee and System Implementation: Under mild assumptions, the greedy top-k construction satisfies f(eS_i) ≥ (1 − α)(1 − 1/e) max_|S|=k f(S) − β.

5 Empirical Evaluation

Across models, tasks, cache budgets, and hardware settings, H2O reduces KV-cache memory while preserving accuracy and improving inference efficiency. Its benefits extend to long-context generation, varying shot counts, quantization, and combinations with other sparsification methods.

  • End-to-End Accuracy and Memory: H2O reduces KV-cache memory footprint by 5–10× while achieving comparable accuracy on a majority of tasks.Experiments cover OPT, LLaMA, and GPT-NeoX across multiple benchmarks and model sizes.
  • End-to-End Accuracy and Memory: With less than 20% KV-cache budget, H2O achieves comparable performance to full KV embeddings across downstream tasks.At 20% budget, H2O also consistently improves over the Local and Sparse Transformer baselines, including long-sequence XSUM and CNN/Daily Mail tasks.
  • Throughput and Latency: 3×, 29×, and 29× higher throughput are achieved over FlexGen, DeepSpeed, and Accelerate, respectively, with a 20% cache budget.The memory reduction enables larger batch sizes and can eliminate the need for offloading in some settings.
  • Ablation and Extended Evaluations: H2O supports inputs up to four million tokens and achieves lower perplexity than StreamLLM across cache sizes.The comparison uses the first text sample of PG-19.
  • Ablation and Extended Evaluations: H2O matches full-model performance across zero-shot to ten-shot inference, while the Local strategy degrades by up to 37.00%.Across downstream tasks, H2O differs from the full model by less than 1.00%.
  • Ablation and Extended Evaluations: Combining H2O with quantization usually outperforms either method alone, and retaining both heavy-hitter and recent tokens restores full-cache performance.Keeping only heavy hitters or only local tokens causes 2.85%–22.75% degradation, while H2O improves generated-text diversity.

6 Conclusion and Discussion

H2O targets the KV cache bottleneck in long-content and large-batch LLM generation by retaining influential Heavy Hitter tokens. The method combines a memory-reducing eviction policy with theoretical guarantees and evaluations showing improved efficiency without compromising generation quality.

  • H2O reduces KV-cache memory by retaining Heavy Hitters, tokens that contribute most to attention computation.The policy is designed for long-content and large-batch generation.
  • H2O formulates KV-cache eviction as a dynamic submodular problem and provides theoretical guarantees for its algorithm.
  • H2O improves end-to-end throughput and decreases wall-clock latency without compromising generation quality across a variety of tasks.

A More Implementation Details

The implementation builds H2O on FlexGen with preallocated KV-cache memory, separate heavy-hitter and recent-token regions, and a circular queue for efficient updates. Experiments vary cache sizes and evaluate summarization, question answering, reasoning, and language-understanding tasks.

  • Implementation Skeleton: The implementation and experiment sections also describe FlexGen's generation loop, attention-forward path, cache buffers, and cache replacement operations.
  • System Details: H2O maintains heavy hitters and recent tokens in a preallocated KV cache, updating recent entries with a circular queue.The implementation keeps the first K entries as heavy hitters and the last K entries as recent tokens.
  • Experiment Details: The evaluation varies KV-cache sizes from 4% to 60% of prompt length and uses 1000 samples for XSUM and CNN/Daily Mail.
  • Experiment Details: Additional lm-eval-harness experiments cover COPA, MathQA, OpenBookQA, PiQA, RTE, and Winogrande.
  • Pseudocode: During prefill, the implementation selects K heavy hitters and K recent tokens when the prompt is at least 2K tokens long.

B.2 Discussions and Limitations

The discussion reports practical and conceptual implications of H2O, including attention-score bias, energy savings, and a dynamic submodular formulation. It also identifies remaining efficiency challenges from the large parameter count of LLMs.

  • Previous Attempts: Accumulated attention scores can favor less recent tokens because earlier tokens receive more attention scores.
  • Previous Attempts: Using averaged attention scores instead resulted in performance degradation, while many Heavy Hitters appeared at sentence beginnings.
  • Social Impact: H2O may reduce energy costs and support green AI by improving the efficiency of LLM generation.
  • Theoretical Framework: The dynamic submodular framework may apply to other dynamic problems involving submodular optimization with slight changes over time.
  • Limitations: LLM generative inference remains challenging because a substantial portion of model parameters resides in MLP blocks.The paper suggests future work on using Heavy Hitters in MLP-aware offloading policies.

C Extended Experiments

Extended experiments examine H2O's effects on text diversity, quantization combinations, extremely long inputs, zero-shot and one-shot inference, and the Top-K baseline.

  • Extended experiments assess text diversity, quantization combinations, extremely long inputs, zero-shot and one-shot inference, and the Top-K baseline.

C.1 Increased Diversity of Generated Text

H2O preserves and can increase generated-text diversity under restricted KV-cache budgets, whereas retaining only recent tokens can produce repetitive or meaningless outputs. Across inference settings, H2O maintains performance with substantially less cache and improves related sparse-cache strategies.

  • Qualitative diversity: H2O generates more diverse content than the Local recent-token strategy for OPT-6.7B and LLaMA-7B despite using less KV-cache budget.The Local strategy repeats punctuation or phrases, while H2O describes varied entities and environments.
  • Quantitative diversity: 0.0051 Self-BELU is achieved by H2O on LLaMA-7B, compared with 0.0057 for the full model and 0.0436 for the local method.The paper interprets H2O’s lower Self-BELU score as slightly increased generated-text diversity.
  • Inference performance: Up to 5× KV-cache reduction is achieved by H2O in zero-shot and one-shot inference while matching full-cache performance; the Local method fails.Some more difficult tasks require a higher 30–40% cache budget.
  • Long-context generation: H2O outperforms StreamLLM on long-context generation and successfully streams to four million tokens.StreamLLM’s fixed initial-token and local-token strategy loses crucial information when it is not retained.
  • Ablations and combinations: Adding H2 to the Top-K strategy improves accuracy by up to 2.00% across four tasks, while ablations show Heavy-Hitters contribute more than local tokens.H2O also maintains full-model performance with a 20% memory budget across different shot counts.

C.10 Heavy-Hitter in MLP Blocks

Heavy-Hitters also emerge among MLP neurons: a small subset is activated extremely often while most neurons are rarely activated. Their positions overlap across input contents, appear early in training, and are important for model performance.

  • Power-law activation: MLP-neuron activation frequencies follow a power-law distribution, with a small subset activated by nearly all input tokens and most neurons rarely activated.The observation uses Wiki-Text-103 with OPT-6.7B.
  • Recovery: Removing MLP Heavy-Hitters causes substantial degradation, but the degradation can be recovered with only 1% of the training data.The paper also reports overlap across input contents and early emergence during training.
  • Functional importance: Removing MLP Heavy-Hitters increases GPT-2 perplexity from 19.32 to 31.78.The pruned neurons are those with activation frequency exceeding 20%.
  • Cross-content overlap: MLP Heavy-Hitter positions show significant concurrence across Wiki-Text-103, Penn Treebank, and Amazon Review inputs.This comparison uses the pretrained OPT-1.3B model.
  • Early-bird property: MLP Heavy-Hitters emerge as early as 4% of the training budget, after which their positions change gradually and minimally.The distribution develops power-law behavior at this early checkpoint.

D Theoretical Analysis

The paper models KV-cache eviction as dynamic submodular maximization and analyzes H2O as a greedy policy under submodularity assumptions. The theoretical framework formalizes attention computation, diminishing returns, and cache-budgeted set updates.

  • Theoretical guarantee: Under a submodular attention assumption, greedily constructing the token set without cache-size limitation is near-optimal in submodular value.The paper presents this as an informal version of Lemma 3.1.
  • Dynamic submodularity: H2O formulates LLM-generation KV-cache eviction as dynamic submodular maximization and provides theoretical guarantees for its approximate greedy policy.The section also describes the policy’s generative-process deployment and efficient implementation.
  • Formal framework: The formal dynamic framework treats the objective as a set function whose submodularity depends on previously selected tokens and the current index.It is presented as a weaker formulation than the stronger framework defined earlier.
  • Recursive attention: The recursive attention definitions use query rows and prefix key matrices to compute per-token normalization and attention quantities before cache eviction.K≤i,∗ selects the first i key rows; D_i is the normalization factor and o_i is the attention of token i.
  • Cache eviction: With a cache budget k<n, the formal H2O eviction policy dynamically maps a source set S_i−1 to a target set S_i using a score function.The policy maintains |S_i|≤k.
  • Diminishing returns: Submodularity captures diminishing returns: adding a token to a smaller covered set can contribute more information than adding it to a larger set.The paper motivates this property using token-associated concepts and neural coverage.

D.9 Robust Greedy with Error Propagation

The section develops greedy selection under exact and approximate values, showing that bounded value errors propagate through dynamic set construction while preserving an explicit approximation guarantee.

  • Greedy selection: The greedy algorithm adds the element with maximum marginal increment until the selected set reaches budget k.This procedure starts from S0 = ∅ and returns Sk after k iterations.
  • Exact-value guarantee: For monotone submodular functions, greedy selection achieves a (1 − 1/e) approximation to the optimum.The guarantee is stated for the set returned by Algorithm 3.
  • Approximate values: Approximate marginal values can be used to expand the set while incurring an additive error of at most ϵ relative to the corresponding greedy guarantee.The resulting performance is at least (1 − 1/e) · opti − ϵ under the stated conditions.
  • Dynamic assumptions: The dynamic analysis assumes monotonicity and dynamic conditions that relate successive set values and optimal costs across iterations.These conditions include fSi,i(Si) ≥ (1 − θ) · fSi−1,i−1(Si) and opti ≥ (1 − γ) · opti+1.
  • Error propagation: Under approximate evaluation, the propagated value bound becomes (1 − 1/e) · (1 − θ)^i+1(1 − γ)^i+1opti+1 − (i + 1) · ϵ0.This bound follows from the approximate monotonicity, dynamic, and value conditions.

D.14 Theoretical Result

This section gives formal guarantees for full-information and approximate greedy eviction, then connects the theory to sparse attention and formulates an optimization framework with a sparsity penalty.

  • Greedy eviction guarantees: Under a mild assumption, full-information greedy attention selection produces sets Si with a formal near-optimality guarantee for any target size k.The guarantee applies when attention scores are computed greedily from full information.
  • Greedy eviction guarantees: Under a mild assumption and budget k, the approximate top-k greedy eviction policy satisfies a formal guarantee for each generated set eSi.This result is identified as the formal version of Theorem 4.4.
  • Attention sparsity: At inference time, attention matrices are reported to be over 95% sparse, with only 5% of the KV cache sufficient to decode the same output token at each generation step.The passage presents this as an observation from Figure 2.
  • Attention sparsity: The sparsity formulation defines an (α, τ, k)-good distribution by bounding thresholded support outside a fixed set S0 by α · k.The condition is stated for every x in the distribution’s domain.
  • Optimization formulation: The section introduces a novel sparsity penalty and combines exponential loss, sparsity loss, and regularization as L(x) := Lexp(x) + Lsparse(x) + Lreg(x).The penalty is described as an ℓ1-type sparsity penalty.
  • Optimization algorithm: A greedy-type approximate Newton algorithm computes an approximate Hessian and updates the parameters with a Newton step over repeated iterations.The procedure subsamples D, computes the gradient exactly, forms A⊤eDA, and updates xt+1.
Loading 2306.14048v3…