{"id":2858,"date":"2025-04-04T07:03:00","date_gmt":"2025-04-04T07:03:00","guid":{"rendered":"https:\/\/mailitics.com\/index.php\/2025\/04\/04\/kernel-case-study-flash-attention\/"},"modified":"2025-04-04T07:03:00","modified_gmt":"2025-04-04T07:03:00","slug":"kernel-case-study-flash-attention","status":"publish","type":"post","link":"https:\/\/mailitics.com\/index.php\/2025\/04\/04\/kernel-case-study-flash-attention\/","title":{"rendered":"Kernel Case Study: Flash Attention"},"content":{"rendered":"

Kernel Case Study: Flash Attention
\n \t

\n
<\/BR>
\n
\n \t

\n
<\/BR><\/p>\n

\n

The attention<\/mdspan> mechanism is at the core of modern day transformers. But scaling the context window of these transformers was a major challenge, and it still is even though we are in the era of a million tokens + context window (Qwen 2.5\u00a0[1]<\/a>). There are both considerable compute and memory bound complexities in these models when we scale the context window (A naive Attention Mechanism<\/a> scales quadratically in both compute and memory requirements). Revisiting Flash Attention lets us understand the complexities of optimizing the underlying operations on GPUs and more importantly gives us a better grip on thinking what\u2019s next.<\/p>\n

Let\u2019s quickly revisit a naive attention algorithm to see what\u2019s going on.<\/p>\n

\"\"
Attention Algorithm. Image by Author<\/figcaption><\/figure>\n

As you can see if we are not being careful then we will end up materializing a full NxM attention matrix into the GPU HBM. Meaning the memory requirement will go up quadratically to increasing context length.<\/p>\n

\n

If you wanna learn more about the GPU memory hierarchy and its differences, my previous post on Triton<\/a> is a good starting point. This would also be handy as we go along in this post when we get to implementing the Flash Attention<\/a> kernel in triton. The flash attention paper<\/a> also has some really good introduction to this.<\/p>\n<\/blockquote>\n

Additionally, when we look at the steps involved in executing this algorithm and its pattern of accessing the slow HBM, (which as explained later in the post could be a major bottleneck as well) we notice a few things:<\/p>\n

    \n
  1. We have Q, K and V in the HBM initially<\/li>\n
  2. We need to access Q and K initially from the HBM to compute the dot product<\/li>\n
  3. We write the output scores back to the HBM<\/li>\n
  4. We access it again to execute the softmax, and optionally for Causal attention, like in the case of LLMs, we will have to mask this output before the softmax. The resulting full attention matrix is written again into the HBM<\/li>\n
  5. We access the HBM again to execute the final dot product, to get both the attention weights and the Value matrix to write the output back to the slow GPU memory<\/li>\n<\/ol>\n

    I think you get the point. We could smartly read and write from the HBM to avoid redundant operations, to make some potential gains. This is exactly the primary motivation for the original Flash Attention algorithm.<\/p>\n

    Flash Attention initially came out in 2022 [2]<\/a>, and then a year later came out with some much needed improvements in 2023 as Flash Attention v2 [3]<\/a> and again in 2024 with additional improvements for Nvidia Hopper and Blackwell GPUs [4<\/a>] as Flash Attention v3 [5<\/a>]. The original attention paper identified that the attention operation is still limited by memory bandwidth rather than compute. (In the past, there have been attempts to reduce the computation complexity of Attention from O(N**2) to O(NlogN)<\/a> and lower through approximate algorithms)<\/p>\n

    Flash attention proposed a fused kernel which does all of the above attention operations in one go, block-wise, to get the final attention output without ever having to realize the full N**2 attention matrix in memory, making the algorithm significantly faster. The term `fused` simply means we combine multiple operations in the GPU SRAM before invoking the much slower journey across the slower GPU memory, making the algorithm performant. All the while providing the exact attention output without any approximations.<\/p>\n

    This lecture, from Stanford CS139, demonstrates brilliantly how we can think of the impact of a well thought out memory access pattern can have on an algorithm. I highly recommend you check this one out if you haven\u2019t already.<\/p>\n

    \n
    \n