The "attention" of large models may seem like a matrix multiplication, but in reality, it is the most memory intensive and slow down step in training and inference. Flash Attention significantly improves both memory usage and speed for the same attention results by changing the order of calculations. This article explains exactly what problem it solved.
1、 Why is attention slow and memory intensive
Standard attention requires first calculating the similarity matrix between query Q and key K, and then weighting and summing the values V. Assuming the length of the sequence is N, the size of this intermediate matrix is N × N. The longer the sequence, the faster it swells - this is the origin of the statement 'attention complexity is O (N ²)'.
The problem is that traditional implementations will write the intermediate result of N × N back to the graphics memory (HBM), and then read it back from the memory, repeatedly reading and writing. The GPU's video memory bandwidth is limited, and what really slows down the speed is often not the computation itself, but these back and forth transfers.
2、 The core idea of Flash Attention
The approach of Flash Attention (proposed by Tri Dao et al. in 2022) can be summarized as "blocking+fusion":
-Tiling: Cut Q, K, and V into small pieces and move them into the on-chip memory (SRAM) of the computing unit, completing attention calculations on the chip instead of reading and writing the entire video memory.
-Kernel Fusion: Calculate similarity softmax、 The weighted sum of these originally separate steps is merged into one GPU kernel function, reducing the round trip of intermediate results.
-Recombination: During backpropagation, the huge intermediate matrix is not saved, but is recalculated when needed - replacing a large amount of video memory with a little extra computation.
The key is that this is not an 'approximation'. The result calculated by Flash Attention is mathematically equivalent to standard attention, with the only difference being the execution order and memory scheduling. Subsequent versions (FlashAttention 2/3) further optimized parallel partitioning and continued to accelerate on longer sequences and updated hardware.
3、 What did it bring
-The memory usage decreases from increasing with the square of the sequence length to approaching linearity;
-Because it reduces the reading and writing of video memory, the actual speed is actually faster, rather than using precision for efficiency;
-Make longer contextual windows feasible in engineering.
4、 The difference between sparse attention and sparse attention
Sparse attention, sliding window attention, and other methods are "less computational" - sacrificing a portion of visible range or accuracy to reduce complexity. Flash Attention does not reduce computational complexity, it optimizes memory access efficiency, so it can usually be used in conjunction with these methods.
5、 Summary
The significance of Flash Attention is to remind us that the bottleneck of large models is not only in "calculating too much", but also in "moving too much". It has now almost become one of the default components of modern large-scale model training and inference frameworks.
【 Reference Source 】 Comprehensive compilation of industry information released publicly