Flash Attention builds on the ideas behind online softmax. While online softmax requires two fetches from HBM, Flash Attention only requires one. Just like online softmax, the output matrix can be written as a recurrence formula:
Take one query row and score it against every key, with the usual scaling:
The attention output for that row is a weighted average of the rows of :
Both and the denominator need the whole row before a single output value can be written, and that is what we want to avoid. So take the keys one at a time and define three partial quantities after the first of them:
The prime marks a partial value: each one is what you would get if the row simply ended at element . Note that by construction. So if each of the three can be updated from its previous value plus the one new element, a single pass over and produces the right answer.
The running max. Nothing to derive:
The denominator. This is the online softmax step. Split off the new element, then rewrite the older terms against , which is the max they were actually summed against:
The only real step is . The correction factor does not depend on , so it comes out of the sum and what is left is exactly . One multiply fixes every term already accumulated.
The output. is a ratio, and the division is what makes it look awkward. So set the division aside and give the numerator its own name:
is the same sum as with attached to each term, so the identical split goes through:
Now put the division back. Substitute on both sides:
and divide by :
That is the recurrence. The rescale factor reads as one job: undoes the normalizer was divided by, shifts it to the new max, and dividing by normalizes against the new denominator. Starting values are , , .
Hardware reads memory in tiles, not single elements. Nothing about the recurrence changes: with tiles of width , the max is taken over the tile and the two single new terms become sums over the tile.
Here now counts tiles rather than elements, and is the -th score inside tile . Only one tile of and the three carries live in fast memory at any time.
Every query row runs this loop independently, and rows share the tile of and that is already in fast memory. So the kernel below works on a block of rows at once, and the carries become vectors: mi and li hold one running max and one denominator per row, and Oi holds one partial output row each.
The visualization shows how data moves from HBM to a single block in SRAM, how the intermediate values are calculated, and finally how the output is moved back to HBM. Each block operates identically on different tiles.
template<int Br, int Bc, int D>
__global__ void fa1(const float* Q, const float* K, const float* V,
float* O, float* l, float* m, int N, float scale) {
int bh = blockIdx.x, tid = threadIdx.x;
Q += bh*N*D; K += bh*N*D; V += bh*N*D; O += bh*N*D;
l += bh*N; m += bh*N;
__shared__ float Qi[Br][D], Oi[Br][D], Kj[Bc][D], Vj[Bc][D];
__shared__ float Sij[Br][Bc], mi[Br], li[Br];
int Tr = (N+Br-1)/Br, Tc = (N+Bc-1)/Bc;
j = 0One K/V tile is loaded into SRAM and stays there while every Q block is swept against it.
Kj, Vj = K[j], V[j]: // j = 0
for (int x = tid; x < Bc*D; x += blockDim.x) {
Kj[x/D][x%D] = K[(0*Bc + x/D)*D + x%D];
Vj[x/D][x%D] = V[(0*Bc + x/D)*D + x%D];
}
__syncthreads();
i = 0It carries empty mi, li, Oi from HBM.
Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]: // j = 0, i = 0
for (int x = tid; x < Br*D; x += blockDim.x) {
Qi[x/D][x%D] = Q[(0*Br + x/D)*D + x%D];
Oi[x/D][x%D] = O[(0*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[0*Br+tid]; li[tid] = l[0*Br+tid]; }
__syncthreads();
Sij = Qi @ Kj.T: // j = 0, i = 0
for (int x = tid; x < Br*Bc; x += blockDim.x) {
int r = x/Bc, c = x%Bc;
float s = 0.f;
for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
Sij[r][c] = s*scale;
}
__syncthreads();
mi = max(mi, Sij.max(-1))li = rescale(li) + exp(Sij - mi).sum(-1)Oi = rescale(Oi) + softmax(Sij) @ Vj: // j = 0, i = 0
if (tid < Br) {
float m_ij = -INFINITY;
for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
float m_new = fmaxf(mi[tid], m_ij);
float alpha = expf(mi[tid] - m_new);
float lsum = 0.f;
for (int c = 0; c < Bc; c++) {
Sij[tid][c] = expf(Sij[tid][c] - m_new);
lsum += Sij[tid][c];
}
float l_new = alpha*li[tid] + lsum;
for (int d = 0; d < D; d++) {
float pv = 0.f;
for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
}
mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();
m[i], L[i], O[i] = mi, li, Oi: // j = 0, i = 0
for (int x = tid; x < Br*D; x += blockDim.x)
O[(0*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[0*Br+tid] = mi[tid]; l[0*Br+tid] = li[tid]; }
__syncthreads();
i = 1It carries empty mi, li, Oi from HBM
Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]: // j = 0, i = 1
for (int x = tid; x < Br*D; x += blockDim.x) {
Qi[x/D][x%D] = Q[(1*Br + x/D)*D + x%D];
Oi[x/D][x%D] = O[(1*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[1*Br+tid]; li[tid] = l[1*Br+tid]; }
__syncthreads();
Sij = Qi @ Kj.T: // j = 0, i = 1
for (int x = tid; x < Br*Bc; x += blockDim.x) {
int r = x/Bc, c = x%Bc;
float s = 0.f;
for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
Sij[r][c] = s*scale;
}
__syncthreads();
mi = max(mi, Sij.max(-1))li = rescale(li) + exp(Sij - mi).sum(-1)Oi = rescale(Oi) + softmax(Sij) @ Vj: // j = 0, i = 1
if (tid < Br) {
float m_ij = -INFINITY;
for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
float m_new = fmaxf(mi[tid], m_ij);
float alpha = expf(mi[tid] - m_new);
float lsum = 0.f;
for (int c = 0; c < Bc; c++) {
Sij[tid][c] = expf(Sij[tid][c] - m_new);
lsum += Sij[tid][c];
}
float l_new = alpha*li[tid] + lsum;
for (int d = 0; d < D; d++) {
float pv = 0.f;
for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
}
mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();
m[i], L[i], O[i] = mi, li, Oi: // j = 0, i = 1
for (int x = tid; x < Br*D; x += blockDim.x)
O[(1*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[1*Br+tid] = mi[tid]; l[1*Br+tid] = li[tid]; }
__syncthreads();
j=1K0, V0 are dropped and the next K/V tile takes their place, and the same Q blocks come round again, now carrying the partial results left behind by j = 0.
Kj, Vj = K[j], V[j]: // j = 1
for (int x = tid; x < Bc*D; x += blockDim.x) {
Kj[x/D][x%D] = K[(1*Bc + x/D)*D + x%D];
Vj[x/D][x%D] = V[(1*Bc + x/D)*D + x%D];
}
__syncthreads();
i = 0It carries partial (not empty) mi, li, Oi from HBM
Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]: // j = 1, i = 0
for (int x = tid; x < Br*D; x += blockDim.x) {
Qi[x/D][x%D] = Q[(0*Br + x/D)*D + x%D];
Oi[x/D][x%D] = O[(0*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[0*Br+tid]; li[tid] = l[0*Br+tid]; }
__syncthreads();
Sij = Qi @ Kj.T: // j = 1, i = 0
for (int x = tid; x < Br*Bc; x += blockDim.x) {
int r = x/Bc, c = x%Bc;
float s = 0.f;
for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
Sij[r][c] = s*scale;
}
__syncthreads();
mi = max(mi, Sij.max(-1))li = rescale(li) + exp(Sij - mi).sum(-1)Oi = rescale(Oi) + softmax(Sij) @ Vj: // j = 1, i = 0
if (tid < Br) {
float m_ij = -INFINITY;
for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
float m_new = fmaxf(mi[tid], m_ij);
float alpha = expf(mi[tid] - m_new);
float lsum = 0.f;
for (int c = 0; c < Bc; c++) {
Sij[tid][c] = expf(Sij[tid][c] - m_new);
lsum += Sij[tid][c];
}
float l_new = alpha*li[tid] + lsum;
for (int d = 0; d < D; d++) {
float pv = 0.f;
for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
}
mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();
m[i], L[i], O[i] = mi, li, Oi: // j = 1, i = 0
for (int x = tid; x < Br*D; x += blockDim.x)
O[(0*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[0*Br+tid] = mi[tid]; l[0*Br+tid] = li[tid]; }
__syncthreads();
i = 1It carries partial mi, li, Oi from HBM
Qi, mi, li, Oi = Q[i], m[i], L[i], O[i]: // j = 1, i = 1
for (int x = tid; x < Br*D; x += blockDim.x) {
Qi[x/D][x%D] = Q[(1*Br + x/D)*D + x%D];
Oi[x/D][x%D] = O[(1*Br + x/D)*D + x%D];
}
if (tid < Br) { mi[tid] = m[1*Br+tid]; li[tid] = l[1*Br+tid]; }
__syncthreads();
Sij = Qi @ Kj.T: // j = 1, i = 1
for (int x = tid; x < Br*Bc; x += blockDim.x) {
int r = x/Bc, c = x%Bc;
float s = 0.f;
for (int k = 0; k < D; k++) s += Qi[r][k]*Kj[c][k];
Sij[r][c] = s*scale;
}
__syncthreads();
mi = max(mi, Sij.max(-1))li = rescale(li) + exp(Sij - mi).sum(-1)Oi = rescale(Oi) + softmax(Sij) @ Vj: // j = 1, i = 1
if (tid < Br) {
float m_ij = -INFINITY;
for (int c = 0; c < Bc; c++) m_ij = fmaxf(m_ij, Sij[tid][c]);
float m_new = fmaxf(mi[tid], m_ij);
float alpha = expf(mi[tid] - m_new);
float lsum = 0.f;
for (int c = 0; c < Bc; c++) {
Sij[tid][c] = expf(Sij[tid][c] - m_new);
lsum += Sij[tid][c];
}
float l_new = alpha*li[tid] + lsum;
for (int d = 0; d < D; d++) {
float pv = 0.f;
for (int c = 0; c < Bc; c++) pv += Sij[tid][c]*Vj[c][d];
Oi[tid][d] = (alpha*li[tid]*Oi[tid][d] + pv) / l_new;
}
mi[tid] = m_new; li[tid] = l_new;
}
__syncthreads();
m[i], L[i], O[i] = mi, li, Oi: // j = 1, i = 1
for (int x = tid; x < Br*D; x += blockDim.x)
O[(1*Br + x/D)*D + x%D] = Oi[x/D][x%D];
if (tid < Br) { m[1*Br+tid] = mi[tid]; l[1*Br+tid] = li[tid]; }
__syncthreads();
Counting trips across HBM:
Flash Attention 2 improves on it by putting on the outer loop, so each block of is written just once.