Calculating KV Cache Memory Footprint

The KV Cache Memory Footprint

The Memory Bottleneck

While KV caching saves computation by avoiding redundant attention math, it introduces a significant memory bottleneck. Unlike model weights, which are static, the KV cache grows linearly with every token generated, eventually competing with the model itself for VRAM.

Calculating the KV cache footprint is essential for any ML engineer. As we generate text, we store Key and Value vectors for every token across every layer. This creates a memory requirement that grows linearly with the sequence length, eventually becoming the primary constraint on our system's capacity.

• KV cache grows linearly with sequence length
• Memory capacity is the primary constraint for long-context LLMs
• Accurate calculation is vital for capacity planning

Deconstructing the Equation

Variables of the Cache

To calculate the size, we need to understand five key architectural parameters:

• L: Number of layers
• H_kv: Number of KV heads (Note: GQA uses fewer than query heads)
• D_head: Dimension of each head
• S: Sequence length
• P: Precision in bytes

Let's break down the components. The '2' accounts for both Keys and Values. 'L' is the depth of the model. Crucially, 'H_kv' refers to the number of KV heads—in architectures like Llama-3, this is much smaller than the number of query heads due to Grouped-Query Attention. Finally, we multiply by the head dimension, the sequence length, and the precision in bytes.

• 2x multiplier accounts for both Keys and Values
• H_kv is usually smaller than Query heads in modern models (GQA)
• Precision (P) is typically 2 bytes for BF16/FP16

Interactive Scaling: Llama-3 8B

Llama-3 8B Case Study

Adjust the Context Length (S) slider to see how the KV cache grows for a Llama-3 8B model (32 layers, 8 KV heads, 128 head dim, BF16 precision).

Use the slider to explore how context length impacts the memory footprint of a Llama-3 8B model. Notice the horizontal line representing the model weights at 15 Gigabytes. As you move toward 128,000 tokens, watch how the cache eventually becomes larger than the model itself, demanding massive VRAM capacity.

• At 128K context, the cache size (16GB) exceeds the model weights (15GB)
• Scaling sequence length shifts the bottleneck from compute to memory

Exercise: Llama-3 70B Footprint

Calculate the 70B Cache

Apply the formula to a Llama-3 70B model for an 8K context window. Type your final answer in Gigabytes (GB).

• Layers: 80
• KV Heads: 8
• Head Dim: 128
• Precision: BF16 (2 bytes)
• Context: 8,192

Now it's your turn. Using the architectural specs provided, calculate the KV cache size for a Llama-3 70B model with an 8,000 token context. Enter your result in Gigabytes to check your accuracy.

• Memory = 2 * L * H_kv * D_head * S * P
• Correct calculation results in ~2.5 GB

Scaling for Real-World Serving

Batch Size & Fragmentation

In production, you aren't serving one user. The total VRAM required is Memory per request × Batch Size. Additionally, standard allocators waste space through fragmentation.

Let's look at the bigger picture. While model weights are shared across all users, every concurrent request needs its own KV cache. Increase the batch size to see how quickly VRAM vanishes. Also, notice the 'waste' layer—this is memory fragmentation, which tools like PagedAttention aim to eliminate.

• Total VRAM = (KV Cache + Weights) * Batch Size is WRONG; Weights are shared, KV Cache is not.
• Fragmentation can add 20-30% overhead.
• PagedAttention is the industry standard to solve fragmentation.