Fused Attention SRAM Tiles

Problem

Size the query tile of a fused attention kernel. The SRAM must hold one query block and two key/value blocks of head_dim elements each, so block_rows = max(1, sram_bytes // (2 * elem_bytes * head_dim)) and num_blocks = ceil(seq_len / block_rows).

The signature is fused_attention_tiles(sram_bytes, elem_bytes, head_dim, seq_len). Return [block_rows, num_blocks].

Hint

Flash attention streams K/V tiles through SRAM instead of materializing the score matrix in HBM.

Test cases

case 1

in: [16384,2,64,1024]

expected: [64,16]

case 2

in: [1024,2,128,512]

expected: [2,256]

case 3

in: [100000,2,64,1000]

expected: [390,3]

case 4

in: [64,4,8,10]

expected: [1,10]

Starter codePython

def fused_attention_tiles(sram_bytes, elem_bytes, head_dim, seq_len):
    # Your code here
    pass

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