Minimal implementation

from dataclasses import dataclass

GiB = 1024 ** 3

@dataclass
class Config:
    layers: int
    kv_heads: int
    head_dim: int

def kv_cache_bytes(cfg, tokens, users, bytes_per_elem=2):
    # K and V, for every layer, KV head, token and user
    return (2 * cfg.layers * cfg.kv_heads * cfg.head_dim
            * bytes_per_elem * tokens * users)

llama_70b = Config(layers=80, kv_heads=8, head_dim=128)
print(kv_cache_bytes(llama_70b, 131072, 100) / GiB)  # 4000.0

In production

Papers

• GQA: Training Generalized Multi-Query Transformer Models from Multi-Head CheckpointsAinslie et al., 2023
• Efficient Memory Management for Large Language Model Serving with PagedAttentionKwon et al., SOSP 2023
• DeepSeek-V2, which introduces multi-head latent attention (MLA)DeepSeek-AI, 2024

Interview drill (optional)