What exactly is FlashAttention?

Key Takeaways

What exactly is Flash Attention? Let's try to understand in a simple way. Attention is how a model compares every word with every other word. So, for N tokens, it builds an N by N grid of scores. Double your text and that grid grows four times bigger. Now, here is the surprise. This GPU is rarely slowed down by the math itself. The real bottleneck is moving data around inside the chip. A GPU has two kinds of memory. HBM, the big main memory, huge but slow. And SRAM, right on the chip, tiny but blazing fast. Naive attention computes that whole giant score matrix, writes it out to slow HBM, then reads it all back for the next step. Watch the red arrows. That traffic is what kills speed. Flash Attention flips the approach. It cuts Q, K, and V into small tiles and processes one block at a time entirely inside fast SRAM. Watch the highlight sweep across the matrix block by block. But SoftMax normally needs the full row of scores, so Flash Attention keeps a running SoftMax. It updates the totals as each block arrives. The giant matrix is never built at all. And here is the best part. The answer is exactly the same. No approximation. Just compare the memory reads, way shorter bar, so attention runs several times faster. So, that is Flash Attention. Same math, smarter movement. Keep the work in fast on-chip memory, skip the giant matrix, and long context suddenly gets cheap. If this made it click, hit subscribe. A lot more is coming.