Loading editor...

Kernel visualization: Flash Attention 2

For a single query row , stream the key/value tiles of width . Carry a running max , denominator , and output :

Initial values: , , . Only one tile of and the three carries live in fast memory at any time.

Steps

One CUDA block is responsible for Br rows of Q from start to finish, so it begins by claiming its shared memory and working out which rows those are.

template<int Br, int Bc, int D>
__global__ void fa2(const float* Q, const float* K, const float* V,
                    float* O, float* L, int N, float scale) {
    int i = blockIdx.x, bh = blockIdx.y, tid = threadIdx.x;
    Q += bh*N*D; K += bh*N*D; V += bh*N*D; O += bh*N*D; L += 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 Tc = (N+Bc-1)/Bc;
}

Visualizing operations in a single block

Zoom in on one block: Qi is loaded once and stays for the whole run, and the carries start empty.

  • Qi = Q[i]; Oi = 0; mi = -inf; li = 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] = 0.f;
}
if (tid < Br) { mi[tid] = -INFINITY; li[tid] = 0.f; }
__syncthreads();

Iteration j = 0

The first tile of K and V arrives, and the running max, the running sum and the output all move for the first time.

  • Kj, Vj = K[j], V[j]: 
// i = 0, 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();
  • Sij = Qi @ Kj.T: 
// i = 0, j = 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: 
// i = 0, j = 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];
    }
    li[tid] = 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*Oi[tid][d] + pv;
    }
    mi[tid] = m_new;
}
__syncthreads();

Iteration j = 1

The second tile arrives, so everything held from the first tile is now stale and gets rescaled to the new max before this tile is added.

  • Kj, Vj = K[j], V[j]: 
// i = 0, 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();
  • Sij = Qi @ Kj.T: 
// i = 0, j = 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: 
// i = 0, j = 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];
    }
    li[tid] = 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*Oi[tid][d] + pv;
    }
    mi[tid] = m_new;
}
__syncthreads();

Write the data back to HBM

The tiles are done, so only now is Oi divided by li, and mi and li leave together as the single array L.

  • O[i], L[i] = Oi/li, mi + log(li): 
// i = 0
if (tid < Br) {
    for (int d = 0; d < D; d++) Oi[tid][d] /= li[tid];
    L[0*Br + tid] = mi[tid] + logf(li[tid]);
}
__syncthreads();
for (int x = tid; x < Br*D; x += blockDim.x)
    O[(0*Br + x/D)*D + x%D] = Oi[x/D][x%D];

Same operations in another block

Every other block of Q runs these same steps, and because no block needs anything from another, they all run at the same time.

Operation Action
Qi = Q[i]; Oi = 0; mi = -inf; li = 0
Kj, Vj = K[j], V[j]
Sij = Qi @ Kj.T
mi, li, Oi = ...
Kj, Vj = K[j], V[j]
Sij = Qi @ Kj.T
mi, li, Oi = ...
O[i], L[i] = Oi/li, mi + log(li)

Counting trips across HBM:

  • is read once, and and are written once and never read back
  • and are read twice, once per block of
    • Why is nothing carried back and forth now? sits on the outer loop, so one block of is built from start to finish in SRAM. It is already the real answer by the time it is written, and and never have to make the trip.

Two smaller wins come with it. The division by happens once at write-back instead of on every step, and and leave as a single array , which is what the backward pass needs.