My Personal Favorite FlashAttention (+ KV-Caching) Visualization
July 15, 2026
Note: To be frank, these were some notes that were hiding in my notability that I stumbled across whilst doing a self-review of some key attention concepts. Hopefully they can be of some use to you!
FlashAttention is extremely well-known, but despite this, I can rarely find an explanation of it that truly satisfies me. I’m the kind of person who likes to get into the nitty-gritty of the exact math and computation behind the kernel. This PDF is my understanding of the core algorithm, along with some mathematics that will hopefully help you grok why tiling in this way is correct.
“Teaching” Notes:
-
I start with a review of the attention computation itself. Sidenote: why don’t more blogs include computation graphs of the operations that they are writing about? I find that drawing something out in a consistent computation graph format is one of the best possible ways to truly understand an operation.
-
This covers base attention, then I show how you can do some fancy indexing to do multi-head attention (“MHA”). This is followed by GQA, which is extremely similar with the minor difference of there being fewer KV-heads.
-
THEN we dive into FlashAttention. Despite being such a famous paper, it is intensely difficult to find a satisfying, step by step explanation of the core kernel and algorithm. Here, I tried my best.
-
I first start out with an explanation of why naive attention is slow.
-
Then I introduce a complementary element of the “online softmax” which seems unrelated at first. But actually understanding the math behind this is super important. Combining the two for-loops leads in to the “fusion” idea.
-
Then I have a diagram showing the core kernel, and I structure it showing the “6 main pillars” in green. These are the core things in HBM memory that we are tiling over! Right next to the visual is the actual pseudocode snippet so people can build their understanding bit by bit.
-
I try my best to show all of the intermediate elements that are produced in the subsequent parts of the inner for-loop. This is the part I think most people gloss over! The creation of these intermediate matrices and vectors in SRAM is intensely important and one of the hardest parts about the algorithm.
Hope this helps all you self-studiers out there! Please contact me if there’s any part of my notes you disagree with or find awry.