KV Cache Dtype Savings

Problem

Compute the byte savings from storing the KV cache in a smaller dtype. Given the fp16 total, scale it by new_bytes_per_value / old_bytes_per_value, then report the new total, the bytes saved, and the percentage saved.

The signature is kv_cache_dtype_savings(fp16_bytes, new_bytes_per_value=1, old_bytes_per_value=2). Return [new_total, saved_bytes, pct_saved].

Hint

The saving depends only on the ratio of dtype widths, not on the model size.

Test cases

case 1

in: [1000000,1,2]

expected: [500000,500000,50]

case 2

in: [134217728,1,2]

expected: [67108864,67108864,50]

case 3

in: [1000,2,2]

expected: [1000,0,0]

case 4

in: [1000,1,4]

expected: [250,750,75]

Starter codePython

def kv_cache_dtype_savings(fp16_bytes, new_bytes_per_value=1, old_bytes_per_value=2):
    # Your code here
    pass

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