"""
kv_sizes.py: how big is the notebook (KV cache) for real models?
notebook bytes per token = 2 (K and V) x floors x notebook heads x head size x bytes per number
Configs are from each model's published config.json.
"""
def per_token(layers, kv_heads, head_dim, nbytes, k_and_v=2):
    return k_and_v * layers * kv_heads * head_dim * nbytes

MODELS = [
    # name,                                   bytes per token
    ("GPT-2 small shape, float32 (our toy)",    per_token(12, 12, 64, 4)),
    ("Llama 3.1 8B, if every head kept notes",  per_token(32, 32, 128, 2)),
    ("Llama 3.1 8B, BF16 (8 shared notebooks)", per_token(32, 8, 128, 2)),
    ("Llama 3.1 8B, FP8 notebook",              per_token(32, 8, 128, 1)),
    ("Llama 3.1 70B, BF16",                     per_token(80, 8, 128, 2)),
    ("DeepSeek-V3, if every head kept notes",   61 * 128 * (192 + 128) * 2),
    ("DeepSeek-V3, BF16 (compressed notes)",    61 * (512 + 64) * 2),
]
print(f"{'model':<42} {'per token':>10} {'8K-token chat':>14}")
for name, b in MODELS:
    print(f"{name:<42} {b/1024:8,.0f} KB {b*8192/1e9:11.2f} GB")

print("\nhow many users fit on one 80 GB H100 with Llama 3.1 8B in BF16?")
weights = 8.03e9 * 2                      # 16 GB of tables
room = 80e9 - weights - 6e9               # keep ~6 GB for everything else
for name, b in [("BF16 notebook", per_token(32, 8, 128, 2)), ("FP8 notebook", per_token(32, 8, 128, 1))]:
    tokens = room / b
    print(f"  {name}: {room/1e9:.0f} GB free -> {tokens/1e3:,.0f}K tokens of notes"
          f" -> {tokens/8192:5.0f} users at 8K, {tokens/131072:4.1f} users at 128K")

print("\none decode step, Llama 3.1 8B, 48 users at 8K each:")
kv = 48 * 8192 * per_token(32, 8, 128, 2)
print(f"  tables {weights/1e9:.0f} GB + notebooks {kv/1e9:.0f} GB -> the notebooks are {kv/weights:.1f}x the model")