Prefill Decode FLOP Ratio

Problem

Estimate the prefill-to-decode FLOP ratio for a transformer. Prefill of a prompt with seq_len tokens costs about 2 * model_params * seq_len FLOPs and one decode step costs about 2 * model_params, so the ratio reduces to seq_len.

The signature is prefill_decode_flops_ratio(seq_len, model_params). Return a float.

Hint

Model size cancels out, so one prefill equals seq_len decode steps of compute.

Test cases

case 1

in: [128,1000000000]

expected: 128

case 2

in: [1,7000000000]

expected: 1

case 3

in: [2048,7000000000]

expected: 2048

case 4

in: [512,100000000]

expected: 512

Starter codePython

def prefill_decode_flops_ratio(seq_len, model_params):
    # Your code here
    pass

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