KV Cache Memory › Anatomy of KV Cache Storage

Michael BrenndoerferJanuary 7, 20265 min read
Language AI HandbookAll chapters
Part 42 of 67Inference Optimization
Chapter 314KV Cache Memory › Anatomy of KV Cache Storage
Reading section 2 of 14 · Anatomy of KV Cache Storage · 5 minSection 2 of 14 · 5 min

Book contents

Section 2 of 14 · 5 min read

KV Cache Memory › Anatomy of KV Cache StorageOne section · Sections in this chapter2 / 14 ↓
Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Before calculating memory requirements, we need to understand precisely what the cache stores. Recall from the previous chapter that during each forward pass, we cache the key and value projections for reuse in subsequent tokens. This caching strategy turns what would be repeated linear projections into simple memory lookups. However, every lookup requires that the data exist in GPU memory, which means someone must allocate space for it before the first token is ever generated. Understanding the precise shape and structure of this stored data is the first step toward calculating its memory footprint.

The keys and values we cache are not arbitrary tensors. They have very specific shapes determined by the attention mechanism's design. When attention processes a token, it creates a query vector to ask "what should I attend to?" Corresponding key and value vectors for that token answer questions from future tokens. The query is used immediately and discarded, but the keys and values must persist because future tokens will compare against them. Think of keys as index cards in a filing cabinet and values as the documents those cards refer to. Every token you generate files a new index card. Every subsequent token must search the entire filing cabinet before it can produce its output.

For a single attention layer processing a single token, the projections produce:

  • Key tensor: shape [batch_size, num_kv_heads, 1, head_dim]
  • Value tensor: shape [batch_size, num_kv_heads, 1, head_dim]

The "1" in these shapes represents the single token being processed during generation. This single-token slice is what gets appended to the existing cache with each forward pass. The batch dimension allows processing multiple independent sequences simultaneously. The head dimension captures the multi-head structure of attention, where different heads specialize in attending to different aspects of the input. The head_dim represents the dimensionality of each head's representation space, controlling how expressive each head's keys and values can be.

After processing tt tokens, the accumulated cache for that layer has:

  • Key cache: shape [batch_size, num_kv_heads, t, head_dim]
  • Value cache: shape [batch_size, num_kv_heads, t, head_dim]

Notice how the sequence dimension grows from 1 to tt as tokens accumulate. This growth is the source of the linear memory scaling that dominates long-context inference. Each new token adds one slice along the sequence dimension, which must be stored for all subsequent generation steps. There is no compression or summarization: every token's key and value vectors survive intact until the generation completes.

The full KV cache spans all transformer layers. A model with LL layers maintains LL key caches and LL value caches. This multiplicative effect of layers is important. A 32-layer model requires 32 times the cache of a single layer. This makes the number of layers one of the primary drivers of total cache size. Each layer's cache is completely independent, containing that layer's unique representation of the sequence history. The attention mechanism at layer 12 attends over different representations than the attention at layer 24, and both sets of representations must be preserved simultaneously.

The data type determines bytes per element. Most modern inference uses half-precision (FP16 or BF16) at 2 bytes per element, though some systems use FP32 (4 bytes) or quantized representations like INT8 (1 byte) or INT4 (0.5 bytes). This choice creates a direct tradeoff between precision and memory consumption. Halving the bytes per element halves the cache size, which is why quantization techniques have become essential for deploying large models. We will revisit this tradeoff in the quantization chapters.

Out[3]:
Visualization
Line plot showing cumulative KV cache memory growing linearly with number of generated tokens.
KV cache memory growth during LLaMA 7B token generation across 4,096 tokens with batch size 1 in FP16 precision. Memory grows linearly from 0 to 2.0 GB, with each token adding 0.50 MB across all 32 layers. The shaded region emphasizes the cumulative nature of cache growth: every generated token increases the memory footprint for all future tokens in the sequence.