Implementing KV Cache in PyTorch

The Two Phases of Inference

To implement KV caching, we must first distinguish between the Prefill and Decode phases. These phases dictate how the model interacts with the cache and processes input tokens.

Welcome to the implementation phase, where we translate the theory of KV caching into functional PyTorch code. In a real-world inference loop, the model operates in two distinct modes. First, the Prefill phase processes the entire input prompt in a single parallel burst to calculate the initial keys and values. Then, the Decode phase takes over, generating one token at a time and incrementally updating the cache. In the Decode phase, the sequence length of our input is always exactly one. We only compute K and V for that new token and then append it to the existing cache tensors. During Prefill, the sequence length is equal to the prompt length. We compute K and V for every token simultaneously and store them to start our 'past_key_values' history.

Prefill: Parallel processing of the prompt to initialize the cache.
Decode: Sequential generation of tokens, one by one.
The cache grows only during the Decode phase.

The Architecture of past_key_values

In PyTorch-based libraries like Hugging Face, the KV cache is stored as a tuple of tuples. This structure mirrors the layer-by-layer organization of the Transformer.

To manage the cache across layers, we use a specific nested structure. Each layer in the Transformer returns its own pair of Key and Value tensors. These tensors are typically 4D, with the sequence length located at dimension 2. This is the dimension that will grow as we generate more text.

Structure: ( (layer_0_K, layer_0_V), (layer_1_K, layer_1_V), ... )
Tensor Shape: (batch_size, num_heads, sequence_length, head_dim)
Concatenation happens on dim=2 (sequence length).

Implementing the Update Logic

The core of the KV cache implementation is the concatenation step. We check for existing states and append new ones along the sequence dimension.

Let's look at the logic inside the attention layer. When computing attention, we check if a previous cache exists. If it's the first step, we simply set our states to the current K and V. But in subsequent steps, we use torch dot cat to append the new token's data to the end of the sequence dimension.

Check if past_kv is None (initial step).
Use torch.cat() to merge old and new states.
Always concatenate on dim=2.

The Generation Loop

Integrating the cache requires updating the input_ids and the past_key_values at every step of the loop.

Now, let's trace a generation loop. Notice how the input changes after the first step. In the first iteration, we pass the full prompt. But look at step two—we only pass the single 'next_token' and the 'past_key_values' from the previous output. This is what makes inference efficient.

Pass use_cache=True to the model.
The model returns the updated cache automatically.
Input only the last token for the next iteration.

Exercise: Complete the Cache Update

Fill in the missing PyTorch logic to update the Key and Value tensors during a decode step. Assume tensors are shaped [B, H, S, D].

It's time to test your implementation skills. Look at this forward function. I've left the update logic blank. Complete the code to concatenate the previous keys and values with the current ones along the sequence dimension.

Concatenate previous and current states.
Target the correct dimension for sequence length.

Common Pitfalls & Memory Bandwidth

While KV caching saves compute, it introduces new bottlenecks like Memory Fragmentation and Bandwidth saturation.

Implementing KV caching isn't without its challenges. Even though we save trillions of floating-point operations, we are now limited by how fast we can move data. During the decode phase, the GPU spends more time waiting for the large KV cache to load from memory than it does actually computing. Furthermore, repeated concatenation can fragment your GPU memory, leading to 'Out of Memory' errors even when you have space—a problem solved by modern techniques like PagedAttention.

Memory Fragmentation: Frequent 'torch.cat' calls create new memory allocations.
Bandwidth Bound: Moving large cache tensors from GPU memory to the cores is slow.
Solutions: PagedAttention (vLLM) manages memory more efficiently.