Flash Attention Block Count

Problem

Count the tile pairs processed by tiled flash attention: ceil(seq_len / block_size) query blocks times the same number of key blocks.

The signature is flash_attention_blocks(seq_len, block_size). Return an integer.

Hint

Flash attention never materializes the full seq_len x seq_len matrix; it works tile by tile.

case 1

in: [1024,128]

expected: 64

case 2

in: [1000,128]

expected: 64

case 3

in: [128,128]

expected: 1

case 4

in: [100,32]

expected: 16

Starter codePython

def flash_attention_blocks(seq_len, block_size):
    # Your code here
    pass

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