Layer 1 cache: K,V for all past tokens
Layer 2 cache: K,V for all past tokens
... one cache per layer ...
python
def kv_bytes(layers, kv_heads, head_dim, seq, batch, dtype_bytes=2):
# 2x for both K and V
return 2 * layers * kv_heads * head_dim * seq * batch * dtype_bytes
mha = kv_bytes(layers=32, kv_heads=32, head_dim=128, seq=4096, batch=1)
gqa = kv_bytes(layers=32, kv_heads=8, head_dim=128, seq=4096, batch=1)
print(f"MHA: {mha/1e9:.2f} GB/seq GQA: {gqa/1e9:.2f} GB/seq")
weights_gb = 14
free_gb = 80 - weights_gb
print("concurrent 4k sequences (GQA):", int(free_gb*1e9 / gqa))