After learning about prefill, decode, and the KV cache, I realized that most LLM optimizations revolve around them. But instead of getting lost in all the optimizations at once, let's start with the foundation: Attention. In this post, we'll start with Multi-head Attention and then dive into Flash Attention.
Self-attention is where the contextual meaning is calculated, so we have the weights, which we call Q, K, and V
Now, in optimization, all the research work is around memory-bound, compute-bound, and I/O bound
Let's get into these Bounds so we easily understand what's happening inside
Compute Bound (Math is the bottleneck): The CPU is working at 100% maximum capacity doing heavy math or logic. The data is already in memory, and there's no waiting for the hard drive. The CPU simply cannot think any faster.
Memory Bound (Moving is the bottleneck): The CPU is fast enough, but it is constantly waiting to move data in and out of the RAM/VRA1M. The “pipeline” between the memory and the CPU is too narrow, so the CPU starves for data.
I/O Bound (Waiting is the bottleneck): The computer is waiting for something outside the CPU to finish. The CPU is mostly idle. This happens when you are downloading files, reading from a slow hard drive, or waiting for the internet. However, in the world of AI and GPUs, when we say a model is 'Memory Bound,' we specifically mean it is I/O bound by the slow pipeline between the GPU's massive VRAM and its tiny, lightning-fast SRAM.
Flash Attention
Before flash attention, this is how multi-head attention works, and find the reason why flash attention is needed. I wrote the multi-head attention code in PyTorch, so if you want to check it on my GitHub.
So here are two types of memory
Inside your GPU, there are two completely different places to store numbers:
VRAM/HBM: All Data is stored here: the model weights, KV cache, like holding the entire massive AI model, but it is relatively slow to access.
SRAM: This is the GPU's super-fast workspace. It is tiny (usually only about 192 Kilobytes per processor block), but it is blazing fast (roughly 100x faster than VRAM).
========================================================================
INSIDE GPU CHIP
========================================================================
[VRAM/HBM] <------(I/O - Input/Output) ------> [ MATH CORES (SRAM) ]
Holds giant Q,K,V matrices. Tiny,incredibly fast
Size: 24 Gigabytes. calculators.
Speed: SLOW to access. Speed: LIGHTNING fast.
========================================================================The golden rule of AI speed is: Keep the numbers on the fast Workbench (SRAM) as long as possible, and only go to the slow Warehouse (VRAM) when absolutely forced to.
Standard Multi-Head Attention Working :
Here is exactly what standard PyTorch does when we run Q @ K^T, then Softmax, then @ V.
Fetch: Bring Q and K from the (VRAM) to the (SRAM).
Math: Calculate S=Q×KT on the fast SRAM.
The Pain Point (Write to VRAM): The standard kernel2 finishes its job. It takes the giant S matrix and writes it all the way back to the slow (VRAM) to save it.
Fetch again: The Softmax kernel wakes up. It has to fetch that massive S matrix all the way back from the (VRAM) to the Fast (SRAM).
Math: Calculate Softmax on the SRAM.
The I/O-Bound Pain Point (Write to VRAM again): The Softmax kernel writes the large P matrix back to the slow VRAM.
Fetch again: The final Matrix Multiply kernel fetches P and V from the VRAM...
Math: Calculates the final Output.
Why is this I/O Bound? The actual math (multiplying numbers) takes a microsecond. But moving those massive matrices back and forth between the VRAM and the SRAM takes milliseconds. The GPU spends 90% of its time just acting as a truck driver, moving data. It is completely I/O (memory) bound.
=======================================================================
INSIDE THE GPU CHIP
=======================================================================
[VRAM] <-------- Slow Pipe ---------> [SRAM] -----> [MATH CORES]
- Holds the Model Weights (W_q) -The ONLY place
- Holds the Input Data (X) where X * W_q actually
- Does absolutely ZERO math. happens.
========================================================================Here comes into play the FlashAttention
FlashAttention says: “ Once I bring Q, K, and V out of the VRAM storage locker and put them on the SRAM counter, I refuse to put them back in the locker until the entire attention math is 100% finished.”
FlashAttention works:
Smart Fetch: It does not bring the whole Q, K, and V to the SRAM. It brings only a small block, e.g (Words 1 to 64), to the fast (SRAM). Then, for that specific block of Q, it streams through the entire K and V matrices in smaller blocks.
Math: It calculates S=Qblock×KTblock on the fast SRAM.
THE MAGIC (No HBM/VRAM write): It does NOT write S back to the VRAM. It keeps it right there on the fast SRAM.
Math: It applies Softmax right there on the SRAM. (Note: To do this, they had to invent a clever new math trick called “Online Softmax”3 so you can do softmax in chunks without seeing the whole row at once).
Math: It multiplies by Vblock right there on the SRAM.
Store: Only the final, completed Output block is written back to the VRAM.
Flash Attention Visualization
[HBM/VRAM] [SRAM (Counter)] [Math Cores]
Q(128x128)--(SLOW TRIP 1)--> Q_tile (32x32)
K(128x128)--(SLOW TRIP 2)--> K_tile (32x32) --> Multiply Q_tile *
K_tile^T
V(128x128) --(SLOW TRIP 3)--> V_tile (32x32) Result = S_tile
S_tile (STAYS HERE)
|
v
Divide by sqrt(d)
Apply "Online Softmax"
Result = P_tile
|
v
Multiply P_tile * V_tile
Result = O_tile
(O_tile is added to a running total accumulator staying in SRAM)
--LOOP-- (The kernel instantly grabs the NEXT 32x32 tiles of Q, K, V from VRAM)
--LOOP-- (Does the exact same math inside SRAM, adds to accumulator)
--LOOP-- (Until the whole 128x128 matrix is done)
<--(ONLY 1 TRIP!)-- Final_Output (128x128)
Final Output saved!The Secret Trick: Online Softmax
Question: If FlashAttention refuses to write the S matrix back to VRAM, how does it calculate Softmax?
Regular Softmax is impossible here because it needs to see the entire row at once to find the maximum value (to prevent number explosions) and the total sum (for the final division).
Online Softmax solves this by processing the math in chunks using a clever rescaling trick. It keeps just three running variables: a running max (m), a running sum (l), and a running output (O).
Online Softmax Calculation
Here are the exact mathematical formulas for Online Softmax as it processes each block j of the K and V matrices
Variables
Sj = Current block of attention scores QxK(j)T
Vj = Current block of Values
mj = Running maximum
lj = Running sum of exponentials
oj = Running output accumulator
The Update Formulas (for each block j ):
Step 1: Find the max of the current block m′(j)=max(S(j))
Step 2: Update the global running max m′(j)= max(S(j))
Step 3: Update the running sum l(j) =l(j−1)⋅exp(m(j−1)−m(j))+∑exp(S(j)−m(j))
Step 4: Update the running output o(j) =o(j−1)⋅exp(m(j−1)−m(j))+exp(S(j)−m(j))×V(j)
The Final Output (after all blocks are processed):
Output=O(final) / l(final)
(Note: The term exp(m(j−1)−m(j)) is the magic rescaling factor. If a new maximum is found, this term becomes a fraction (< 1), perfectly shrinking the previous work to match the new scale without needing to look back at old data.)
Example: How Online Softmax is calculated
[ INPUTS HELD IN VRAM, FETCHED TO SRAM IN CHUNKS ]
Q = [1, 2]
Block 1: K1 = [[0, 1], [1, 0]], V1 = [4, 5]
Block 2: K2 = [[1, 1], [1, 0]], V2 = [6, 7]
[ SRAM DASHBOARD INITIALIZATION ]
m = -∞
l = 0
O = 0
=========================================================
[ BLOCK 1 ARRIVES IN SRAM ]
=========================================================
S_block = Q @ K1^T = [1, 2] @ [[0, 1], [1, 0]] = [2, 1]
m_new = max(-∞, 2, 1) = 2
# Rescale previous accumulations (results in 0 since m was -∞)
l = 0 * exp(-∞ - 2) = 0
O = 0 * exp(-∞ - 2) = 0
# Add new block's contribution
P_block = exp([2, 1] - 2) = [exp(0), exp(-1)] ≈ [1.000, 0.368]
l = 0 + sum([1.000, 0.368]) = 1.368
O = 0 + ([1.000, 0.368] @ [4, 5]) = 0 + (4.000 + 1.840) = 5.840
m = 2
[ SRAM DASHBOARD STATE ]
m = 2.000 | l = 1.368 | O = 5.840
(Notice: Block 1's S_matrix is completely discarded from SRAM!)
=========================================================
[ BLOCK 2 ARRIVES IN SRAM ]
=========================================================
S_block = Q @ K2^T = [1, 2] @ [[1, 1], [1, 0]] = [3, 1]
m_new = max(2, 3, 1) = 3 NEW MAX FOUND!
# THE MAGIC: Rescale previous accumulations to match new max
Correction_Factor = exp(m_old - m_new) = exp(2 - 3) = exp(-1) ≈ 0.368
l = 1.368 * 0.368 = 0.503
O = 5.840 * 0.368 = 2.149
# Add new block's contribution
P_block = exp([3, 1] - 3) = [exp(0), exp(-2)] ≈ [1.000, 0.135]
l = 0.503 + sum([1.000, 0.135]) = 1.638
O = 2.149 + ([1.000, 0.135] @ [6, 7]) = 2.149 + (6.000 + 0.945) = 9.094
m = 3
[ SRAM DASHBOARD STATE ]
m = 3.000 | l = 1.638 | O = 9.094
=========================================================
[ ALL BLOCKS DONE - FINAL NORMALIZATION ]
=========================================================
Final_Output = O / l = 9.094 / 1.638 ≈ 5.551
(Result perfectly matches standard softmax, but we never
wrote the intermediate S matrix back to slow VRAM!)
