KV Cache

The KV cache is the single most important optimization for efficient LLM inference — reusing previously computed key and value vectors from earlier tokens, instead of wastefully recomputing them from scratch at every single new generation step.

The Problem It Solves

Recall autoregressive generation from Seq2Seq Model: generating token 100 requires attending over all 99 previous tokens. Without any caching, generating each new token would mean recomputing key and value vectors (from Query, Key, Value) for the entire sequence so far, from scratch, at every single step — an enormously wasteful amount of redundant computation, since tokens 1 through 99's key/value vectors never actually change once computed.

Why Keys and Values for Past Tokens Never Change

Because of causal masking (see Masked Self-Attention), earlier tokens never attend to later ones — token 5's key and value vectors depend only on token 5's own representation and the tokens before it, never on anything that comes after. This means, critically, that once token 5's key and value vectors are computed, they remain exactly correct forever, no matter how many more tokens get generated afterward — there's no reason to ever recompute them.

The Optimization

This is a substantial, practically essential speedup for any reasonably long generation — without it, generating long outputs would be dramatically, often prohibitively, slower.

Code — The Conceptual Pattern

import torch

class KVCacheDemo:
    def __init__(self):
        self.cached_keys = []
        self.cached_values = []

    def generate_step(self, new_token_embedding, compute_kv_fn):
        new_k, new_v = compute_kv_fn(new_token_embedding)   # compute K, V for ONLY the new token
        self.cached_keys.append(new_k)
        self.cached_values.append(new_v)

        all_keys = torch.stack(self.cached_keys)      # reuse everything computed so far
        all_values = torch.stack(self.cached_values)    # + the one new addition
        return all_keys, all_values   # used for this step's attention computation

# Every real generation step, only ONE new token's K/V is computed --
# everything before it is reused directly from the cache

# In practice, most Hugging Face generation calls enable this by default:
outputs = model.generate(input_ids, max_length=100, use_cache=True)

The Memory Cost Tradeoff

The KV cache trades compute for memory — every generated token's key and value vectors must be stored in GPU memory for the duration of that generation, growing linearly with sequence length. This memory cost is exactly why techniques like grouped-query attention (mentioned in LLM Architecture) were developed — reducing the number of distinct key/value projections directly shrinks how much cache memory each additional token requires.

Common Mistakes

• Assuming the KV cache also caches query vectors — only keys and values are cached, since each new step's query is genuinely new (computed fresh from the newest token) and has nothing analogous to reuse.
• Underestimating KV cache memory as a real production constraint — for very long generations or many simultaneous requests, KV cache memory usage can become a significant, sometimes limiting factor in inference infrastructure design.

Interview Relevance

Q: "Why is the KV cache valid — why don't earlier tokens' key and value vectors need to be recomputed as generation continues?" Causal masking guarantees earlier tokens never attend to later ones, so an earlier token's key and value vectors depend only on itself and tokens before it — nothing that happens later in generation can ever change them. Once computed, they remain exactly correct for the rest of generation, so caching and reusing them (instead of wastefully recomputing) is both valid and a major efficiency win.

Practice Question

Without a KV cache, would generating a 500-token response require more or less total computation than generating a 100-token response, relatively speaking? How does the KV cache change this relationship?

Related DL Notes

• Masked Self-Attention
• Query, Key, Value
• Efficient Inference