Transformer Attention Total FLOPs

Problem

Count total attention FLOPs per layer: the quadratic score and weighted-sum products cost 4 * batch_size * seq_len^2 * d_model, and the QKV plus output projections cost 8 * batch_size * seq_len * d_model^2, treating a multiply-add as 2 FLOPs.

The signature is attention_total_flops(batch_size, seq_len, d_model). Return an integer.

Hint

Short sequences are dominated by projections; long sequences are dominated by the seq_len^2 term.

Test cases

case 1

in: [1,128,512]

expected: 301989888

case 2

in: [8,512,768]

expected: 25769803776

case 3

in: [1,1,64]

expected: 33024

case 4

in: [2,32,128]

expected: 9437184

Starter codePython

def attention_total_flops(batch_size, seq_len, d_model):
    # Your code here
    pass

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