FlashAttention-2

FlashAttention computes exact attention without ever writing the full N-by-N attention matrix to GPU memory. The naive approach materializes scores for every pair of tokens, costing quadratic memory and flooding the slow high-bandwidth memory with reads and writes. FlashAttention instead streams over blocks of keys and values, keeping running softmax statistics in fast on-chip SRAM, so it produces identical results with far less memory traffic. FlashAttention-2 re-engineers that kernel to close most of the remaining gap to peak GPU throughput.

Version 2 makes three concrete changes. It reduces non-matmul floating-point work, the rescaling bookkeeping of the online softmax, because tensor cores are far faster at matrix multiplies than at other operations. It parallelizes across the sequence-length dimension, not only batch and heads, so long sequences keep every streaming multiprocessor busy. And it improves the work partitioning among warps within a thread block to cut shared-memory synchronization. Together these roughly double throughput over the original, reaching 50 to 70 percent of theoretical floating-point peak on A100-class hardware.

The result is exact, not approximate, matching standard attention up to floating-point ordering, so it is a drop-in speed and memory win that enables longer contexts at fixed memory. Its successor FlashAttention-3 further exploits Hopper asynchrony and FP8; the family is now the default attention backend in most training stacks.