Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Grouped Query Attention › Memory Savings Analysis

Michael BrenndoerferUpdated August 6, 20254 min read

Language AI HandbookAll chapters

Part 29 of 67Modern Decoder Models

Chapter 209Grouped Query Attention › Memory Savings Analysis

Reading section 5 of 14 · Memory Savings Analysis · 4 minSection 5 of 14 · 4 min

Section 5 of 14 · 4 min read

Grouped Query Attention › Memory Savings AnalysisOne section · Sections in this chapter5 / 14 ↓

Sections in this chapter14

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

GQA's memory benefits depend on the ratio of query heads to key-value heads. During autoregressive generation, the model caches key and value tensors from previous positions to avoid redundant computation. The size of this KV cache determines how many sequences can be processed in parallel and how long those sequences can be.

KV Cache Memory FormulaLink Copied

For a single transformer layer using GQA, the KV cache stores keys and values for all $g$ groups across all cached positions:

\text{Cache}_{\text{GQA}} = 2 \times g \times d_k \times T \times B \times \text{bytes}

where:

$2$ : accounts for storing both keys and values (two tensors per group)
$g$ : number of key-value groups in GQA
$d_k$ : dimension per head (typically 64 or 128 in modern models)
$T$ : current sequence length (grows during generation)
$B$ : batch size (number of sequences processed in parallel)
$\text{bytes}$ : bytes per element (2 for float16/bfloat16, 4 for float32)

For a full model with $L$ layers, multiply by $L$ . The total cache size is $2 \times L \times g \times d_k \times T \times B \times \text{bytes}$ .

Memory Reduction FactorLink Copied

Comparing GQA to MHA, which stores $h$ key-value heads instead of $g$ , the memory reduction is simply the ratio of heads to groups:

\text{Reduction factor} = \frac{h}{g}

where:

$h$ : number of query heads (and number of KV heads in MHA)
$g$ : number of KV groups in GQA

For GQA-8 with 64 query heads, the reduction factor is $64/8 = 8\times$ . This means the KV cache is 8 times smaller than it would be with full MHA, letting either 8 times longer sequences, 8 times larger batches, or some combination of both.

In[11]:

Code

def memory_analysis(
    num_query_heads: int,
    head_dim: int,
    num_layers: int,
    seq_lengths: list,
    batch_size: int = 1,
):
    """Analyze KV cache memory for different GQA configurations."""
    bytes_per_element = 2  # fp16

    kv_head_options = [num_query_heads, 8, 4, 2, 1]

    results = {"seq_lengths": seq_lengths, "configs": []}

    for num_kv_heads in kv_head_options:
        config_data = {
            "name": "MHA"
            if num_kv_heads == num_query_heads
            else "MQA"
            if num_kv_heads == 1
            else f"GQA-{num_kv_heads}",
            "kv_heads": num_kv_heads,
            "memory_gb": [],
        }

        for seq_len in seq_lengths:
            cache_bytes = (
                2
                * num_layers
                * num_kv_heads
                * head_dim
                * seq_len
                * batch_size
                * bytes_per_element
            )
            config_data["memory_gb"].append(cache_bytes / 1e9)

        results["configs"].append(config_data)

    return results


# Analyze for a 70B-scale model
memory_data = memory_analysis(
    num_query_heads=64,
    head_dim=128,
    num_layers=80,
    seq_lengths=[1024, 2048, 4096, 8192, 16384],
    batch_size=1,
)

Out[12]:

Visualization

Line plot showing KV cache memory growth with sequence length for MHA, GQA variants, and MQA. — KV cache memory scaling with sequence length for different attention variants. GQA-8 (used by LLaMA 2 70B) provides 8x memory reduction compared to MHA while retaining more capacity than MQA. At 16K tokens, this translates to 4.4 GB vs 35.2 GB for MHA.

Out[13]:

Console

KV Cache Memory at 16K Sequence Length (70B model, batch=1)
=======================================================

Variant      KV Heads     Memory (GB)     vs MHA
-------------------------------------------------------
MHA          64           42.95             1.0x smaller
GQA-8        8            5.37             8.0x smaller
GQA-4        4            2.68             16.0x smaller
GQA-2        2            1.34             32.0x smaller
MQA          1            0.67             64.0x smaller

The analysis reveals why GQA-8 has become the standard choice for large models. At 16K tokens, MHA would require 35 GB just for the KV cache, exceeding the memory of most GPUs. GQA-8 reduces this to 4.4 GB, fitting comfortably on modern hardware while retaining 8 independent key-value representations.