FlashAttention — the same attention,
made memory-aware

00 Start here — the bill Walk 02 quietly ran up

Remember the attention block from Walk 02. For a sequence of lengthT, it does this:

# nanoGPT's manual attention — Walk 02, the slow path
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))  # (B, nh, T, T)
att = att.masked_fill(mask == 0, float('-inf'))
att = F.softmax(att, dim=-1)                                     # (B, nh, T, T)
y   = att @ v                                                    # (B, nh, T, hs)

Look at the shapes: att is (B, nh, T, T). For every pair of tokens, one number. At T = 8192 that's a 67-million-entry matrix, per head, per layer — and it gets written to slow global memory (HBM), read back for the softmax, written again, read again for the@ v. The multiplies are cheap. Shuffling that giant matrix in and out of HBM is the cost, and it grows with T².

01 The move: tile it, and never write the big matrix down

This is Walk 03's lesson — keep data in fast on-chip memory, cross the slow cliff as little as possible — applied to attention. FlashAttention chops Q, K, V into tiles small enough to live in SRAM, and computes attention one tile-pair at a time, accumulating the final output as it goes. The full T×T scores exist only momentarily, one small block at a time, on-chip.

// Everything lives in SRAM (on-chip), never a full T×T matrix in HBM
extern __shared__ float sram[];
int tile_size = Bc * d;          // one tile of Q, K, or V
float* Qi = sram;                // a block of queries
float* Kj = &sram[tile_size];    // a block of keys
float* Vj = &sram[tile_size * 2];// a block of values
float* S  = &sram[tile_size * 3];// scores for THIS tile only — Bc×Br, not T×T

Notice S: it's sized Bc × Br — one tile, maybe 32×32 — notT × T. That single decision is the whole algorithm. But it raises a problem: softmax needs the whole row at once (you divide by the sum over all keys). If you only ever see one tile of keys at a time, how do you normalize? That's the next section, and it's the clever bit.

02 The trick that makes it possible: online softmax

Softmax over a row is exp(sᵢ − max) / Σ exp(sⱼ − max) — you subtract the row max for numerical safety, then divide by the row sum. Both the max and the sum need the entire row. FlashAttention computes them incrementally, updating a running max m and running sum l as each new tile arrives — the same math a streaming average uses.

// S = QKᵀ for this tile, and track the running row-max
float row_m = -INFINITY;
for (int y = 0; y < Bc; y++) {
    float sum = 0;
    for (int x = 0; x < d; x++)
        sum += Qi[(tx * d) + x] * Kj[(y * d) + x];   // one query · one key
    sum *= softmax_scale;
    S[(Bc * tx) + y] = sum;
    if (sum > row_m) row_m = sum;                    // rowmax(S), online
}

// P = exp(S − row_m), and the running row-sum
float row_l = 0;
for (int y = 0; y < Bc; y++) {
    S[(Bc * tx) + y] = __expf(S[(Bc * tx) + y] - row_m);
    row_l += S[(Bc * tx) + y];                       // rowsum(P), online
}

For the current tile it computes the scores S = QKᵀ, this tile's local max (row_m), the exponentials P = exp(S − row_m), and this tile's local sum (row_l). All in SRAM. Nothing full-width is stored.

First get a feel for the softmax itself — the same function attention runs over its scores. Drag the logits and watch the probabilities redistribute; the running max and sumthe kernel tracks are just the incremental way to compute exactly this:

The full derivation — logits, softmax, and the cross-entropy loss it feeds — is inStudy · Module 2 — Probability & the loss.

03 Reconcile: merge each tile into the running answer

When a new tile gives a bigger max than we've seen, everything accumulated so far was normalized against the old max — so it must be rescaled. That's these lines: combine old and new stats, correct the running output O, add the new tile's contribution, carry the updatedm and l forward.

// merge THIS tile's stats with everything seen so far
float row_m_new = max(row_m_prev, row_m);
float row_l_new = (__expf(row_m_prev - row_m_new) * row_l_prev)
                + (__expf(row_m     - row_m_new) * row_l);

// rescale the running output O and add this tile's contribution — no T×T matrix ever stored
for (int x = 0; x < d; x++) {
    float pv = 0;                                     // Pij · Vj
    for (int y = 0; y < Bc; y++)
        pv += S[(Bc * tx) + y] * Vj[(y * d) + x];
    O[... i ...] = (1 / row_l_new)
        * (row_l_prev * __expf(row_m_prev - row_m_new) * O[... i ...]   // rescale old
        +              __expf(row_m      - row_m_new) * pv);            // add new
}
m[...] = row_m_new;   l[...] = row_l_new;             // carry the stats forward

The output O is updated in place, tile by tile. At no point does aT × T matrix exist in HBM. When the loops finish, O holds exactly whatsoftmax(QKᵀ) @ V would have produced — computed without ever paying to store the intermediate.

04 The whole shape of the kernel

Stepping back, the entire forward pass is two nested loops over tiles:

for (int j = 0; j < Tc; j++) {                 // outer loop over KEY/VALUE tiles
  // pull one tile of K and V from slow HBM into fast SRAM, once
  for (int x = 0; x < d; x++) {
    Kj[(tx * d) + x] = K[qkv_offset + (tile_size * j) + (tx * d) + x];
    Vj[(tx * d) + x] = V[qkv_offset + (tile_size * j) + (tx * d) + x];
  }
  __syncthreads();

  for (int i = 0; i < Tr; i++)  {              // inner loop over QUERY tiles
    // ... the whole attention math for this (query-tile × key-tile) block ...
  }
  __syncthreads();
}

05 Why doing more work is faster

FlashAttention performs extra arithmetic — the rescaling, and in the backward pass it evenrecomputes the scores instead of storing them. By a FLOP count it does more than naive attention. Yet it runs several times faster and uses memory that grows with T instead ofT². The only way that reconciles:

This is only possible because attention was memory-bound — the math units were idle, waiting on HBM. Trading cheap idle FLOPs for expensive memory traffic is a winning trade every time the limiter is memory. FlashAttention is the most famous instance of the exact principle Walk 03 taught kernel by kernel.

06 The taste to take away