Understanding Flash Attention: Writing the Algorithm from Scratch in Triton
Why is Flash Attention so fast? Find out how Flash Attention works. Afterward, we'll polish our understanding by writing a GPU kernel of the algorithm in Triton.
Alex Dremov
AI RESEARCHER • EPFLI'm a Machine Learning Researcher and Engineer. Here I write posts on the intersection of deep learning theory and high-performance system engineering
Featured
MACHINE LEARNING, PAPER · 8 MONTHS AGO · 13 MIN READ
Rethinking Quantization-Aware Training: Why Your QAT Length is Probably Wrong
Training quantized neural networks involves a fundamental trade-off: how should you divide your compute budget between full-precision pretraining and quantization-aware training?
Read the Article →Speed Up PyTorch With Custom Kernels. But It Gets Progressively Darker
It's all about making your models run faster, from flicking a magic “compile” switch to writing your own custom GPU code. In each step, we’ll implement an innocent softmax function, but things are about to get dark by the end.
Simple Ways to Speed Up Your PyTorch Model Training
If all machine learning engineers want one thing, it's faster model training — maybe after good test metrics.
History
All Posts2025
2023
2022
IOS & SWIFT
4 MIN READUsing Threads in Swift
4 MIN READUsing Threads in Swift
IOS & SWIFT
4 MIN READTop 7 Subtle Swift Features
4 MIN READTop 7 Subtle Swift Features
2021
TOOLS
2 MIN READNote-taking apps
2 MIN READNote-taking apps
CODE
15 MIN READThe Mystery of Mach-O Object Structure
15 MIN READThe Mystery of Mach-O Object Structure
2020