FlashAttention (v2)

FlashAttention is an attention algorithm that computes attention faster and with far less memory traffic, without approximating the result. It works one tile at a time, small enough to fit in the GPU’s fast on-chip memory (SRAM), so the full token-by-token score matrix never travels out to the slower high-bandwidth memory (HBM). FlashAttention-2, released in 2023, is a rewrite of that algorithm that runs about twice as fast as the original.

The simulation below walks the same twelve-token attention through both schemes, one four-by-four tile at a time, with the standard version parking every score it computes in high-bandwidth memory.

FlashAttention-2’s gains come from keeping the GPU busier. It trims bookkeeping work that isn’t matrix multiplication, and it splits a single attention head across thread blocks along the sequence dimension instead of parallelizing only over batch and heads. It also repartitions work between warps, the groups of 32 threads a GPU executes together, so that they exchange less data through shared memory.

The FlashAttention-2 paper reports 50-73% of an A100’s theoretical peak and 225 TFLOPs/s when training a GPT-style model. Its output matches standard attention exactly, so adopting it is a backend choice rather than a modeling decision.

PyTorch can select it under scaled_dot_product_attention, and Hugging Face Transformers exposes it as attn_implementation="flash_attention_2". Attention this cheap is much of what makes long context windows affordable to serve. Later releases follow newer hardware, with FlashAttention-3 targeting Hopper GPUs and FlashAttention-4 extending to Blackwell.

Related Resources

Tutorial

Hugging Face Transformers: Leverage Open-Source AI in Python

As the AI boom continues, the Hugging Face platform stands out as the leading open-source model hub. In this tutorial, you'll get hands-on experience with Hugging Face and the Transformers library in Python.

intermediate ai

For additional information on related topics, take a look at the following resources:

• Python AI: How to Build a Neural Network & Make Predictions (Tutorial)
• How to Use Ollama to Run Large Language Models Locally (Tutorial)
• PyTorch vs TensorFlow for Your Python Deep Learning Project (Tutorial)
• Building a Neural Network & Making Predictions With Python AI (Course)
• Python Deep Learning: PyTorch vs Tensorflow (Course)
• Hugging Face Transformers (Quiz)
• Python AI: How to Build a Neural Network & Make Predictions (Quiz)
• How to Use Ollama to Run Large Language Models Locally (Quiz)