KV-Cache: Accelerating Autoregressive Inference

The Inefficiency of Stateless Autoregressive Generation

As you know, autoregressive models like GPT generate text one token at a time. To predict the next token, the model needs to consider the entire sequence generated so far.

Let's break down the "naive" or stateless generation process:

1. Input: The model receives the entire sequence of tokens generated so far (e.g., "The cat sat on").
2. Process: It runs a full forward pass, computing the attention scores for every token in the sequence.
3. Output: It uses the prediction for the very last token to choose the next one (e.g., "the").
4. Append & Repeat: The new token is appended to the sequence ("The cat sat on the"), and the entire process repeats from step 1.

The inefficiency lies in step 2. To generate just one new token, the model re-computes the Key (K) and Value (V) vectors for all the previous tokens, even though it just calculated them in the previous step. This is a massive amount of redundant work.

The following video introduces the LLM inference workload and highlights this one-token-at-a-time process, setting the stage for understanding the inefficiency.

Quantifying the Waste: The "Triangle of Waste"

The video you just watched mentions storing tokens on the GPU. The core of the problem is how that stored information is used. In a stateless model, the K and V vectors are recalculated from scratch every time. This computational load grows quadratically with the sequence length, leading to O(T²) complexity for generating T tokens.

Let's visualize this waste.

As you saw, for each new token, we only truly need to compute the K and V vectors for that single new token. All the previous K and V vectors remain unchanged. This insight is the key to the solution.

The Solution: Caching Keys and Values

The KV-cache is a simple but powerful optimization that eliminates this redundant work. The idea is to store, or "cache," the Key and Value tensors for each token after they are computed.

The generation process is now split into two phases:

1. Prefill (or Prompt Processing): When the model first receives the prompt (e.g., "The cat sat"), it performs one initial forward pass. It computes the K and V vectors for all tokens in the prompt and stores them in the cache.
2. Decode (or Generation): For every subsequent token generation step:
   • The model only takes the single, most recent token as input.
   • It computes the Query (Q), Key (K), and Value (V) vectors for only this one token.
   • It retrieves the previously computed K and V vectors from the cache.
   • It concatenates the new K and V vectors with the cached ones.
   • The new Q vector attends to this full history of K and V vectors to predict the next token.
   • The newly computed K and V are appended to the cache for the next step.

By doing this, the amount of new work at each step becomes constant (just processing one token), transforming the overall complexity from quadratic O(T²) to linear O(T).

Let's see how this efficient workflow looks in practice.

The Trade-Off: Computation vs. Memory

The KV-cache offers a dramatic speedup, but it comes at a cost: memory. The cache needs to store the Key and Value tensors for every token in the sequence, across every attention head and every transformer layer.

The memory footprint of the KV-cache can be calculated as:

$\text{Cache Size (bytes)}=2\times L\times N_h\times T_s\times d_h\times P$

Where:

• 2: For storing both Keys and Values.
• $L$: Number of decoder layers.
• $N_h$: Number of attention heads.
• $T_s$: Sequence length (context window).
• $d_h$: Dimension of each attention head.
• $P$: Bytes per parameter (e.g., 2 for float16, 4 for float32).

As you can see, the memory usage grows linearly with the sequence length. This is a primary reason why models with very long context windows require GPUs with massive amounts of VRAM.

This article provides a great summary of the advantages and disadvantages.

From Concept to Code

While you don't need to implement it now, your background in engineering makes it valuable to see how the KV-cache translates to code. The core logic involves modifying the forward pass of the attention module to accept an optional cache argument.

• If the cache exists, concatenate the newly computed K and V vectors to it.
• If not (i.e., the prefill step), initialize the cache with the new K and V.
• Return the model's output and the updated cache.

Most modern libraries like Hugging Face's transformers handle this automatically. When you call the .generate() method, the use_cache=True parameter is enabled by default, transparently providing this massive speedup.

The video below briefly shows the code diffs and data flow, which should be intuitive given your PyTorch experience.

Conclusion

You have now mastered a fundamental optimization for LLM inference. The KV-cache is not a complex architectural change but a clever engineering solution that makes interactive, low-latency generation feasible.

Key Takeaways:

• The Problem: Stateless autoregressive inference is inefficient, with O(T²) complexity due to redundant re-computation of Key and Value vectors for the entire sequence at every step.
• The Function: The KV-cache stores previously computed Key and Value tensors, avoiding this re-computation.
• The Process: Inference is split into a Prefill step (processing the prompt) and subsequent Decode steps (processing one token at a time).
• The Result: The computational complexity of generation is reduced from quadratic to linear, O(T), enabling significant speedups.
• The Trade-Off: This speed comes at the cost of increased memory usage, as the cache size grows linearly with the sequence length.

What's Next?

The memory cost of the KV-cache is a significant bottleneck. So, the natural next question is: can we reduce it? In our next lesson, we will explore Grouped-Query Attention (GQA) and Multi-Query Attention (MQA), two architectural modifications designed specifically to shrink the size of the Key and Value tensors, and thus the KV-cache, without a major loss in performance. This is a perfect example of how architectural design and inference optimization go hand-in-hand.