Skip to contents

FlashAttention tiled exact attention

Source: R/flsh2_native.R
Flsh2.Rd

Block-tiled streaming attention with online softmax: the running row maximum, row sum and accumulator are rescaled by exp(m_old - m_new) as key/value blocks stream through, so the output equals softmax of Q K transpose over sqrt(d) times V exactly while only one tile of scores is ever materialised. There is exactly one implementation: this function delegates to morie_geron_flash_attention.

Usage

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

Arguments

Q

Query matrix, N by d.

K

Key matrix, M by d.

V

Value matrix, M by dv.

block_size

Tile side, positive integer. Default 2.

causal

Mask keys after the query position. Default FALSE.

Value

List with output, direct_output, max_abs_error, row_max, row_sum, n_blocks, peak_score_memory, naive_score_memory, memory_ratio, estimate, n, method.

References

Dao, T., Fu, D. Y., Ermon, S., Rudra, A. and Re, C. (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 35, arXiv:2205.14135, Algorithm 1.

Examples

Flsh2(matrix(0, 1, 1), matrix(c(1, 3), 2, 1), matrix(c(1, 3), 2, 1))$output
#>      [,1]
#> [1,]    2