Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Multi-Query Attention › The KV Cache Memory Problem

Michael BrenndoerferUpdated August 4, 20254 min read

Language AI HandbookAll chapters

Part 29 of 67Modern Decoder Models

Chapter 210Multi-Query Attention › The KV Cache Memory Problem

Reading section 2 of 11 · The KV Cache Memory Problem · 4 minSection 2 of 11 · 4 min

Section 2 of 11 · 4 min read

Multi-Query Attention › The KV Cache Memory ProblemOne section · Sections in this chapter2 / 11 ↓

Sections in this chapter11

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 diving into MQA, we need to understand why KV cache memory is such a critical concern for large language model inference. During autoregressive generation, models cache the keys and values computed for previous tokens to avoid redundant computation. This cache grows with every token generated.

For a model with $L$ layers, $h$ attention heads, head dimension $d_k$ , and sequence length $T$ , the KV cache size is:

\text{KV Cache Size} = 2 \times L \times h \times T \times d_k \times \text{bytes per element}

where:

$2$ : accounts for storing both keys and values (each requiring the same memory)
$L$ : the number of transformer layers, each maintaining its own KV cache
$h$ : the number of attention heads per layer, each with independent K and V tensors
$T$ : the current sequence length (this grows with each generated token)
$d_k$ : the dimension per head, determining the size of each key/value vector
Bytes per element: depends on numerical precision (4 for FP32, 2 for FP16/BF16)

The formula multiplies these factors because we need to store one key vector and one value vector for each layer and head at every sequence position. The memory grows linearly with sequence length, which becomes problematic for long-context generation.

Let's calculate the KV cache requirements for a concrete example.

In[3]:

Code

# KV cache size calculation for different model scales
def calculate_kv_cache_size(
    num_layers: int,
    num_heads: int,
    head_dim: int,
    seq_length: int,
    batch_size: int = 1,
    bytes_per_element: int = 2,  # FP16/BF16
) -> dict:
    """Calculate KV cache memory requirements in bytes and GB."""
    # Standard MHA: each head has its own K and V
    kv_cache_bytes = (
        2  # K and V
        * num_layers
        * num_heads
        * seq_length
        * head_dim
        * batch_size
        * bytes_per_element
    )

    return {
        "bytes": kv_cache_bytes,
        "megabytes": kv_cache_bytes / (1024**2),
        "gigabytes": kv_cache_bytes / (1024**3),
    }


# Example: LLaMA-2 7B scale model
llama_7b_config = {
    "num_layers": 32,
    "num_heads": 32,
    "head_dim": 128,
    "seq_length": 4096,
}

kv_cache = calculate_kv_cache_size(**llama_7b_config)

Out[4]:

Console

LLaMA-2 7B KV Cache (4K context, batch=1):
  Size: 2.00 GB
  Per-token overhead: 512.00 KB
  At 8,192 tokens: 4.00 GB
  At 16,384 tokens: 8.00 GB
  At 32,768 tokens: 16.00 GB
  At 131,072 tokens: 64.00 GB

For a 7B parameter model, the KV cache alone consumes over 1 GB for a 4K context window. At 128K tokens (common in modern models), this balloons to over 32 GB, which exceeds the memory of most consumer GPUs. The problem compounds further with larger batch sizes for throughput optimization.

The situation becomes even more acute for larger models. Let's compare across scales.

In[5]:

Code

# Model configurations showing query heads (for KV cache comparison assuming MHA)
# Note: LLaMA-2 70B actually uses GQA with 8 KV heads, not full MHA
models = {
    "LLaMA-2 7B": {"layers": 32, "heads": 32, "head_dim": 128},
    "LLaMA-2 13B": {"layers": 40, "heads": 40, "head_dim": 128},
    "LLaMA-2 70B": {"layers": 80, "heads": 64, "head_dim": 128},
    "GPT-3 175B": {"layers": 96, "heads": 96, "head_dim": 128},
}

seq_lengths = [4096, 32768, 131072]

Out[6]:

Visualization

Bar chart showing KV cache size in GB for different model scales at varying context lengths. — KV cache memory requirements grow rapidly with model size and context length. A 70B model serving 128K context requires over 160 GB of memory just for the cache, exceeding the capacity of even enterprise GPUs like the A100-80GB.

This memory pressure is the driving motivation behind MQA. When the cache alone exceeds available GPU memory, you can't run inference at all, regardless of how fast your hardware is.