KV Cache Memory Footprint Across Attention Variants

The problem, from first principles

[02-autoregressive-decoding-kv-cache] built a cache that grows one row per generated token. Every real serving deployment needs to know, ahead of time, how much GPU memory that cache will actually consume — this is the calculation every serving-capacity plan starts from. Compute the KV cache memory footprint (in bytes) for a batch of sequences, given model shape parameters and which attention variant (MHA, MQA, or GQA) is in use.

From theory to code

kv_heads = n_heads (MHA)  |  1 (MQA)  |  n_kv_heads (GQA)
bytes_per_token = 2 * n_layers * kv_heads * d_head * bytes_per_element
total_bytes = batch_size * seq_len * bytes_per_token

Constraints

• variant is one of 'mha', 'mqa', 'gqa'; for 'mha' effective kv_heads = n_heads, for 'mqa' it's 1, for 'gqa' it's the given n_kv_heads.
• Account for K and V separately (factor of 2).
• Multiply across n_layers, seq_len, and batch_size.
• Return total bytes, and also bytes-per-token-per-layer for reference.

Hints