FlashAttention: tiled online-softmax, exact (hmfa)

Usage

morie_geron_flash_attention(Q, K, V, block_size = 2, causal = FALSE)

Arguments

Q, K, V: Matrices.
block_size: Tile side (>=1).
causal: Mask future keys.

Value

List with output, direct_output, max_abs_error, n_blocks, peak_score_memory.

Examples

morie_geron_flash_attention(Q = 0.5, K = c(1, 2, 3, 4, 5, 6, 7, 8), V = c(1, 2, 3, 4, 5, 6, 7, 8))
#> $output
#>          [,1]
#> [1,] 6.607765
#> 
#> $direct_output
#>          [,1]
#> [1,] 6.607765
#> 
#> $max_abs_error
#> [1] 8.881784e-16
#> 
#> $row_max
#> [1] 4
#> 
#> $row_sum
#> [1] 2.494945
#> 
#> $n_blocks
#> [1] 4
#> 
#> $block_size
#> [1] 2
#> 
#> $peak_score_memory
#> [1] 2
#> 
#> $naive_score_memory
#> [1] 8
#> 
#> $memory_ratio
#> [1] 4
#> 
#> $causal
#> [1] FALSE
#> 
#> $estimate
#> [1] 8.881784e-16
#> 
#> $n
#> [1] 1
#> 
#> $method
#> [1] "tiled online-softmax attention (FlashAttention), exact"
#>