Attention Memory Estimate

Problem

Estimate the activation memory of a self-attention layer.

The QKV projections hold 3 * batch * tokens * dim values and the attention matrix holds batch * tokens * tokens values. Return the total number of values multiplied by bytes_per_value.

Hint

QKV plus the tokens x tokens attention matrix, all in values.

Test cases

case 1

in: [2,16,64,4]

expected: 26624

case 2

in: [1,4,8,4]

expected: 448

case 3

in: [1,1,1,2]

expected: 8

Starter codePython

def attention_memory(batch, tokens, dim, bytes_per_value):
    # Your code here
    pass

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