Attention: The Kernel That Ate the World & FlashAttention v1

Attention: The Kernel That Ate the World & FlashAttention v1

Standard Transformer Attention computes:

\[\text{Attention}(Q, K, V) = \text{Softmax}\left(\frac{Q K^T}{\sqrt{d_k}}\right) V\]

For a sequence length \(N\) and head dimension \(d\) (\(N \gg d\)): - \(S = Q K^T\) generates an \(N \times N\) attention score matrix. - Traditional PyTorch implementations materialize \(S\) and \(P = \text{Softmax}(S)\) in global HBM memory (\(O(N^2)\) memory reads and writes). - For long contexts (\(N = 32k\) or \(128k\)), intermediate matrices consume gigabytes of memory and choke HBM bandwidth.

The Breakthrough: Online Softmax & Tiling

Tri Dao et al. introduced FlashAttention: compute the exact attention output without ever writing the \(N \times N\) intermediate attention matrix to global memory.

Traditional Attention (Memory-Bound):
  Q, K -> [Global HBM Write: S] -> [Global HBM Read: S] -> Softmax -> [Global HBM Write: P] -> [Global HBM Read: P] -> V -> Output
FlashAttention (Fused into SRAM):
  Load Blocks of Q, K, V into On-chip SRAM -> Compute Local Gemm -> Online Softmax Update -> Accumulate Output -> Single Global Write

The Online Softmax Formula

Traditional softmax requires two passes over the sequence: pass 1 to find the maximum \(m = \max(x_i)\), and pass 2 to compute \(\sum e^{x_i - m}\). Online softmax updates the running maximum \(m\) and normalization sum \(l\) incrementally as blocks stream through shared memory:

\[m_{\text{new}} = \max(m_{\text{prev}}, m_{\text{block}})\]

\[l_{\text{new}} = e^{m_{\text{prev}} - m_{\text{new}}} l_{\text{prev}} + e^{m_{\text{block}} - m_{\text{new}}} l_{\text{block}}\]

\[O_{\text{new}} = e^{m_{\text{prev}} - m_{\text{new}}} O_{\text{prev}} + P_{\text{block}} V_{\text{block}}\]