KV Cache Size With Grouped Query Attentionmedium

KV Cache Size With Grouped Query Attention

medium 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

Implement KV-cache size estimation specifically accounting for Grouped Query Attention's reduced KV head count, the real reason GQA models have a measurably smaller serving memory footprint.

Task

Implement gqa_kv_cache_size_mb(num_layers, num_kv_heads, head_dim, seq_len, batch_size, bytes_per_value). Total elements: 2 (K and V) * num_layers * num_kv_heads * head_dim * seq_len * batch_size. Multiply by bytes_per_value, convert to MB (divide by 1024*1024).

Function Signature

gqa_kv_cache_size_mb(num_layers: int, num_kv_heads: int, head_dim: int, seq_len: int, batch_size: int, bytes_per_value: float) -> float

Examples

Example 1: few kv heads small cache
Input: num_layers = 32
num_kv_heads = 8
head_dim = 128
seq_len = 2048
batch_size = 1
bytes_per_value = 2
Output: 256

Example 2: fewer kv heads even smaller
Input: num_layers = 32
num_kv_heads = 2
head_dim = 128
seq_len = 2048
batch_size = 1
bytes_per_value = 2
Output: 64

Example 3: longer context scales linearly
Input: num_layers = 32
num_kv_heads = 8
head_dim = 128
seq_len = 8192
batch_size = 1
bytes_per_value = 2
Output: 1024

Constraints

all inputs positive

Python3Saved ✓

def gqa_kv_cache_size_mb(num_layers, num_kv_heads, head_dim, seq_len, batch_size, bytes_per_value):

# Your implementation here

pass

Case 1: few kv heads small cache

Input: num_layers = 32
num_kv_heads = 8
head_dim = 128
seq_len = 2048
batch_size = 1
bytes_per_value = 2

Expected: 256

Case 2: fewer kv heads even smaller

Input: num_layers = 32
num_kv_heads = 2
head_dim = 128
seq_len = 2048
batch_size = 1
bytes_per_value = 2

Expected: 64

Case 3: longer context scales linearly

Input: num_layers = 32
num_kv_heads = 8
head_dim = 128
seq_len = 8192
batch_size = 1
bytes_per_value = 2

Expected: 1024