The KV Cache Mechanism

The Autoregressive Redundancy

In standard autoregressive generation, each new token requires the context of all previous tokens. Without optimization, a Transformer would recompute hidden states for the entire sequence at every step, leading to quadratic time complexity.

Welcome. To understand why we need KV caching, we first must look at the inefficiency of standard autoregressive generation. As the model generates each new token, it must look back at everything it has already produced. Without a cache, the Transformer re-processes every single previous token to generate the next one, causing the computational cost to explode as the sequence grows.

Prefill vs. Decode Phases

LLM inference is split into two distinct phases: Prefill and Decode. Each phase interacts with hardware resources differently.

Let's distinguish between the two phases of inference. First is the Prefill phase, where the entire prompt is processed in parallel. This is compute-bound, meaning the GPU's arithmetic units are working at full capacity. Next is the Decode phase, where tokens are generated one by one. Here, the bottleneck shifts from raw math to memory bandwidth as we fetch cached data. During Decode, we only calculate Q, K, and V for the single new token. We then reach into the cache to grab the historical K and V tensors. In Prefill, we compute and store the Key and Value tensors for every token in the prompt. This 'warms up' the cache for the generation steps to follow.

Architectural Mechanism

The KV cache avoids recomputing the Key (K) and Value (V) tensors for historical tokens by storing them in memory and concatenating them with the current step's tensors.

Let's walk through the architectural workflow at step 'n'. First, we take the current token and compute its Query, Key, and Value tensors. We don't re-calculate these for previous tokens. Instead, we retrieve the historical Keys and Values from the cache. We concatenate them together and update the cache with our new K and V. Finally, we perform attention using the single new Query against the full sequence of Keys and Values.

Complexity: O(N²) vs O(N)

By caching, we transform the per-token complexity of the attention mechanism from quadratic to linear relative to the sequence length.

Visualizing the complexity shift is crucial. Without caching, every new token requires re-processing the entire sequence, leading to this steep quadratic curve. With KV caching, the work per token increases only linearly. Use the slider to see how the computational gap widens as the context length grows.

The Memory Bandwidth Bottleneck

While we save FLOPs, we introduce a memory bandwidth bottleneck. The GPU cores (ALUs) often wait for data to be moved from VRAM.

Even though we are doing less math, the GPU might still be slow. Why? Because the arithmetic logic units, or ALUs, are incredibly fast. However, the KV cache is stored in High Bandwidth Memory, or HBM. During the decode phase, the ALUs sit idle while waiting for the massive cache tensors to travel across the memory bus. This is the 'Memory Bandwidth' bottleneck.

KV Cache in PyTorch

In PyTorch, the KV cache is managed through the past_key_values argument. Developers must pass the cache back into the model to maintain state.

Let's look at a practical implementation using Hugging Face Transformers. Notice the 'past_key_values' variable. By setting 'use_cache' to True, the model outputs the current KV tensors. To generate the next token, we pass those same tensors back in. Try completing the code to handle the cache update.