Calculate KV-Cache Memory Footprinteasy

Calculate KV-Cache Memory Footprint

easy Inside an LLM 15 min

NumPy allowed · plain Python earns 10 bonus points

Tokens attend to tokens

Q · Kᵀ / √d → softmax → weighted values, computed in the page.

Open lesson →

01tokens

02Q K V

03scores

04attention

modelslearnrelationships

models32%

learn38%

relationships30%

Q · Kᵀ / √d scores normalized with softmax

Calculate the exact VRAM footprint in bytes required by an LLM KV-cache during multi-turn autoregressive serving.

Task

Implement calculate_kv_cache_bytes(num_layers, num_kv_heads, head_dim, seq_len, batch_size, bytes_per_param=2). Return the total byte count accounting for both Key and Value tensors across all layers, heads, positions, and batch slots.

Function Signature

calculate_kv_cache_bytes(num_layers: int, num_kv_heads: int, head_dim: int, seq_len: int, batch_size: int, bytes_per_param: int = 2) -> int

Examples

Example 1: Single Layer Minimal Config
Input: batch_size = 1
bytes_per_param = 2
head_dim = 64
num_kv_heads = 2
num_layers = 1
seq_len = 128
Output: 65536

Example 2: Llama-2 7B Configuration
Input: batch_size = 4
bytes_per_param = 2
head_dim = 128
num_kv_heads = 32
num_layers = 32
seq_len = 2048
Output: 4294967296

Constraints

All arguments must be positive integers (`>` 0). Raise ValueError otherwise.
The factor of 2 accounts for storing both Key and Value matrices.
bytes_per_param defaults to 2 (FP16 / BF16 precision).

Python3Saved ✓

def calculate_kv_cache_bytes(num_layers, num_kv_heads, head_dim, seq_len, batch_size, bytes_per_param=2):

"""Calculate total memory in bytes required to store Key-Value cache:

2 * num_layers * num_kv_heads * head_dim * seq_len * batch_size * bytes_per_param.

Raise ValueError if any parameter is <= 0."""

# Your implementation here

pass

Case 1: Single Layer Minimal Config

Input: batch_size = 1
bytes_per_param = 2
head_dim = 64
num_kv_heads = 2
num_layers = 1
seq_len = 128

Expected: 65536

Case 2: Llama-2 7B Configuration

Input: batch_size = 4
bytes_per_param = 2
head_dim = 128
num_kv_heads = 32
num_layers = 32
seq_len = 2048

Expected: 4294967296