FlashAttention-2: Faster Attention with Better…

FlashAttention-2 asks why an exact, I/O-aware attention kernel can still leave much of an A100 idle. Read it to see how reducing non-matmul work, parallelizing along sequence tiles, and giving each warp disjoint query and output rows improves GPU utilization without changing the attention result.

Reading focus: Why scalar softmax bookkeeping and synchronization are expensive relative to tensor-core matrix multiplication on the paper's A100 setup. How sequence-level thread-block parallelism supplies more independent work when batch size and head count are too small to occupy the GPU. Why splitting query rows across warps while sharing key and value tiles removes the cross-warp reduction of partial output rows.

ICLR 2024. Tri Dao. 50 min read, hard difficulty.