KV-Cache Optimization and AI Memory Bottlenecks

KV-cache optimization is a set of techniques used to compress, manage, and efficiently retrieve stored attention states during large language model (LLM) inference. When AI models generate text, they normally have to recalculate the relationships between all previous words for every new word they produce. The Key-Value (KV) cache prevents this redundant computation by storing previous calculations in memory. However, as prompts and responses grow longer, this cache becomes massive, shifting the bottleneck from computation power to memory bandwidth. Optimizing the KV-cache shrinks the size of this stored data and organizes it better, drastically reducing the amount of information that needs to be transferred between GPU memory and processing cores, which directly accelerates AI inference speeds.

Understanding the KV-Cache in AI Inference

Large language models generate text token by token in an autoregressive manner. To predict the next token, the model's attention mechanism calculates Query (Q), Key (K), and Value (V) vectors for the sequence.

Without a cache, generating the 100th token requires the model to recalculate the Key and Value vectors for the previous 99 tokens all over again. The KV-cache solves this by computing the K and V vectors for each token just once and saving them in the GPU's High Bandwidth Memory (HBM). For every subsequent token generated, the model only calculates the K and V for the new token and appends it to the existing cache.

The Memory Bandwidth Bottleneck

While the KV-cache saves immense computational effort, it introduces a severe memory problem. AI inference is divided into two phases: the prefill phase (processing the initial prompt) and the decoding phase (generating the response one token at a time).

The decoding phase is strictly memory-bound. For every single token the model generates, the GPU must transfer the entire model's weights and the entirety of the growing KV-cache from the HBM to the GPU's compute cores (SRAM). The actual math required to predict the next token takes very little time; the GPU spends the vast majority of its time simply waiting for the KV-cache data to travel across the memory bus. As context lengths reach tens or hundreds of thousands of tokens, the KV-cache grows linearly, saturating the memory bandwidth and slowing generation speeds to a crawl.

How KV-Cache Optimization Solves the Bottleneck

KV-cache optimization directly targets the memory bandwidth bottleneck by reducing the total volume of data that must be moved and eliminating wasted memory space. Because less data needs to be fetched from the HBM for each token generation step, the GPU spends less time waiting on data transfers and more time computing.

1. PagedAttention

Traditionally, systems pre-allocated a contiguous block of memory for the KV-cache based on the maximum possible length of a sequence. Because most responses do not reach the maximum length, up to 80% of this reserved memory went unused, creating fragmentation and limiting how many users could be served simultaneously. PagedAttention, inspired by virtual memory in operating systems, breaks the KV-cache into small, non-contiguous blocks. Memory is allocated only as needed. This eliminates wasted space, allowing larger batch sizes and reducing the overhead of memory management.

2. Architectural Changes (MQA and GQA)

Standard Multi-Head Attention (MHA) requires a separate Key and Value head for every Query head, resulting in a massive cache.

Multi-Query Attention (MQA): Forces all Query heads to share a single Key head and a single Value head. This drastically shrinks the KV-cache size and the bandwidth required to read it.
Grouped-Query Attention (GQA): A middle ground where Query heads are divided into groups, and each group shares a single Key and Value head. This maintains high model accuracy while still cutting KV-cache memory transfers by a factor of 4 to 8.

3. KV-Cache Quantization

By default, KV-cache values are stored in 16-bit floating-point formats (FP16 or BF16). Quantization compresses these values into lower-precision formats, such as 8-bit (INT8 or FP8) or even 4-bit integers. Halving the bit-width from 16-bit to 8-bit immediately cuts the memory footprint in half. Consequently, it requires half the memory bandwidth to load the cache into the compute cores, effectively doubling the speed of the memory-bound portion of inference.

4. Token Eviction and Pruning

Not all tokens in a prompt are equally important for generating future text. Token eviction strategies, such as StreamingLLM or H2O (Heavy-Hitter Oracle), monitor attention scores and dynamically delete the K and V vectors of unimportant tokens from the cache. By keeping only the initial tokens (which anchor the model's attention) and the most recent or most relevant tokens, the KV-cache size remains capped at a fixed limit, preventing memory bandwidth degradation regardless of how long the conversation lasts.