FlashAttention: Tiling and Online Softmax

Tiling and the Challenge of Softmax

From your experience optimizing GEMM, you are already well-acquainted with the power of tiling. By partitioning large matrices into smaller blocks that fit into fast on-chip SRAM, we can drastically reduce the number of costly trips to global memory and maximize data reuse.

However, a naive application of tiling fails for self-attention. The computation C = A B can be tiled because matrix multiplication is based on addition, which is associative. The sum of block-wise products is the same as the product of the full matrices.

The attention formula, however, contains a softmax operation:
$\text{Attention}(Q,K,V)=\text{softmax}(Q{K}^T)V$
Softmax is not associative. The value of softmax(x_i) depends on every other element in its row, as it requires a global denominator ${\sum }_j{e}^{x_j}$. You cannot simply compute softmax on one tile of $Q{K}^T$ and then combine it with the softmax of another tile to get the correct result. This dependency seems to force us to compute the entire $n\times n$ matrix before we can proceed, which is precisely the problem we want to avoid.

The following video provides a refresher on how tiling benefits I/O-bound algorithms and sets up the challenge that softmax presents.

The central question, then, is: how can we reformulate the softmax calculation to work in a block-wise, or "online," fashion?

The Online Softmax Trick

The breakthrough that enables FlashAttention is a method to compute softmax in a streaming manner, updating the result as we process blocks of the input one by one. This is often called "online softmax" or using an "online normalizer."

Let's break down how it works. Recall the "safe softmax" trick from our previous module, where we subtract the maximum value $m$ for numerical stability:
$\text{softmax}(x_i)=\frac{e^{x_i-m}}{{\sum }_j{e}^{x_j-m}}\text{where}m={\max }_j(x_j)$
A standard implementation does this in three passes: one to find $m$, a second to find the denominator $d={\sum }_j{e}^{x_j-m}$, and a third to compute the final values. This involves multiple reads of the input data from HBM.

The online approach merges these passes. Imagine processing the input vector $x$ in blocks. As we process a new block, we can update our running statistics—the current maximum value and the denominator sum—from the previous blocks.

Let's say we have processed some blocks and have a running maximum $m_{\text{old}}$ and a running denominator sum $d_{\text{old}}$. When a new block arrives, we find its local maximum, $m_{\text{local}}$. The new overall maximum is $m_{\text{new}}=\max (m_{\text{old}},m_{\text{local}})$.

If $m_{\text{new}}$ is different from $m_{\text{old}}$, our previous sum $d_{\text{old}}$ is now incorrect because it was normalized with the wrong maximum! The key insight is that we can correct it. We can "rescale" the old denominator by multiplying it by $e^{m_{\text{old}}-m_{\text{new}}}$ and then add the contribution from the new block.

The update rule looks like this:
$d_{\text{new}}=d_{\text{old}}\cdot e^{m_{\text{old}}-m_{\text{new}}}+d_{\text{local}}$
where $d_{\text{local}}$ is the sum for the current block, computed using $m_{\text{new}}$. This allows us to compute the correct final denominator by iterating through the data just once.

The following resource provides a fantastic mathematical intuition and a numerical walkthrough of this exact process.

From Online Softmax to Fused Tiled Attention

The online softmax trick allows us to compute the denominator in a single pass. FlashAttention takes this one step further. Instead of just computing the attention matrix $P$, it fuses the final multiplication with the $V$ matrix directly into the same loop.

To do this, we introduce a third running statistic: the output vector, $o$. The update rule for the output vector is similar in spirit to the denominator update. The previously computed output $o_{\text{old}}$ is rescaled to account for the new maximum, and the contribution from the current block is added.

The full, single-pass tiled algorithm can be summarized as:

1. Initialize running statistics: $m=-\infty$, $d=0$, $o=\vec{0}$.
2. Loop through blocks of $K$ and $V$:
   a. Load a block of $Q$, $K$, and $V$ into fast SRAM.
   b. Compute the dot product for that block: $S_{\text{block}}=Q{K}_{\text{block}}^T$.
   c. Find the new overall maximum, $m_{\text{new}}$.
   d. Update the running statistics $d$ and $o$ using the rescaling trick.
3. After the loop, perform a final normalization of $o$.

This entire process is visualized in the diagram below, which should now be much clearer. The "Outer Loop" corresponds to iterating through the tiles of $K$ and $V$, and the "Inner Loop" iterates through tiles of $Q$. Inside the "Compute Block", the online update of our running statistics takes place.

To solidify your understanding, the next resource provides an exceptional, step-by-step manual calculation. It takes a small example and processes it tile-by-tile, showing exactly how the running statistics are updated, perfectly matching the algorithm.

The data flow within a Streaming Multiprocessor (SM) during this process is highly orchestrated. Different groups of threads (warps) are assigned specific tasks: some load data from HBM to shared memory, MMA warps perform the matrix multiplications, and other warps handle the online softmax and correction calculations before the final result is written out.

Conclusion

In this lesson, we have demystified the core mechanisms of FlashAttention. You have moved from understanding why standard attention is slow to understanding how FlashAttention fixes it without any approximation.

Key Takeaways:

• Tiling: The input matrices $Q,K,V$ are processed in blocks that fit into fast on-chip SRAM, minimizing slow global memory access.
• Online Softmax: This is the key algorithmic trick. By maintaining and updating running statistics ($m,d,o$), we can compute the exact, numerically stable softmax and attention output in a single pass over the tiles.
• Kernel Fusion: The matrix multiplications and the softmax calculation are fused into a single GPU kernel. This avoids writing the massive $n\times n$ intermediate attention matrix to global memory, thereby breaking the quadratic memory bottleneck and solving the I/O problem.

You are now equipped with the conceptual foundation of one of the most important algorithms in modern AI systems. The next logical step is to translate this understanding into code.

In our next lesson, you will begin implementing a simplified FlashAttention kernel. We will start with the inner loop, focusing on computing the output for a single block and correctly updating the running statistics, bringing the theory you learned today one step closer to practical, high-performance execution on the GPU.