Learning FlashAttention the Hard Way. Part 1: The Algebraic Foundation [D]
Mirrored from r/MachineLearning for archival readability. Support the source by reading on the original site.
I'm writing a short series of tutorials on FlashAttention: from theory to efficient CUDA kernels.
Part 1 is the theoretical foundation. It walks through a modern algebraic formalism showing that FlashAttention is an associative operation, which lets you treat it as a regular reduction on the GPU and apply all the same scheduling optimizations. Some recent MLSys and CVPR papers lean on this framing, and I find it much more powerful than the original.
Overview:
- Safe softmax, Welford's variance, and FlashAttention are the same secretly-associative operation
- The twisted monoid (transport of structure), why the max-rescale coupling doesn't break associativity
- The qk_scale = log2(e)/√D you already see in FA-2 and the Triton kernels, derived from scratch
- Numerical analysis: overflow bounds, error limits, and why tiling never amplifies error
- Bird's 3rd Homomorphism Theorem as a test for whether any loop is secretly associative
[link] [comments]
More from r/MachineLearning
-
For the people who got reviews back from neurips, cvpr, eccv, etc and also tested their paper through an agentic reviewer like the stanford one, how different were the reviews? [D]
Aug 14
-
Building text to ASCII diffusion model , need advice and guidance [P]
Aug 14
-
A collision-entropy floor for watermark/retrieval AI-text detection. Looking for a sanity check before I take this further [D]
Aug 14
-
Are supervised and unsupervised learning still relevant today? [D]
Aug 14
Discussion (0)
Sign in to join the discussion. Free account, 30 seconds — email code or GitHub.
Sign in →No comments yet. Sign in and be the first to say something.