Attention FLOPs (Quadratic)

Problem

Count the FLOPs of the two quadratic attention matrix products (scores QK^T and the weighted sum of V), treating a multiply-add as 2 FLOPs: 4 * batch_size * seq_len^2 * d_model.

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

Hint

Both score computation and weighted sum scale with seq_len squared.

Test cases

case 1

in: [1,128,512]

expected: 33554432

case 2

in: [8,512,768]

expected: 6442450944

case 3

in: [1,1,64]

expected: 256

case 4

in: [2,32,128]

expected: 1048576

Starter codePython

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

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