Tiled Attention Row Output

Problem

Compute one attention output row tile by tile with the online softmax recurrence, never materializing the full score vector. Scores use the standard 1 / sqrt(d) scale, and each block of block_size keys rescales the running accumulator when a new maximum appears.

The signature is tiled_attention_row(q, K, V, block_size). Return the output vector for the query row q; the result must match a full softmax weighted sum.

Hint

When a new maximum arrives, multiply the accumulator and denominator by exp(old_m - new_m) before adding the tile contribution.

Test cases

case 1

in: [[1,0],[[1,0],[0,1]],[[1,0],[0,1]],1]

expected: [0.6697615493,0.3302384507]

case 2

in: [[1,0],[[1,0],[0,1],[1,1]],[[1,0],[0,1],[1,1]],2]

expected: [0.8022241854,0.5988879073]

case 3

in: [[0,1],[[1,0],[0,1]],[[2,3],[4,5]],5]

expected: [3.3395230987,4.3395230987]

case 4

in: [[1,1],[[1,0],[0,2]],[[1,1],[2,0]],1]

expected: [1.6697615493,0.3302384507]

Starter codePython

import math
def tiled_attention_row(q, K, V, block_size):
    # Your code here
    pass

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