Attention Recompute Memory Savings

Problem

Compare the memory of materializing the full attention matrix against recomputing it block by block. Full storage holds batch_size * num_heads * seq_len * seq_len values, while the tiled version holds batch_size * num_heads * seq_len * block_size values.

The signature is attention_recompute_savings(batch_size, num_heads, seq_len, block_size, bytes_per_value=2). Return [full_bytes, recompute_bytes, saved_fraction].

Hint

The saved fraction depends only on block_size / seq_len, not on batch size, heads, or dtype.

Test cases

case 1

in: [1,12,1024,128,2]

expected: [25165824,3145728,0.875]

case 2

in: [8,16,512,64,2]

expected: [67108864,8388608,0.875]

case 3

in: [1,1,10,10,4]

expected: [400,400,0]

case 4

in: [2,4,8,2,2]

expected: [1024,256,0.75]

Starter codePython

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

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