KV Cache Size

Problem

Compute the byte size of the key-value cache: 2 (K and V) * num_layers * batch_size * num_kv_heads * seq_len * head_dim * bytes_per_value.

The signature is kv_cache_bytes(num_layers, num_kv_heads, head_dim, seq_len, batch_size=1, bytes_per_value=2). Return an integer.

Hint

The cache stores both keys and values for every layer and KV head.

case 1

in: [12,12,64,512]

expected: 18874368

case 2

in: [32,8,128,2048,1,2]

expected: 268435456

case 3

in: [1,1,4,10,1,4]

expected: 320

case 4

in: [6,6,32,100,2,2]

expected: 921600

Starter codePython

def kv_cache_bytes(num_layers, num_kv_heads, head_dim, seq_len, batch_size=1, bytes_per_value=2):
    # Your code here
    pass

4 test cases · run your code against them in the editor