{"id":504,"date":"2025-11-25T09:29:33","date_gmt":"2025-11-25T01:29:33","guid":{"rendered":"https:\/\/high-flyer.in.suopu.cc\/?p=504"},"modified":"2025-11-25T09:29:33","modified_gmt":"2025-11-25T01:29:33","slug":"flashattention%ef%bc%9a%e5%85%b7%e6%9c%89-io-%e6%84%9f%e7%9f%a5%ef%bc%8c%e5%bf%ab%e9%80%9f%e4%b8%94%e5%86%85%e5%ad%98%e9%ab%98%e6%95%88%e7%9a%84%e6%96%b0%e5%9e%8b%e6%b3%a8%e6%84%8f%e5%8a%9b%e7%ae%97","status":"publish","type":"post","link":"https:\/\/high-flyer.in.suopu.cc\/en\/blog\/504\/","title":{"rendered":"FlashAttention: A Novel Attention Algorithm with IO Awareness, Fast and Memory-Efficient"},"content":{"rendered":"

At the heart of the Transformer model is the self-attention mechanism (self-attention), which has both time and storage complexity over the length of the sequence of\u00a0O(N2)<\/span>O<\/span>(<\/span>N<\/span>2<\/span><\/span><\/span><\/span><\/span><\/span><\/span>)<\/span><\/span><\/span><\/span><\/span>\u00a0Level. As the size of Large Language Models (LLMs) continues to grow, equipping LLMs with longer contextual backgrounds presents a very significant challenge in terms of engineering implementation.<\/p>\n

A team of researchers from Stanford University's Department of Computer Science and the State University of New York at Buffalo has published a new attention algorithm called FlashAttention, which not only runs 2-4 times faster than PyTorch's standard attention, but also requires 5-20 times less memory. FlashAttention2, Flash Decoding, and FlashDecoding will be released later with even more dramatic performance speedups.<\/p>\n

Phantom Cube Basic Research and DevelopmentLarge Model Training Tool HAI-LLM<\/a>\u00a0FlashAttention has been adopted across the board to dramatically improve graphics card utilization and achieve excellent training performance. In this series of articles, we will talk about the technology behind FlashAttention and our practical experience.<\/p>\n

paper address<\/strong>\uff1ahttps:\/\/arxiv.org\/abs\/2205.14135<\/a><\/p>\n

Project Source Code<\/strong>\uff1ahttps:\/\/github.com\/Dao-AILab\/flash-attention<\/a><\/p>\n

contexts<\/h2>\n

Traditional attention algorithms whose memory efficiency is\u00a0O(N2)<\/span>O<\/span>(<\/span>N<\/span>2<\/span><\/span><\/span><\/span><\/span><\/span><\/span>)<\/span><\/span><\/span><\/span><\/span>\u00a0The. Some past approaches to optimizing attention mechanisms have used approximations such as sparse approximations, low-rank approximations, and combinations thereof. While these methods can reduce the computation to linear or near-linear (O(N)<\/span>O<\/span>(<\/span>N<\/span>)<\/span><\/span><\/span><\/span><\/span>), but they are overly concerned with reducing the number of floating-point operations (FLops) performed per second and tend to ignore the overhead from memory accesses (IO).<\/p>\n

GPU FLOPS have been growing faster than memory throughput (TB\/s) for years. In our practice of optimizing model training on A100 or equivalent graphics cards, we have found that memory throughput is the bottleneck that affects further training efficiency, and FLOPS and memory throughput need to be tightly coupled in order to fully improve training efficiency. This requires more detailed design at the software level.<\/p>\n

As shown in the figure below:<\/p>\n

\"01\"<\/a><\/span><\/p>\n

The graph above shows the throughput and capacity of different tiers of memory for the CPU and GPU. As you can see memory is not a single component, it is layered in nature and the general rule is: the faster the memory, the more expensive it is and the smaller the capacity.<\/p>\n

Take A100 as an example: A100 GPU has 40~80GB of High Bandwidth Memory (HBM) with a bandwidth of 1.5-2.0 TB\/s, while every 108 stream processors has 192KB of SRAM, with an estimated bandwidth of 19TB\/s. It can be seen that although the SRAM capacity is much smaller, the speedup is 10 times higher. It can be seen that although the SRAM capacity is much smaller, the speed is increased by 10 times, so how to utilize SRAM efficiently is the key to speed up the attention algorithm.<\/p>\n

Standard Attention Algorithm<\/h2>\n

Let's first look at the computational logic behind the standard attention algorithm:<\/p>\n

\"02\"<\/a><\/span><\/p>\n

It can be seen that the standard attention algorithm essentially treats HBM load\/store operations as 0-cost (it is not IO-aware).<\/p>\n

The following figure shows the complete computational time consumption statistics for one Attention operator in the GPT-2 model:<\/p>\n

\"03\"<\/a><\/span><\/p>\n

It can be seen that masking, softmax and dropout operations take up a large amount of time, while matrix multiplication (Matmul), which mainly utilizes FLOPS, takes up only a fraction of the time. Therefore, the FlashAttention algorithm, which is optimized to be hardware IO-aware, is proposed to drastically reduce redundant HBM IO and leverage SRAM for computational acceleration.<\/p>\n

FlashAttention<\/h2>\n

FlashAttention The idea is that since the standard attention algorithm has to write S back to the HBM and this step is only for reloading the computed Softmax, we can save it in SRAM and then write the final result back to the HBM when all the intermediate steps have been performed. as shown in the following figure:<\/p>\n

\"04\"<\/a><\/span><\/p>\n

You can see that FlashAttention fuses multiple operations together by loading from the HBM only once, performing the fused arithmetic operation, and then writing the result back to the HBM.The fusion operation employs the following two main techniques:<\/p>\n