Paged Attention Block Table

Problem

Compute the paged-attention metadata for a sequence: the number of fixed-size KV blocks and how many slots of the last block are used. Blocks are filled completely before a new one is allocated.

The signature is paged_attention_block_table(seq_len, block_size). Return [num_blocks, last_block_used]; an empty sequence uses no blocks.

Hint

One token past a block boundary allocates a fresh block.

Test cases

case 1

in: [100,16]

expected: [7,4]

case 2

in: [16,16]

expected: [1,16]

case 3

in: [17,16]

expected: [2,1]

case 4

in: [0,16]

expected: [0,0]

Starter codePython

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

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