Compare KV Cache Layouts: MHA, GQA, MQA and MLA

The KV cache usually limits how many tokens a server can hold in memory. Attention variants exist mainly to shrink it:

• MHA: every head stores its own K and V: 2 × layers × heads × head_dim values per token
• GQA: heads share K/V in groups: 2 × layers × kv_heads × head_dim
• MQA: GQA with a single KV head
• MLA (DeepSeek-V2/V3): stores one compressed latent of size latent_dim plus a small decoupled RoPE key of size rope_dim per layer, with no factor of 2: layers × (latent_dim + rope_dim)

Implement kv_bytes_per_token(layout, layers, head_dim=0, heads=0, kv_heads=0, latent_dim=0, rope_dim=0, bytes_per_value=2) for layout in "mha", "gqa", "mqa" and "mla", and max_tokens(memory_gib, bytes_per_token): how many tokens fit in that many GiB (2³⁰ bytes), rounded down.