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

TL;DR

InfiniGen tackles the KV-cache memory bottleneck in long-context LLM inference on offloading-based systems. It introduces a dynamic KV cache management framework that speculates the next layer's attention using the attention input from the previous layer and a skewed Q/K obtained through offline SVD, fetching only essential KV entries to the GPU while keeping most KV in CPU memory. The approach uses a Prefetching pipeline with a prefill stage selecting a subset of columns and a decoding stage applying a threshold alpha; a CPU-based KV cache pool is managed with a counter-based eviction to respect memory limits. Empirical results on OPT and Llama-2 show up to 3x speedups and up to 32.6 percentage points accuracy improvement, with better scalability for longer sequences and larger models.

Abstract

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.

Figures (20)

• Figure 1: Transformation from matrix $\mathbf{V}^{T}$ to matrix $Q$ in terms of SVD. The orthogonal matrix $A$ maximizes the difference in magnitude between the column vectors of $Q$.
• Figure 2: Total size of the KV cache and model weights of OPT-30B for different sequence lengths and batch sizes. The batch size of (a) is 16, and the sequence length of (b) is 2048. The dotted line represents the size of the model weights.
• Figure 3: Comparison between different execution styles of Transformer blocks.
• Figure 4: Cosine similarity between the attention weights of the base model with full cache and (a) H$_2$O or (b) Optimal. H$_2$O and Optimal use 200 tokens for attention computation. We use OPT-6.7B and a random sentence with 2000 tokens from the PG-19 dataset rae:pot20.
• Figure 5: Histogram that shows the number of key tokens needed to achieve 0.9 out of 1.0 total attention weight for (a) Layer 0 and (b) Layer 18 of the OPT-6.7B model. The bin width is set to 16. We observe that the distribution dynamically changes across the layers.