Attention and matrix kernels

In one paragraph

Attention and matrix multiplies are where most of a language model's GPU time goes. FlashAttention computes attention block by block without writing the full score matrix to memory, using an online softmax that is exact. Matrix kernels for mixture-of-experts models and low-precision formats add their own requirements: many small multiplies of uneven size, and scale factors applied inside the kernel.

How it works

• Online softmax. Each row keeps a running maximum and sum; each new block of scores rescales the partial result. Memory traffic grows linearly with sequence length instead of quadratically.
• The FlashAttention line. FlashAttention, then FlashAttention-2 (better parallelism), FlashAttention-3 (Hopper asynchrony and FP8) and FlashAttention-4 (Hopper and Blackwell).
• FlexAttention. Attention variants (causal, sliding window, document masks, score modifications) written as small Python functions and compiled into a fused kernel.
• Paged-attention kernels. Read the attention cache from non-contiguous blocks, matching how serving engines store it.
• Mixture-of-experts. Grouped GEMMs run many uneven expert multiplies in one launch; fused MoE kernels combine routing, expert compute and the weighted sum.
• Low precision. FP8 and NVFP4 GEMMs apply per-block scales inside the kernel; DeepGEMM is a published FP8 example.

Design choices and trade-offs

Using the attention kernel shipped with the serving engine is almost always right. A custom variant is justified when a new attention pattern is needed that existing kernels do not support, and FlexAttention may cover it first.

When not to use it

When the model and engine already use a FlashAttention-class kernel for the hardware.

How we evaluate it

Exactness against a reference implementation at the tolerances stated in profiling and verifying kernels, then speed on the real sequence lengths and batch sizes.

Sources

• Dao et al. (2022), FlashAttention (opens in a new tab); FlashAttention repository (opens in a new tab); FlashAttention-4 (opens in a new tab)
• FlexAttention (opens in a new tab), DeepGEMM (opens in a new tab)