The "KV Cache" of Modern AI: How Large Models Accelerate Generation Through "Memory"

When conversing with Large Language Models (LLMs), you may have noticed a phenomenon: the speed at which the model generates answers is usually constant, but after you input an extremely long text, the "thinking time" before the model outputs the first character—known as Time to First Token (TTFT)—increases significantly.

Behind this phenomenon lies a crucial engineering optimization technique: KV Cache (Key-Value Cache). It is the cornerstone enabling all mainstream LLMs (such as GPT-4, Claude, and Llama) to achieve efficient real-time generation.

Why Is KV Cache Needed?

To understand KV Cache, we must first understand the Attention mechanism within the Transformer architecture.

During text generation, an LLM operates as an "autoregressive" process: it predicts only the next token at a time. To predict the $(N+1)$-th token, the model must review all previous $N$ tokens.

In standard Attention computation, each token is transformed into three vectors: Query (Q), Key (K), and Value (V).
- Query: Represents "what information I am currently looking for."
- Key: Represents "what information I contain."
- Value: Represents "the specific content I can provide if selected."

Without caching, every time a new token is generated, the model must recompute the K and V vectors for all previous tokens. This means:
- When generating the 1st word, K/V is computed once;
- When generating the 2nd word, K/V is recomputed for the previous 1 word plus the current word;
- When generating the 100th word, K/V is recomputed for the previous 99 words plus the current word.

This computational load grows quadratically ($\mathcal{O}(n^2)$) with sequence length, leading to massive computational waste and severe latency.

How KV Cache Works: Trading Space for Time

The core idea behind KV Cache is simple: Since previous tokens do not change during the input process, their K and V vectors remain identical in every iteration, making repeated calculations unnecessary.

The specific workflow is as follows:
1. Prefill Phase: When you send a prompt, the model processes all input tokens at once, computes their K and V vectors, and stores them in GPU memory (VRAM).
2. Decoding Phase:
- The model computes Q, K, and V only for the current latest token.
- The newly generated K and V are appended to the cache.
- The current Q is used to compute attention weights via dot product with all Ks in the cache: $\text{softmax}(\frac{QK^T}{\sqrt{d}})$.
- A weighted sum of all Vs in the cache is calculated based on these weights to produce the final output.

Through this method, the computational cost per step is reduced from $\mathcal{O}(n^2)$ to $\mathcal{O}(n)$.

The Cost of KV Cache: The "Memory Hog"

While KV Cache significantly boosts speed, it introduces a major challenge: VRAM pressure.

KV Cache needs to store the K and V vectors for every token across every layer and every attention head. The formula for its memory usage is approximately:
$$\text{Memory} = 2 \times \text{layers} \times \text{heads} \times \text{dim} \times \text{seq_len} \times \text{precision}$$

For example, for a typical Llama-3-8B model (FP16 precision), if the context length reaches 32K tokens, the KV Cache alone could occupy several gigabytes of VRAM. This is why many AI services impose limits on context length or encounter OOM (Out of Memory) errors when processing long texts.

How to Optimize KV Cache?

To alleviate VRAM pressure, the industry has developed various optimization strategies:

1. MQA (Multi-Query Attention) and GQA (Grouped-Query Attention)

Traditional Multi-Head Attention assigns one Key head and one Value head to each Query head.
- MQA: All Query heads share a single Key and Value head. This directly reduces the size of the KV Cache by a factor of $H$ (where $H$ is the number of heads).
- GQA: A compromise solution. Query heads are grouped, with each group sharing a pair of KV heads. This significantly reduces memory usage while maintaining model performance (mainstream models like Llama-3 adopt GQA).

2. PagedAttention (vLLM)

Traditional caching requires contiguous memory space, which can lead to fragmentation and waste (similar to early operating system memory management). vLLM's PagedAttention stores the KV Cache in non-contiguous physical memory blocks using paging, greatly improving VRAM utilization and supporting higher concurrency (Throughput).

3. Quantization

Quantizing the KV Cache from FP16 to INT8 or FP8 can directly halve memory usage without significant loss in precision.

Conclusion

KV Cache is one of the key engineering breakthroughs that transformed LLMs from "laboratory toys" into "practical products." By sacrificing VRAM space to eliminate redundant computations, it allows AI to output text smoothly at a speed comparable to human reading. When we discuss the context window size and inference costs of AI, we are essentially discussing how to manage this memory space known as "KV Cache" more efficiently.