flash attention

in the last piece i drew attention as a spider feeling its web. that's the idea. but i glossed over what actually happens when you run it on a real gpu, and it turns out the slow part isn't the feeling. it's the fetching.

the real bottleneck

the math is cheap, the moving is not

here's the thing that surprised me when i first dug in. a modern gpu can do far more multiplications per second than it can fetch numbers from its main memory. so for a lot of work the chip is just sitting there, waiting on data. attention is exactly that kind of work. the multiplies are not the problem. how often you walk back to memory is.

the way i think about it: there are two places to keep numbers. a tiny workbench right in front of the compute, fast to reach but small (this is SRAM), and a warehouse across the lot, huge but slow to walk to (this is HBM). everything lives in the warehouse. every trip there and back is the cost that actually shows up in your latency.

where the trips go

standard attention writes the whole table down

let's trace the standard version. you score every query against every key, which gives you an n by n table (that's S = QKᵀ). you softmax each row, then multiply by the values. the naive implementation does these one at a time, and in between it walks the whole table back to the warehouse and fetches it again. for a 4,000-token prompt that table is sixteen million numbers, carried back and forth several times. that walking, not the math, is what you wait on.

the fix

tile it, and keep a running softmax

flashattention's move is simple to say: cut everything into tiles small enough to live on the workbench. a block of queries, with blocks of keys and values streamed past it. for each tile you compute its little patch of scores, softmax it, and fold the result into a running total. the full table never has to exist all at once.

there's one wrinkle, and it's the clever bit. softmax needs every score in a row to normalize, but you're only ever holding one tile of that row at a time. the fix is the online softmax: keep a running maximum and a running sum, and whenever a new tile shows up with a bigger value, gently rescale what you've already accumulated so the running answer stays exactly correct. this is the part i'd make sure i could re-derive from scratch.

build it from scratch

naive, then flash

here they are side by side. same inputs, same output down to floating-point noise. the only real difference is that the second one never builds the full n by n table.

import numpy as np

def softmax(x):
    e = np.exp(x - x.max(-1, keepdims=True))
    return e / e.sum(-1, keepdims=True)

def attention_naive(Q, K, V):
    d = Q.shape[-1]
    S = Q @ K.T / np.sqrt(d)       # (n, n), the whole table lives in memory
    return softmax(S) @ V              # read it back, twice more

↳ clean and correct, and exactly what falls over at long context. S is n by n, and it gets walked to the warehouse and back at every step.

def flash_attention(Q, K, V, Bq=64, Bk=64):
    n, d  = Q.shape
    scale = 1.0 / np.sqrt(d)
    O = np.zeros((n, d))
    for i in range(0, n, Bq):                  # a block of queries
        q   = Q[i:i+Bq]
        m   = np.full((q.shape[0], 1), -np.inf)   # running max
        l   = np.zeros((q.shape[0], 1))         # running sum
        acc = np.zeros((q.shape[0], d))          # running output
        for j in range(0, n, Bk):              # stream the keys / values
            k, v  = K[j:j+Bk], V[j:j+Bk]
            s     = (q @ k.T) * scale          # (Bq, Bk), one tile, fits in SRAM
            m_new = np.maximum(m, s.max(-1, keepdims=True))
            p     = np.exp(s - m_new)          # this tile, rescaled to the new max
            a     = np.exp(m - m_new)          # shrink what we had by the same factor
            l     = a * l + p.sum(-1, keepdims=True)
            acc   = a * acc + p @ v
            m     = m_new
        O[i:i+Bq] = acc / l                    # normalize at the very end
    return O                                  # identical to attention_naive(Q, K, V)

↳ the biggest thing on the workbench is now a single (Bq, Bk) tile, not the full n by n table. the running max, sum, and output are the online softmax holding the answer together. on a real gpu this is one fused kernel, but the logic is exactly what you see here.

what it buys

exact, and much faster

the part i want to underline: this is not an approximation. it isn't dropping any scores the way sparse attention does. you get the identical output, just computed in an order that barely touches slow memory. that's a rare kind of free lunch in systems, and it's why flashattention went from a paper to the default almost overnight.

why it matters

long context got cheaper without changing the model

attention's cost was the wall in front of longer inputs. flashattention didn't move the wall by making the model approximate. it noticed the wall was mostly memory traffic and walked around it. if you take one thing from this, take this: before you change what you compute, look hard at what you're moving. it's often the cheapest speedup you'll find.