Context Memory Estimate

Problem

Estimate the memory of the attention probability tensor for a full context window: batch_size * num_heads * seq_len^2 * bytes_per_value.

The signature is attention_memory_bytes(batch_size, num_heads, seq_len, bytes_per_value=2). Return an integer.

Hint

The seq_len^2 attention matrix is why long-context inference scales quadratically.

case 1

in: [1,12,512]

expected: 6291456

case 2

in: [8,16,1024,2]

expected: 268435456

case 3

in: [1,8,2048,4]

expected: 134217728

case 4

in: [2,4,128,2]

expected: 262144

Starter codePython

def attention_memory_bytes(batch_size, num_heads, seq_len, bytes_per_value=2):
    # Your code here
    pass

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