Inference FLOPs Estimate

Problem

Estimate the multiply-add FLOPs of one transformer forward pass. Per layer: QKV projections 6 * seq * d^2, output projection 2 * seq * d^2, attention scores and weighted sum 4 * seq^2 * d, and the two MLP matrices 4 * seq * d * d_ff. Sum these and multiply by num_layers.

The signature is inference_flops(seq_len, d_model, num_layers, d_ff). Return an integer.

Hint

One multiply-add counts as 2 FLOPs; the attention term grows quadratically with sequence length.

Test cases

case 1

in: [128,512,12,2048]

expected: 10066329600

case 2

in: [1,512,12,2048]

expected: 75522048

case 3

in: [512,768,12,3072]

expected: 96636764160

case 4

in: [32,64,2,256]

expected: 6815744

Starter codePython

def inference_flops(seq_len, d_model, num_layers, d_ff):
    # Your code here
    pass

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