Dao et al. (2022), "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness," speed up the standard self-attention computation on GPUs without changing the attention mechanism's mathematical result. What kind of optimization does FlashAttention make, and what does it NOT do?
- It reduces the number of reads and writes between the GPU's high-bandwidth memory and its much faster on-chip SRAM by tiling the computation and avoiding materializing the full attention-score matrix, while still computing exactly the same attention output as the standard formula, not an approximation of it
- It approximates the full quadratic attention computation with a sparse or low-rank attention pattern, trading some accuracy in the attention output for a reduction in the number of floating-point operations performed
- It reduces the number of floating-point operations attention requires by lowering the numerical precision of the query and key vectors, at the cost of a less numerically exact attention output
- It restructures the attention computation to run entirely within CPU memory instead of GPU memory, trading GPU compute for cheaper CPU compute
Why A? And why not the others?
Correct answer: A. It reduces the number of reads and writes between the GPU's high-bandwidth memory and its much faster on-chip SRAM by tiling the computation and avoiding materializing the full attention-score matrix, while still computing exactly the same attention output as the standard formula, not an approximation of it
Dao et al. observe that standard attention implementations are limited less by the number of floating-point operations they perform and more by how much data must be moved between the GPU's large but comparatively slow high-bandwidth memory (HBM) and its much smaller but much faster on-chip SRAM; FlashAttention restructures the computation using tiling, computing attention in blocks and using recomputation during the backward pass, so that it reads from and writes to HBM far less often, all while still computing exactly the same attention output the standard formula would produce -- it is explicitly an exact algorithm, not an approximation. The option describing a sparse or low-rank approximation is wrong because FlashAttention preserves the exact standard attention computation rather than substituting a different, approximate attention pattern. The option describing lower numerical precision as the source of the speedup is wrong because FlashAttention's gains come from reducing memory movement through tiling, not from reducing precision, and it does not sacrifice numerical exactness of the result. The option describing running entirely in CPU memory is wrong because FlashAttention's optimization is specifically about the memory hierarchy within the GPU itself (HBM versus on-chip SRAM), not about moving computation to the CPU.
Source: Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (2022), arXiv:2205.14135