Paged KV Cache Footprint

Problem

Compute allocated KV cache bytes with paged memory: each sequence is rounded up to a whole number of blocks, so padded_tokens = ceil(seq_len / block_size) * block_size.

The signature is paged_kv_cache_bytes(num_layers, batch_size, num_kv_heads, head_dim, seq_len, block_size, bytes_per_value=2). Total bytes = 2 * num_layers * batch_size * num_kv_heads * head_dim * padded_tokens * bytes_per_value; return an integer.

Hint

Only the token count gets rounded up; the block size itself never changes the per-token cost.

Test cases

case 1

in: [12,1,12,64,500,16,2]

expected: 18874368

case 2

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

expected: 768

case 3

in: [2,1,2,8,8,8,2]

expected: 1024

case 4

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

expected: 134217728

Starter codePython

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

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