Does FlashAttention change what the model computes?

No — it produces mathematically the same attention output as the standard formula. It only changes how the computation is scheduled on the GPU to reduce memory traffic, so it's a pure speed/memory optimization.

Why is memory bandwidth the bottleneck rather than FLOPs?

Modern GPUs can perform far more arithmetic per second than they can move data between HBM and on-chip SRAM. Standard attention implementations spend most of their time waiting on memory reads/writes of the large intermediate score matrix, not on the matrix multiplications themselves.

What is 'online softmax'?

A technique for computing softmax incrementally over blocks of data using running maximum and sum statistics, so the full attention matrix never needs to be stored in memory all at once.

Do I need to write custom CUDA code to use FlashAttention?

No — it's available as a pip-installable library and is also built into PyTorch's scaled_dot_product_attention and Hugging Face Transformers, so most users get it automatically or with a single flag.

What's new in FlashAttention-3?

It targets NVIDIA Hopper-generation GPUs (H100) specifically, adding support for FP8 low-precision computation and asynchronous warp-specialized execution to better exploit the hardware's capabilities.

How does FlashAttention relate to KV caching during inference?

They address different stages — FlashAttention speeds up the attention computation itself, while KV caching avoids recomputing key/value projections for previously generated tokens; production inference engines use both together.

AI Fundamentals

Flash Attention

By Tri Dao et al.

AdvancedTechnique10.9K learners

FlashAttention is an IO-aware exact attention algorithm that reorders and fuses the computation of transformer self-attention to minimize reads and writes to GPU high-bandwidth memory, delivering significant speedups and memory savings…

#FlashAttention #AIFundamentals#Technique#Advanced#SparseAttention #KVCache #GroupedQueryAttention #RotaryPositionalEmbedding #ArtificialIntelligence#Glossary#SkillVeris

Definition

Overview

Introduced by Tri Dao and collaborators in 2022, FlashAttention tackles a subtle but critical inefficiency in standard transformer attention implementations: the bottleneck is often not raw floating-point compute but memory bandwidth, specifically the repeated movement of large intermediate attention matrices between a GPU's fast on-chip SRAM and its much slower high-bandwidth memory (HBM). Naive attention implementations materialize the full n×n attention score matrix in HBM, incurring expensive read/write traffic at every layer. FlashAttention avoids this by using tiling: it splits the query, key, and value matrices into blocks small enough to fit in SRAM, computes partial attention outputs for each block using an online softmax technique that avoids ever storing the full attention matrix, and accumulates the final result incrementally. Because it is IO-aware rather than algorithmically approximate, FlashAttention produces numerically identical (up to floating point rounding) results to standard attention — it is a systems-level optimization, not a sparsity or approximation technique like Sparse Transformers or Linformer. The original paper reported 2-4x training speedups and dramatically reduced memory usage (from quadratic to linear in sequence length) compared to standard PyTorch attention implementations, which in turn enabled training and inference with much longer context windows. FlashAttention-2 (2023) improved GPU occupancy and parallelism further, and FlashAttention-3 (2024) added support for Hopper-generation GPU features like FP8 and asynchronous execution. FlashAttention has since become a near-universal default: it ships as a fused kernel in PyTorch's scaled_dot_product_attention, is integrated into Hugging Face Transformers, vLLM, and virtually every major open-weight and proprietary LLM training and inference stack. It's a foundational reason modern LLMs can be trained and served efficiently at context lengths of tens or hundreds of thousands of tokens.

Key Concepts

Computes exact (not approximate) attention with no change to model outputs
Uses tiling and online softmax to avoid materializing the full n×n attention matrix
Reduces memory usage from quadratic to linear in sequence length
Delivers 2-4x+ training/inference speedups versus naive attention implementations
FlashAttention-2 improved GPU parallelism and work partitioning
FlashAttention-3 adds Hopper GPU support including FP8 and warp-specialization
Integrated as a default kernel in PyTorch, Hugging Face Transformers, and vLLM
Open-source implementation widely adopted across the LLM training ecosystem

Use Cases

Speeding up transformer training runs on GPU clusters

Enabling long-context LLM training and fine-tuning that would otherwise run out of memory

Reducing inference latency and cost for production LLM serving

Powering efficient attention kernels inside inference engines like vLLM and TensorRT-LLM

Making long-document and long-context applications (RAG, code assistants) computationally feasible

Frequently Asked Questions

From the Blog

AI & Technology

How Large Language Models Actually Work

LLMs seem magical until you understand what they are: next-token predictors trained on massive text corpora. This guide explains tokenisation, embeddings, the transformer architecture, attention mechanism, and how training works — without requiring a maths degree.