FlashAttention tiled online softmax (Geron Ch 16, morie.fn grflash)

Streams K/V in blocks of block_size, rescaling the running softmax denominator; exact against the plain attention kernel at every block size.

Usage

morie_geron_flash_attention_tile(Q, K, V, block_size = 2)

Arguments

Q, K, V

Attention inputs.

block_size

Positive tile height over the key axis.

Value

List with output, reference_output, max_abs_error, n_blocks, row_max, row_denominator, peak_score_elements, full_score_elements.

Examples

morie_geron_flash_attention_tile(Q = 0.5, K = 5L, V = c(1, 2, 3, 4, 5, 6, 7, 8))
#> $output
#>      [,1] [,2] [,3] [,4] [,5] [,6] [,7] [,8]
#> [1,]    1    2    3    4    5    6    7    8
#> 
#> $reference_output
#>      [,1] [,2] [,3] [,4] [,5] [,6] [,7] [,8]
#> [1,]    1    2    3    4    5    6    7    8
#> 
#> $max_abs_error
#> [1] 0
#> 
#> $n_blocks
#> [1] 1
#> 
#> $row_max
#> [1] 2.5
#> 
#> $row_denominator
#> [1] 1
#> 
#> $peak_score_elements
#> [1] 1
#> 
#> $full_score_elements
#> [1] 1
#> 
#> $block_size
#> [1] 2
#> 
#> $estimate
#>      [,1] [,2] [,3] [,4] [,5] [,6] [,7] [,8]
#> [1,]    1    2    3    4    5    6    7    8
#> 
#> $n
#> [1] 1
#>