InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management

Wonbeom Lee; Jungi Lee; Junghwan Seo; Jaewoong Sim

Source-linked AI summary

InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management

Wonbeom Lee, Jungi Lee, Junghwan Seo, Jaewoong Sim

arXiv:2406.19707v1Jun 28, 2024cs.LGcs.DC

TL;DR

Long-context LLM inference is constrained by KV-cache memory growth and the cost of transferring offloaded cache data to the GPU. InfiniGen speculatively identifies important next-layer tokens and prefetches only their KV entries. It achieves up to 3.00× speedup over existing KV-cache management methods while improving accuracy by up to 32.6 percentage points.

Problem
KV-cache size grows with sequence length and batch size, while transferring an offloaded cache from CPU memory to GPU creates a major inference bottleneck.
Method
InfiniGen uses a minimal rehearsal of the next layer’s attention pattern to prefetch essential KV entries while retaining the cache pool in CPU memory.
Results
Up to 3.00× speedup and a 32.6 percentage point accuracy increase were achieved over existing KV-cache management methods.
Takeaways & Limitations
InfiniGen provides substantially shorter inference latency while preserving language-model performance in offloading-based long-content generation.
Takeaways & Limitations
A fixed KV-cache budget can be insufficient because later queries may require more retained tokens as generation proceeds.

Abstract

from arXiv · show

Transformer-based large language models (LLMs) demonstrate impressive performance across various natural language processing tasks. Serving LLM inference for generating long contents, however, poses a challenge due to the enormous memory footprint of the transient state, known as the key-value (KV) cache, which scales with the sequence length and batch size. In this paper, we present InfiniGen, a novel KV cache management framework tailored for long-text generation, which synergistically works with modern offloading-based inference systems. InfiniGen leverages the key insight that a few important tokens that are essential for computing the subsequent attention layer in the Transformer can be speculated by performing a minimal rehearsal with the inputs of the current layer and part of the query weight and key cache of the subsequent layer. This allows us to prefetch only the essential KV cache entries (without fetching them all), thereby mitigating the fetch overhead from the host memory in offloading-based LLM serving systems. Our evaluation on several representative LLMs shows that InfiniGen improves the overall performance of a modern offloading-based system by up to 3.00x compared to prior KV cache management methods while offering substantially better model accuracy.

1 Introduction

InfiniGen addresses the growing KV-cache memory and transfer costs of long-text LLM inference by speculatively identifying and prefetching only critical entries. It integrates dynamic KV-cache management with CPU offloading and reports up to 3.00× faster performance with better accuracy.

Motivation: KV-cache size grows with output sequence length and batch size, often exceeding model-weight memory and making long-content generation increasingly difficult.The cache stores keys and values for preceding tokens to avoid repeated computation.
Motivation: Offloading the KV cache to CPU memory extends capacity beyond GPU limits but makes CPU-to-GPU transfers a performance bottleneck.The transferred cache can reach hundreds of gigabytes, while limited PCIe bandwidth increases Transformer-block execution time.
Approach: InfiniGen speculates which KV-cache entries are critical for the next attention layer using a minimal rehearsal at the preceding layer.The framework uses the current layer’s attention input together with partial query weights and key-cache information for the subsequent layer.
Approach: InfiniGen retains the KV-cache pool in CPU memory and dynamically removes infrequently used token entries while prefetching essential values to the GPU.Offline weight manipulation emphasizes important query and key channels to make speculation more efficient and precise.
Evaluation: Up to 3.00× faster performance and a 32.6 percentage point accuracy increase were achieved over existing KV-cache management methods.The evaluation covered two representative LLMs across varying model sizes, batch sizes, and sequence lengths.

2 Background

The background describes autoregressive Transformer inference, KV caching, and SVD-based matrix transformations. It motivates InfiniGen’s strategy of using skewed query and key representations to predict important attention tokens efficiently.

Large Language Models: Transformer blocks layer-normalize inputs, project them into query, key, and value matrices, and compute attention as softmax(QK^T)V.The matrices are reshaped across H attention heads with head dimension d, where D = H × d.
Generative Inference and KV Caching: Generative inference has prefill and decoding stages, with each decoding iteration feeding the newly generated token back into the model.Prefill summarizes the prompt and produces the initial decoding input.
Generative Inference and KV Caching: KV caching avoids recomputing previous-token keys and values, but its dimension grows with generated tokens and its size also scales linearly with batch size.At iteration i, the cache dimension is H × (N+i) × d.
Large Language Models: LLM inputs contain large-magnitude outliers concentrated in a few fixed channels across layers.These outliers are associated with intrinsic model properties such as large layer-normalization weights.
Singular Value Decomposition: InfiniGen skews query and key matrices so a small number of channels dominate, allowing important tokens to be predicted from those channels.The transformation can make one transformed query vector much larger than another, supporting attention-score prediction with fewer dimensions.
Singular Value Decomposition: SVD factors a matrix into orthogonal transformations and a diagonal scaling matrix, providing the basis for selecting an orthogonal matrix that separates column magnitudes.For Q = UΣV^T, U and V are orthogonal while Σ contains the singular values.
Singular Value Decomposition: Figure 2 contrasts constant model-weight size with KV-cache growth across sequence lengths and batch sizes.The figure fixes batch size at 16 in one plot, sequence length at 2048 in the other, and marks model-weight size with a dotted line.

3 Motivation

Long-text generation makes the KV cache a major memory and data-transfer bottleneck, especially when caches are offloaded to CPU memory. Existing eviction and fixed-budget strategies struggle because attention patterns and required cache sizes vary across iterations, layers, and queries.

KV Cache in LLM Inference Systems: The KV cache scales with sequence length and batch size, becoming a key memory consumer beyond model weights during long-sequence or batched generation.The model size remains constant while KV-cache size grows linearly with sequence length and batch size.
KV Cache in LLM Inference Systems: Offloading the KV cache to CPU memory expands feasible batch sizes and sequence lengths but makes CPU-to-GPU transfer a performance bottleneck.Transferring hundreds of gigabytes over limited PCIe bandwidth increases Transformer-block execution time.
KV Cache in LLM Inference Systems: Conventional prefetching hides only part of cache-load latency, while quantization does not remove the cache’s linear growth with sequence length.These limitations motivate intelligent KV-cache management that reduces the amount of data loaded.
Challenges in KV Cache Management: Prior eviction methods assume attention patterns persist across iterations, but tokens unimportant now can become important later, reducing similarity to full-cache attention.H2O remains similar for roughly 200 iterations, then performs worse than Optimal as sequences exceed the cache budget.
Challenges in KV Cache Management: The number of key tokens needed to capture 0.9 of total attention varies across layers and queries, so fixed cache budgets can underrepresent the baseline attention pattern.For Layer 18, examples range from 80 to 164 tokens across later queries, while adjacent queries can require markedly different counts.

4 InfiniGen Design

InfiniGen manages long-context KV caches by keeping most entries in CPU memory and speculatively prefetching only tokens likely to matter for the next attention layer. It exploits similarity between consecutive-layer inputs and skewed query/key matrices to reduce transfers while preserving output quality.

4.1 Overview: InfiniGen keeps most KV-cache tokens in CPU memory and dynamically loads only a few important keys and values for GPU attention.This expands the effective selection window without discarding most cached tokens, unlike prior approaches.
4.1 Overview: The runtime combines controllers for partial-weight generation, KV selection, inference, and offline weight skewing with a CPU-memory pool manager.These components coordinate critical-entry prefetching and manage the cache under CPU memory pressure.
4.2 Prefetching Opportunities: Consecutive attention layers have highly similar inputs, allowing Layer i−1 attention inputs to speculate on Layer i attention patterns.Outliers and LayerNorm are identified as reasons for this similarity; the block input is more similar to the preceding block input than to preceding attention or FFN outputs.
4.2 Prefetching Opportunities: InfiniGen multiplies query and key weight matrices by the same orthogonal matrix to concentrate attention-relevant information in fewer columns without changing QK^T.The orthogonal matrix is obtained from query-matrix SVD, and the resulting skewing is applied offline without changing weight dimensions.
4.3 Efficiently Prefetching KV Cache: After partial query and key-cache multiplication, InfiniGen selects tokens whose speculative attention scores exceed the maximum score minus alpha.This adaptive threshold avoids requiring a fixed number of fetched tokens and reduces PCIe transfers and loading latency while maintaining output quality similar to a full KV cache.
4.3 Efficiently Prefetching KV Cache: The method reduces KV-cache loading and computation overhead while using only the PCIe bandwidth required by the selected tokens.Its counter-based and LRU-based variants show comparable model accuracy, while the counter-based design avoids atomic updates for better parallelism.

5 Evaluation

InfiniGen is evaluated across models, cache budgets, sequence lengths, batch sizes, and model sizes. It preserves accuracy while reducing KV-cache transfer costs and improving latency and scalability over prior methods.

Accuracy: InfiniGen consistently preserves accuracy below 10% relative KV cache size, while quantization and H2O show noticeable accuracy drops.Above 10%, InfiniGen closely matches the full-cache baseline.
Sequence Length: InfiniGen maintains perplexity comparable to the full-cache baseline as sequence length increases, while H2O increasingly diverges.H2O uses a fixed budget, whereas InfiniGen dynamically computes attention with only the essential cache.
Inference Latency: 1.63×-32.93× speedups over baselines are achieved by InfiniGen on OPT-13B with 1920 input tokens, 128 output tokens, and batch size 20.The gain mainly comes from reducing KV-cache data loaded from CPU memory.
Batch Size: As batch size increases, InfiniGen’s performance gap widens because its dynamic cache loading avoids the transfer costs that dominate FlexGen latency.InfiniGen throughput rises from 27.36 to 41.99 tokens per second as batch size increases from 4 to 20.
Model Size: InfiniGen outperforms other methods across model sizes, with speedup increasing by 1.17× from 6.7B to 13B and reaching 1.34× over FlexGen for 30B.For 30B, 30% of model parameters are offloaded to CPU memory.

6 Analysis and Discussion

The analysis examines sensitivity, overhead, memory consumption, long-context perplexity, and million-token attention patterns. These results explain how InfiniGen balances accuracy, storage, and transfer costs through dynamic cache selection.

Sensitivity Analysis: Increasing alpha fetches more KV entries, increasing inference latency while improving accuracy.Alpha determines the threshold for loading entries based on speculated attention scores.
Sensitivity Analysis: Partial weight ratio has negligible impact on inference latency but increases memory consumption for partial weights and key cache.The amount of transferred KV cache is independent of the partial weight ratio.
Prefetching Overhead: 96.9% and 91.8% of FlexGen and H2O execution time, respectively, are occupied by data transfer overhead.InfiniGen is only 1.52× slower than Ideal, while other methods show 3.90×-18.55× slowdowns.
Long Context: InfiniGen keeps perplexity close to the full-cache baseline as relative KV cache size decreases on Llama-2-7B-32K.Other methods diverge because of insufficient quantization precision or permanent KV-cache removal.
Million-Token Contexts: As sequence length grows, more query tokens attend to less than 1% of key tokens, enabling InfiniGen to dynamically adjust loaded cache size.Attention spikes can reappear after long intervals, so permanently removing temporarily unimportant tokens may lose later-critical context.

7 Related Work

Related work covers serving-system optimization, offloading, KV-cache management, and efficient inference. InfiniGen complements these approaches by selectively prefetching essential KV entries during decoding.

DNN Serving Systems: Prior serving systems improve latency or throughput through predictable SLOs, preemption, fine-grained batching, and memory optimizations.These methods address system-level serving efficiency rather than InfiniGen’s selective KV-cache prefetching.
DNN Serving Systems: FlexGen offloads weights and KV cache to CPU memory and disk, while InfiniGen works with FlexGen to improve KV-cache offloading and prefetching.InfiniGen is presented as orthogonal to FlexGen.
KV Cache Management: vLLM and StreamingLLM manage fragmentation, duplication, or context length but do not reduce KV-cache size, leaving offloading transfer overhead significant.InfiniGen targets this transfer overhead directly.
Efficient LLM Inference: InfiniGen can combine with kernel fusion and is described as the first approach to prefetch only essential KV entries in offloading-based inference systems.Kernel fusion primarily mitigates quadratic attention-memory overhead during prefill.

8 Conclusion

InfiniGen addresses KV-cache scalability in offloading-based LLM inference by speculatively identifying important tokens and prefetching their cache entries. It shortens latency while preserving language-model performance and scales better across workloads than prior solutions.

Conclusion: InfiniGen uses the previous layer’s attention input to speculatively prefetch important KV-cache entries and manipulates query and key weights for efficient speculation.The framework is designed for offloading-based inference systems.
Conclusion: InfiniGen substantially shortens inference latency while preserving language-model performance.It also provides better scalability with batch size, sequence length, and model size than prior solutions.