Flash Attention from Scratch Part 1: Intro

Intro In this 10-part series, we're going to implement Flash Attention 2 from scratch on Ampere GPUs.

Intro In this 10-part series, we're going to implement Flash Attention 2 from scratch on Ampere GPUs . We'll build an initial implementation and optimize it over 16 kernel iterations , all without importing any external libraries. By the final kernel, we'll reach 99.2% the performance of the official implementation on the A100 and 102.9% on the RTX 3090 (at sequence length 4096, ). You can find the code here . Prerequisites What you'll need: Solid CUDA experience (memory hierarchy, occupancy, tiling, bank conflicts, etc.) Familiarity with attention mechanisms and Flash Attention basics New to

saved by

• Ashish Sinha

related reading

• We reverse-engineered Flash Attention 4modal.com
• How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklogsiboehm.com
• GPUs Go Brrr · Hazy Researchhazyresearch.stanford.edu
• From Online Softmax to FlashAttentioncourses.cs.washington.edu
• Transformers Inference Optimization Toolset | AstraBlogastralord.github.io
• Inside NVIDIA GPUs: Anatomy of high performance matmul kernels - Aleksa Gordićaleksagordic.com
• [2410.20399] ThunderKittens: Simple, Fast, and Adorable AI Kernelsarxiv.org
• [2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenessarxiv.org
• Making Deep Learning go Brrrr From First Principleshorace.io
• Transformer Inference Arithmetic | kipply's blogkipp.ly
• The Best GPUs for Deep Learning in 2023 — An In-depth Analysistimdettmers.com
• Overleaf Examplearxiv.org