The Details of Flash Attention

From the online-softmax derivation to eight CUDA kernels — a 121.8× speedup, and the two optimizations that made things slower.

Overview

The computation of self-attention is a key component in transformer models. It computes the attention weights from the query and key vectors and performs a weighted sum of the value vectors. This step is expensive in both time and memory, because it is memory-bound and compute-bound at the same time.

This post covers both halves of the problem. The first half is the online-softmax derivation that lets attention be computed in a single pass over the data instead of three. The second half is what it takes to actually write that single pass in CUDA — eight kernel variants, ending 121.8× faster than the naive one.

It is not, however, a list of eight optimizations that each make things faster. Two of them made things slower, and both are worth keeping. The code is in Minimal_Flash_Attention, and all timings below are N=1024, d=1024, averaged over 10 runs on an NVIDIA RTX 4500 Ada.

The Attention Computation

The inputs \(Q, K, V\) typically have shapes \((B, \text{nhead}, N, nd)\). Since the batch and head dimensions are independent from the attention weights computation, we focus on the sequence length \(N\) and the embedding dimension \(nd\). The attention formula is

\[\begin{equation} O=\text{Attn}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V~. \end{equation}\]

The k-th row of the output matrix is

\[\begin{equation} O[k, :]=\sum_{i=1}^N X_s[k, i]\cdot V[i, :]~, \end{equation}\]

where

\[\begin{align} X_s[k, i]&=\text{softmax}\{i=1:N\vert x_i=X[k,i]=\frac{Q[k, :]K[i, :]^T}{\sqrt{d_k}}\}\\ &=\{i=1:N|\frac{e^{x_i}}{\sum_{j=1}^N e^{x_j}}\}~. \end{align}\]

Safe Softmax and the Three Passes

We cannot simply exponentiate the dot products, because \(e^{x_i}\) can easily overflow floating-point limits. To fix this we use safe softmax: find the maximum value in the row (\(m_N\)) and subtract it from every element before exponentiation.

That turns the computation into three distinct loops over the data:

  1. Find the global max (\(m_N\)): iterate through the row to find \(m_N = \max(x_i)\).
  2. Compute the global sum (\(d_T\)): iterate again to compute the normalization constant \(d_T = \sum_{i=1}^{N} e^{x_i - m_N}\).
  3. Compute the output: iterate a third time to calculate the weighted sum \(O[k, :] = \sum_{i=1}^{N} \frac{e^{x_i - m_N}}{d_T} \cdot V[i, :]\).

This forces the calculation into multiple passes. We cannot compute the output until we have finished scanning the entire row to find \(m_N\) and \(d_T\). Repeatedly reading and writing the \(N \times N\) matrix is what makes standard attention memory-bound.

So the question is: can we merge these blocks?

Online Softmax

We need to compute the softmax incrementally as we read the data, rather than waiting for the global statistics. Instead of calculating the final max and sum upfront, we define running variables for the \(i\)-th step:

  1. the running maximum: \(m_i\);
  2. the running sum of exponentials: \(d'_i=\sum_{j=1}^{i} e^{x_j - m_i}\).

Now consider loop 1, \(i=1 \rightarrow N\). When moving from step \(i-1\) to \(i\) we introduce a new value \(x_i\). If \(x_i > m_{i-1}\), our maximum changes, and we must rescale our previous accumulations to account for that change.

  1. Updating the max: \(m_i = \max(m_{i-1}, x_i)\)
  2. Updating the sum (\(d'_i\)) by rescaling the previous sum (\(d'_{i-1}\)) by the exponential difference between the old and new max: \(d'_i = d'_{i-1} \cdot e^{m_{i-1} - m_i} + e^{x_i - m_i}\)

Upon reaching the N-th step, we have \(d'_T=d_T\). Then we start loop 2, \(i=1 \rightarrow N\),

\[O[k, :] = \sum_{i=1}^{N} \frac{e^{x_i - m_N}}{d'_T} \cdot V[i, :]\]

Fusing the Output Accumulator

Can we merge these two loops? Yes, but the second loop is dependent on \(m_N\) and \(d'_T\). So we apply the same rescaling logic to the partial output accumulator, defining \(O'_i = \sum_{j=1}^i \frac{e^{x_j - m_i}}{d'_i} \cdot V[j, :]\). As a result, the computation of the output matrix \(O[k, :]\) can be recast into a geometric progression:

\[\begin{align} O'_i &= \sum_{j=1}^i \frac{e^{x_j - m_i}}{d'_i} \cdot V[j, :] \\ & = \sum_{j=1}^{i-1} \frac{e^{x_j - m_i}}{d'_i} \cdot V[j, :] + \frac{e^{x_i - m_i}}{d'_i} \cdot V[j, :] \\ & = \sum_{j=1}^{i-1} \frac{e^{x_j - m_{i-1}}}{d'_{i-1}} \cdot \left( \frac{d'_{i-1}}{e^{x_j - m_{i-1}}} \cdot \frac{e^{x_j - m_i}}{d'_i} \right) \cdot V[j, :] + \frac{e^{x_i - m_i}}{d'_i} \cdot V[i, :] \\ & = O'_{i-1} \cdot \left( \frac{d'_{i-1}}{d'_i} \cdot e^{m_{i-1} - m_i} \right) + \frac{e^{x_i - m_i}}{d'_i} \cdot V[i, :] \end{align}\]

Upon reaching the N-th step with these recursive formulas, we have \(O'_T[k, :] = O[k, :]\). The output can now be computed in a single pass, without ever materializing the \(X\) matrix and without reading and writing it repeatedly to get the global max and sum upfront.

Tiling the Loop

So the computation reduces to this loop, parallelizable over each element of the output matrix \(O[k, j]\):

\[\begin{align*} \textbf{for } i &\textbf{ in } \{1,\cdots, N\} \\ x_i &\leftarrow Q[k,:] K^T[:,i] \\ m_i &\leftarrow \max(m_{i-1}, x_i) \\ d'_i &\leftarrow d'_{i-1} e^{m_{i-1} - m_i} + e^{x_i - m_i} \\ \mathbf{o}'_i &\leftarrow \mathbf{o}'_{i-1} \frac{d'_{i-1} e^{m_{i-1} - m_i}}{d'_i} + \frac{e^{x_i - m_i}}{d'_i} V[i,:]\\ O[k,:] & \leftarrow \mathbf{o}'_N \end{align*}\]

To compute \(O[k, j]\) we need to access column vector \(V[:, j]\), row vector \(Q[k, :]\), and the full \(K[:, :]\) matrix. We could use step size \(Bc\) to reduce the number of iterations:

\[\begin{align*} \textbf{for } i \textbf{ in } &\{1,\cdots, N/Bc\} \\ \mathbf{x}_i &\leftarrow Q[k,:] K^T[:,i-1*Bc:i*Bc] \\ m_i^{local} &\leftarrow \max(\mathbf{x}_i) \\ m_i &\leftarrow \max(m_{i-1}, m_i^{local}) \\ d'_i &\leftarrow d'_{i-1} e^{m_{i-1} - m_i} + \sum_{j=1}^{Bc} e^{\mathbf{x}_i[j] - m_i} \\ \mathbf{o}'_i &\leftarrow \mathbf{o}'_{i-1} \frac{d'_{i-1} e^{m_{i-1} - m_i}}{d'_i} + \sum_{j=1}^{Bc} \frac{e^{\mathbf{x}_i[j] - m_i}}{d'_i} V[j+(i-1)*Bc,:]\\ O[k,:] \leftarrow & \mathbf{o}'_{N/Bc} \end{align*}\]

Let’s assume Q,K,V are Nxd matrices and each thread handles one element in the output matrix O[row, col]. To compute O[row, col], we iterate over the full K matrix through the dimension N, and access row vector Q[row, :] and col vector V[:, col]. We adopt tile based parallelism where the block size is Br x Bc. Br is for the rows in Q, Bc is for the columns in V. This means we need to access Br rows of Q[row:row+Br, :] and Bc columns of V[:, col:col+Bc]. So we have a grid of (N/Br) x (d/Bc), and the global index of a thread is

int row = blockIdx.x * blockDim.x + threadIdx.x;
int col = blockIdx.y * blockDim.y + threadIdx.y;
int steps = (N + Bc - 1) / Bc;
float O = 0.0f; // O[row, col]
for (int i = 0; i < steps; i++) {
    int start = i * Bc;
    int end = min((i + 1) * Bc, N);
    for (int j = start; j < end; j++) {
        // access Q[row, :], K[j, :] and V[j, col]
    }
}

That is the algorithm. The rest of this post is what happens when you write it down in CUDA.

The Baseline

The most direct translation of the loop above gives a kernel where each thread owns one output element and keeps the running max, running sum, and partial output in registers. The per-tile reduction — the local max and the local sum over \(Bc\) elements — is done cooperatively in shared memory, and the block loops over \(j\) in steps of \(Bc\).

It runs, and it is correct, and it takes 258 ms. Almost all of that is wasted work in two places.

The first is that nothing is staged on-chip. Each thread walks the full head dimension of both \(Q\) and \(K\) straight from global memory, and the values it walks are values its neighbours are also walking. threadIdx.x indexes the row of \(Q\), so for a fixed threadIdx.x all \(Bc\) threads at different threadIdx.y read the same row of \(Q\); and threadIdx.y indexes the row of \(K\), so for a fixed threadIdx.y all \(Br\) threads read the same row of \(K\). Per tile iteration the block issues \(Br \cdot Bc \cdot d\) loads of each where \(Br \cdot d\) and \(Bc \cdot d\) would do — a factor of \(Bc\) and \(Br\) respectively, all of it redundant.

The second is the access pattern. Threads in a warp vary threadIdx.x fastest, and threadIdx.x indexes the row of \(Q\), so consecutive threads read inputQ[row * d + index] for rows that differ by 1 — addresses that differ by d floats. Every one of those reads is a separate transaction.

Kernel 2: Coalesced Shared Memory Tiles

The fix for both is the same, and it is the standard one: stage the tiles in shared memory so the block loads each element once, with coalesced global accesses, and all reuse happens on-chip.

__shared__ float Qds[Br * Bc];
__shared__ float Kds[Bc * Bc];
__shared__ float Vds[Bc * Bc];
...
for (int ph = 0; ph < gridDim.x; ph++) {
    Qds[block_index] = inputQ[row * d + threadIdx.x + ph * Bc];
    Kds[block_index] = inputK[k_row_index * d + threadIdx.y + ph * Bc];
    __syncthreads();
    for (int index = 0; index < Bc; index++)
        qk_dot = std::fma(Qds[start + index], Kds[index * Bc + threadIdx.x], qk_dot);
    __syncthreads();
}

Now each thread reads a row of shared memory instead of a strided column of global memory, and the global loads are contiguous across the warp. 258 ms → 32.9 ms, a 7.8× speedup from memory access patterns alone. No arithmetic changed.

This is worth pausing on, because it is the single largest win in the whole series and there is nothing clever in it. Kernel 2 does the same arithmetic as kernel 1. It just stops re-fetching the same values from global memory.

Kernel 3: Register Tiling and Warp Shuffles

Kernel 2 still moves a lot of data through shared memory that never needed to leave registers. After the \(QK^T\) tile is computed, its values are written to shared memory, read back for the max, written back again for the exponentials, and read again for the \(PV\) product. Each round trip costs a __syncthreads().

Kernel 3 gives each thread a small tile of outputs — Rq = 2 rows by Rv = 4 columns — so the scores live in registers across the whole online-softmax update. The reduction over \(Bc\) also moves from shared memory to warp shuffles:

template <template <typename> class ReductionOp, typename T, int thread_group_width = warpSize>
__inline__ __device__ T WarpAllReduce(T val) {
    for (int mask = thread_group_width / 2; mask > 0; mask >>= 1)
        val = ReductionOp<T>()(val, __shfl_xor_sync(0xffffffff, val, mask));
    return val;
}

A shuffle is one instruction and needs no synchronization, where the shared-memory tree in kernels 1 and 2 needed a __syncthreads() at every halving of strip. 32.9 ms → 6.11 ms. The register tile also gives the compiler independent FMA chains to interleave, which is most of the rest of it.

Kernels 4 and 5: One Win and One Regression

Kernel 4 widens the register tile to Rq = 3, Rv = 4 and improves buffer reuse: 4.93 ms, a modest 1.24× over kernel 3.

Kernel 5 adds float4 vectorized loads. It comes in at 5.76 ms — slower than kernel 4.

I did not profile this, so I will not claim a cause with confidence. But it is a familiar shape of result: float4 access changes the shared-memory access pattern, and a layout that is conflict-free at 4-byte granularity need not stay conflict-free at 16. It also raises per-thread register pressure, which costs occupancy. Whatever the balance of those, the honest summary is that the kernel was already close enough to balanced at Br = Bc = 32 that widening the accesses did not pay for itself.

I kept it in the series anyway. A progression of eight kernels that all get faster is a filtered result, not a measured one.

Kernel 6: Shared Memory Layout

The next real win comes from the layout of the shared-memory tiles, not from anything arithmetic.

Shared memory on NVIDIA hardware is 32 banks of 4 bytes, and a warp’s access is served in one cycle only if the 32 threads hit 32 distinct banks. Store the tile row-major as [m][k] and have a warp read a column of it — same k, different m — and every thread’s address differs by the row length. Here that length is numQ = Rq * Br = 128, a multiple of 32, so all 32 threads land on the same bank and the access serializes 32 ways.

Kernel 6 transposes the store, so the stride between adjacent threads becomes Rq = 8 floats instead of a full row:

shareQK[(4 * smem_a_k + id) * numQ + smem_a_m] = a[id];

after which the consumption side

com_a[0] = (float4 &)shareQK[index * numQ + threadIdx.y * Rq];

differs by 8 floats across threadIdx.y. That spreads 16 threads over 4 banks rather than piling all of them onto one. It is not perfectly conflict-free — 8 floats apart is still a 4-way conflict at 4-byte granularity — but it is a long way from a 32-way stall. Small tiles (Br = Bc = 16) with large register tiles (Rq = Rv = 8) keep 8×8 outputs per thread. 5.76 ms → 2.41 ms, recovering kernel 5’s regression and then some.

Kernel 7: Overlapping Load and Compute

Every kernel so far stalls the same way: load a tile, synchronize, compute on it, synchronize, load the next. The global load cannot begin until the previous tile has been fully consumed, and nothing overlaps.

Kernel 7 double-buffers. The shared-memory arrays are indexed by the parity of the current step, so the loads for step ph write into the buffer that step ph-1 is not reading:

com_a[0] = (float4 &)shareQK[index * numQ + threadIdx.y * Rq + (ph - 1) % 2 * numQ * Bd];
...
shareQK[(4 * smem_a_k) * numQ + smem_a_m + (ph % 2) * numQ * Bd] = a[0].x;

Loads issued for the next tile proceed while the current tile’s FMAs execute. 2.41 ms → 2.12 ms. That is only 1.14×, and it is a fair signal that this style of kernel is near its practical limit — four further kernels past kernel 3 have bought 2.9×.

Kernel 8: Tensor Cores

The last variant is the one people expect to win: a WMMA kernel using FP16 tensor cores with 16×16×16 tiles.

It is 15.76 ms — slower than every kernel from 3 onward, and 7× slower than kernel 7.

The reasons are structural rather than a bug. The tile is 16×16, so each mma_sync is preceded and followed by moving data through shared memory with __syncthreads() between every stage: load Q tile, sync, load K tile, sync, matmul, store S, sync, softmax, store P, sync, load V, sync, matmul. The ratio of tensor-core work to synchronization is tiny. One warp per block also means occupancy is low, and for d = 1024 the inner loop over the head dimension re-tiles Q and K for every sequence block.

There is also a precision cost. FP16 accumulation of the scores gives a max difference of 1.65e-02 against the FP32 reference, well outside the 1e-4 tolerance the other kernels meet. Kernel 8 is excluded from the pass/fail validation for that reason.

Real FlashAttention kernels are fast on tensor cores, but not by shrinking a straightforward kernel to a 16×16 tile. They use much larger tiles, cp.async so the loads do not occupy registers, warp specialization so different warps load and compute concurrently, and they keep the accumulation in FP32. Kernel 8 shows the mechanics of WMMA; it is not a recipe.

Results

Kernel Change Time (ms) vs. previous
1 Baseline tiling 258.18 —
2 Coalesced shared-memory tiles 32.90 7.8×
3 Register tiling, warp shuffles 6.11 5.4×
4 Wider register tile, buffer reuse 4.93 1.24×
5 float4 vectorization 5.76 0.86×
6 Transposed shared-memory layout 2.41 2.39×
7 Double-buffered pipeline 2.12 1.14×
8 WMMA tensor cores (FP16) 15.76 0.13×

Kernel 7 is 121.8× faster than kernel 1. But the distribution matters more than the total. Measured in time actually removed, kernel 2 accounts for 88% of the entire improvement (225 ms of the 256 ms total) and kernel 3 for another 10%. Everything past kernel 3 — four more kernels and a regression along the way — accounts for the last 1.6%. Kernel 8 then gives all of it back and more.

The two that lost are the more instructive ones. Vectorization that changes your shared-memory access pattern is not a free win. And tensor cores are only as fast as the data movement around them.

Validation

Correctness is checked against a NumPy reference rather than against the other kernels, so a shared bug cannot hide:

python3 flash.py      # reference attention in NumPy -> python_output_standard.txt
./run.sh              # builds, runs kernels 1-8, compares each to the reference

Kernels 1 through 7 pass at a tolerance of 1e-4. Kernel 8 is compared but not gated on tolerance, for the FP16 reason above.

The derivation in the first half is the standard online-softmax reformulation, and the interesting part of the second half was how much of the result came from the first two kernels and how little from the last five.