Flash Decode Block Split

Problem

Split the KV cache into blocks for flash decoding, which processes each block in parallel and merges partial results. Return [blocks, num_partials], where each block is a [start, end) pair of at most block_size positions.

The signature is flash_decode_split(seq_len, block_size).

Hint

Splitting the KV dimension gives long-context decoding more thread blocks to parallelize.

Test cases

case 1

in: [100,32]

expected: [[[0,32],[32,64],[64,96],[96,100]],4]

case 2

in: [10,10]

expected: [[[0,10]],1]

case 3

in: [10,4]

expected: [[[0,4],[4,8],[8,10]],3]

case 4

in: [0,4]

expected: [[],0]

Starter codePython

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

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