KV Cache Bytes Per Token

Problem

Compute the KV cache bytes added per generated token: 2 (for K and V) * num_layers * num_kv_heads * head_dim * bytes_per_value.

The signature is kv_bytes_per_token(num_layers, num_kv_heads, head_dim, bytes_per_value=2). Return an integer.

Hint

Two tensors per layer, one write per token, multiplied by the dtype width.

case 1

in: [12,12,64,2]

expected: 36864

case 2

in: [32,8,128,2]

expected: 131072

case 3

in: [1,1,4,4]

expected: 32

case 4

in: [12,12,64,1]

expected: 18432

Starter codePython

def kv_bytes_per_token(num_layers, num_kv_heads, head_dim, bytes_per_value=2):
    # Your code here
    pass

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