Naar inhoud springen

Deze video en de bijbehorende tekst zijn in het Engels.

FlashAttention: same attention result, without the giant memory bill

AI Concepts #124

FlashAttention gets exactly the same attention result tile by tile, never storing the huge grid of scores. Here is the memory wall it answers, in plain words.

In detail

What FlashAttention is

FlashAttention is a smarter way to compute attention, the step where a language model scores every word against every other word. Ordinary attention writes all those scores into one big grid and keeps it in memory. FlashAttention works through the grid in small blocks, called tiles. It uses each tile and lets it go before moving on, so the full grid never sits in memory at once. The final answer is exactly the same as the ordinary method. What changes is how much memory the job needs along the way.

What the video shows

The short sets up the problem first: normal attention stores a giant grid of scores, big enough to fill a whole graphics card, the chip that does the heavy maths in AI. Then it shows the fix. FlashAttention computes that grid tile by tile and never stores all of it. Same result, far less memory.

The wall it answers

The grid is big because it is quadratic: its size is the text length multiplied by itself, so twice the text gives four times the grid. The course measured where that leads. It tested one attention layer with 8 heads (parallel copies of attention that each look for different links) on an 8 GB laptop graphics card. At 16,384 tokens, which are small pieces of text, roughly words, the score grid alone would have needed 8 GB, and the card ran out of memory. Not the whole model: one temporary grid, inside one layer.

The course names that wall as the reason FlashAttention exists.

An everyday example

Imagine checking out a huge trolley at the supermarket. One way would be to unload every item onto an enormous table, look at them all together, and only then add them up. What a cashier actually does is scan a few items, add them to a running total, bag them and move on. The bill at the end is identical. The second way just never needs a table big enough for the whole trolley.

FlashAttention is the cashier. The real calculation needs some extra bookkeeping to stay exact, but the idea holds: keep a small running summary, not the whole grid.

A common misconception

FlashAttention is not a shortcut or an approximation. It still compares every word with every other word, and it produces the same attention result as the ordinary method. The course is precise about what kind of improvement it is: a memory saving first, and a speed-up only second. The amount of comparing still grows quickly with the length of the text. What disappears is the need to keep the giant grid sitting in memory.

Why it matters

In the course's test, it was memory that stopped a single layer at 16,384 tokens, not time. By never storing the full grid, FlashAttention goes straight at that limit on how much text attention can handle on a given card. It is also a useful lesson in engineering: sometimes the fix is not to do less maths, but to keep less of it lying around.

If you have heard that long texts are costly because attention is quadratic, this is one of the ways engineers push back on the memory half of that cost.

Learn it step by step in Chapter 9 of our free course AI From Scratch: Attention and the Transformer Block, Derived From an Average.

Ook op

Meer uit dit hoofdstuk

Tekst geschreven met AI-hulp op basis van hoofdstuk 9 van onze gratis cursus.