Paper digest: FlashAttention

In 2022, Tri Dao and colleagues at Stanford made the core step in modern AI models run up to three times faster, with identical results. Their method, FlashAttention, is now built into most of the tools used to train and run large language models.

For businesses, it is one reason AI services can handle long documents and chat histories at a reasonable price.

The papers

Why this paper exists

Most AI language models are Transformers, explained in our Attention Is All You Need digest. Their key step, called attention, compares every word in a text with every other word.

That gets expensive fast. Double the length of the text and the work roughly quadruples. Long documents used to run out of memory or take too long.

Other researchers tried shortcuts that skip some comparisons. Those approximate methods often changed the model's answers and, the FlashAttention paper notes, frequently didn't save much real time anyway.

How it works

The authors' insight was that the slow part wasn't the maths. It was moving data around inside the graphics chip, or GPU, that runs the model.

A GPU has two kinds of memory:

• A large main memory, which is relatively slow to read and write.
• A small amount of very fast memory right next to the processors.

The standard method computed a huge table of word-to-word scores, saved it to the slow main memory, then read it back to finish the calculation. FlashAttention instead splits the work into small blocks that fit in the fast memory. It keeps running totals so the blocks combine into exactly the same answer, and never saves the full table.

During training, it even recalculates some values rather than storing them, because recalculating is quicker than fetching them from slow memory.

A simple way to picture it: instead of carrying a giant spreadsheet back and forth to a storeroom, you work on a few pages at a time at your desk and keep a running total.

Results that moved

The method is exact, so models give the same answers as before. The speed gains reported in the paper's abstract:

Saving memory also let models read longer texts, which improved quality. GPT-2 scored better, and models passed a 16,000-token test called Path-X for the first time, scoring 61.4 percent where earlier Transformers did no better than chance.

Later versions

FlashAttention-2 reorganised how work is shared across the chip, roughly doubling speed again. FlashAttention-3 was tuned for Nvidia's newer H100 chips, reaching up to 75 percent of the chip's maximum speed and adding support for lower-precision numbers.

Limits

FlashAttention doesn't remove the underlying cost of comparing every word with every other. It cuts the memory and data movement, which was the bigger bottleneck. For extremely long inputs, other techniques are still needed.

The gains depend on the hardware and the model's shape, so results vary. It also focuses on training. Running models for users, especially in long chats, has its own bottlenecks that later work addressed.

Who should care

If you build AI systems, FlashAttention is now standard, and most popular tools include it.

If you buy AI services, it explains why longer prompts cost more and sometimes run slower. When comparing providers, test speed and price at the prompt lengths you actually use, such as a full brief or a long WhatsApp chat export.

Speed gains like these also make it cheaper to give models more context, which helps with local language work. Our digests on Creole language AI and connecting models to local documents cover that side.

Bottom line

FlashAttention showed that a smarter way of moving data could make AI much faster without changing what it computes. Watch for the next round of chip-specific versions, since each new GPU generation has needed its own tuning.