Flash Attention from Scratch Part 1: Intro
lubits.ch · 1,461 words · saved by 1 readers
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
related reading
- We reverse-engineered Flash Attention 4modal.com
- 2307.08691arxiv.org
- [2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenessarxiv.org
- Biao's Bloghebiao064.github.io
- How to Optimize a CUDA Matmul Kernel for cuBLAS-like Performance: a Worklogsiboehm.com
- 2205.14135arxiv.org
- About Us - Colfax Researchresearch.colfax-intl.com
- From Online Softmax to FlashAttentioncourses.cs.washington.edu
- GPUs Go Brrr · Hazy Researchhazyresearch.stanford.edu
- Transformers Inference Optimization Toolset | AstraBlogastralord.github.io
- Inside NVIDIA GPUs: Anatomy of high performance matmul kernels - Aleksa Gordićaleksagordic.com
- Modern GPU Programming For MLSys — Modern GPU Programming For MLSysmlc.ai