Grouped-Query KV Storage

Problem

Return the number of bytes needed to store all keys and values for a dense grouped-query KV cache. Include both K and V, every layer, every sequence, and every cached token. Count stored KV heads, not query heads. This calculation excludes allocator padding and metadata.

def kv_storage_bytes(batch, tokens, layers, kv_heads, head_dim, bytes_per_element):

Constraints

• batch,tokens,layers,kv_heads,head_dim are nonnegative integers at most 1024.
• bytes_per_element is 1,2,4,or 8; the exact result is at most 2^53−1.
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  1,
  1,
  1,
  1,
  1,
  2
]
Expected result: 4

Arguments: [
  2,
  4,
  3,
  2,
  8,
  2
]
Expected result: 1536