Quantized KV Cache Size

Problem

Compute quantized KV cache size in bytes: 2 (keys and values) * num_layers * num_kv_heads * seq_len * head_dim * bits/8 * batch. This is the memory each active sequence holds during decoding.

The signature is quantized_kv_cache(num_layers, num_kv_heads, seq_len, head_dim, bits=8, batch=1).

Hint

KV cache bits scale linearly with context length, so int8 halves long-context memory.

Test cases

case 1

in: [32,8,1024,128,8,1]

expected: 67108864

case 2

in: [1,1,4,2,4,1]

expected: 8

case 3

in: [2,4,8,64,16,3]

expected: 49152

case 4

in: [0,4,8,64,8,1]

expected: 0

Starter codePython

def quantized_kv_cache(num_layers, num_kv_heads, seq_len, head_dim, bits=8, batch=1):
    # Your code here
    pass

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