The KV cache is a memory structure used during autoregressive inference (text generation) in large language models. It stores previously computed key (K) and value (V) vectors from the attention mechanism so that they can be reused when generating subsequent tokens, eliminating the redundant recomputation that would otherwise occur at every generation step. This yields substantial speedups – often 5× or more – at the cost of increased memory consumption and added implementation complexity. KV caching is enabled by default in popular inference libraries such as Hugging Face Transformers.
Key Points
Without caching, each generation step recomputes key and value vectors for every token in the prompt and all previously generated tokens, leading to O(n²) compute per step.
The KV cache stores the keys and values from earlier tokens and concatenates only the newly computed vectors, reducing per-step compute to O(n) (linearly proportional to the current sequence length).
The cache is implemented as a pair of buffers (cache_k, cache_v) inside the multi-head attention layer.
On a Mac Mini M4 (CPU) with a 124M parameter model generating 200 tokens, a 5× speedup was measured.
On a T4 GPU generating 300 tokens from a 1.7B‑parameter model, KV caching yielded a 5.21× speedup (11.7 seconds vs. 1 minute 1 second for standard inference).
As sequence length grows, memory consumption increases linearly because the cache retains all previous tokens; truncation (sliding window) or pre-allocation can mitigate this.
The cache must be reset between separate generation sessions to avoid stale data.
In practice, torch.cat for building the cache causes memory fragmentation; production systems often pre-allocate a fixed-size tensor or use a sliding window.
On GPU with small models, the overhead of transferring tensors and managing the cache can negate the benefit; compiled models without explicit caching may perform better.
Concepts
Autoregressive generation: The model produces tokens one at a time, conditioning on all previously generated tokens.
: In scaled dot‑product attention, keys represent the “identifier” of a token; attention scores measure how well a query matches each key.
Key (K)
Value (V): Values carry the actual information from each token; the attention output is a weighted sum of values.
Query (Q): A representation of the current token used to compute attention weights against all keys.
KV cache: A data structure that stores the concatenated key and value tensors produced by previous decoding steps, avoiding recomputation.
Caching strategy: The first forward pass on the full prompt populates the cache. Each subsequent step computes keys and values only for the newest token and appends them to the cache. The model then uses the entire cached K and V for attention.
Position tracking: Because cached keys and values are in token order, the queries for new tokens need to be computed at positions that follow the cached ones. A current_pos counter (or checking the cache shape) ensures correct positional encoding.
Details
The KV cache is integrated into a standard multi‑head attention module. In the constructor two buffers are registered:
The forward method accepts a use_cache flag. When True, the newly computed keys and values (keys_new, values_new) are concatenated onto the existing cache:
A reset_cache method sets both buffers back to None to prevent cross‑session interference. The parent model (GPTModel) maintains self.current_pos as a global position counter. During forward passes with use_cache=True, the position IDs start at current_pos and advance by seq_len; otherwise they start at 0. This ensures that each new query aligns with the correct cached keys and values.
The generation function generate_text_simple_cached orchestrates the process:
Reset the cache.
Run the full prompt once with use_cache=True (populates the cache).
For each new token to generate, feed only the single previous token (model(next_idx, use_cache=True)), append the output to the sequence, and repeat.
In contrast, the uncached version re‑feeds the entire truncated sequence at every step.
A simplified PyTorch implementation of a standalone KV cache class is also common:
class KVCache:
def __init__(self):
self.cache = {"key": None, "value": None}
def update(self, key, value):
if self.cache["key"] is None:
self.cache["key"] = key
self.cache["value"] = value
else:
self.cache["key"] = torch.cat([self.cache["key"], key], dim=1)
self.cache["value"] = torch.cat([self.cache["value"], value], dim=1)
def get_cache(self):
return self.cache
In Hugging Face Transformers, KV caching is enabled via the use_cache=True parameter (default) in model.generate().
Comparison: Standard vs. Cached Inference
Feature
Standard Inference
KV Caching
Computation per step
Recomputes all keys and values from scratch
Computes only keys/values for the new token
Memory usage
Lower per step, but total grows with sequence length
Higher up‑front memory to store cache
Speed
Slows down as sequence lengthens
Roughly constant per step (O(1))
Long‑text handling
Becomes prohibitively slow
Practical for long sequences
Performance and Scaling
Compute: Per‑step cost drops from O(n²) to O(n) because each key/value is computed only once.
Memory: The cache grows by the number of tokens generated, each token contributing num_heads * head_dim entries for both K and V. For long sequences or large models this can become prohibitive.
Measured improvements:
A 124M parameter model on a Mac Mini M4 (CPU) with a 4‑token prompt and 200‑token generation showed roughly 5× speedup. Both cached and uncached outputs were identical (gibberish due to untrained model).
On a T4 GPU generating 300 tokens from a 1.7B‑parameter model, KV caching yielded a 5.21× speedup (11.7 seconds vs. 1 minute 1 second for standard inference).
Optimizations
The naive torch.cat causes repeated memory allocations and fragmentation. Two practical improvements:
Pre‑allocation: Reserve a tensor of fixed maximum length (e.g., torch.zeros(batch_size, num_heads, max_seq_len, head_dim)) and write into slices. This avoids repeated reallocation.
Sliding window cache: Keep only the most recent window_size tokens by slicing: cache_k = cache_k[:, :, -window_size:, :]. This bounds memory growth.
The optimized version (gpt_with_kv_cache_optimized.py) uses both techniques. On a Mac Mini M4 CPU with window size equal to the full context length, the speed comparison held. However, on CUDA devices with small models, the memory allocation and transfer overhead cancels the benefit; there, a compiled model without explicit KV cache may be faster.
In experiments with Qwen3 (0.6B) and Llama 3 (1B) models, the torch.cat approach was retained to avoid ~8 GB of extra memory for large context sizes. The KV cache was moved outside the model to enable torch.compile. On CPUs the cache provided substantial speedup, further boosted by compilation. On GPUs, the best performance came from a regular compiled model without an explicit KV cache – likely because the models were small and pre‑allocation on GPU was not ideal.
Memory‑efficient variants (e.g., sliding window, sparse caching) are an active area of research.